[PATCH v2] shell32/autocomplete: Revamp pwzsRegKeyPath handling so it can deal with arbitrary sizes and make it more robust

[Gabriel Ivăncescu]

Handle heap_alloc failure, reg strings without a \ character at all, try
harder to find the reg path (if only value fails the lookup), and read the
registry value with any size.

Signed-off-by: Gabriel Ivăncescu <gabrielopcode(a)gmail.com>
---

v2: Retrieve the registry value without failing in rare race conditions.

 dlls/shell32/autocomplete.c | 93 +++++++++++++++++++++++++++++----------------
 1 file changed, 61 insertions(+), 32 deletions(-)

diff --git a/dlls/shell32/autocomplete.c b/dlls/shell32/autocomplete.c
index 5fe48e9..99d5ae7 100644
--- a/dlls/shell32/autocomplete.c
+++ b/dlls/shell32/autocomplete.c
@@ -486,40 +486,69 @@ static HRESULT WINAPI IAutoComplete2_fnInit(
     if (This->options & ACO_AUTOSUGGEST)
         create_listbox(This);
 
-    if (pwzsRegKeyPath) {
-	WCHAR *key;
-	WCHAR result[MAX_PATH];
-	WCHAR *value;
-	HKEY hKey = 0;
-	LONG res;
-	LONG len;
-
-	/* pwszRegKeyPath contains the key as well as the value, so we split */
-	key = heap_alloc((lstrlenW(pwzsRegKeyPath)+1)*sizeof(WCHAR));
-	strcpyW(key, pwzsRegKeyPath);
-	value = strrchrW(key, '\\');
-	*value = 0;
-	value++;
-	/* Now value contains the value and buffer the key */
-	res = RegOpenKeyExW(HKEY_CURRENT_USER, key, 0, KEY_READ, &hKey);
-	if (res != ERROR_SUCCESS) {
-	    /* if the key is not found, MSDN states we must seek in HKEY_LOCAL_MACHINE */
-	    res = RegOpenKeyExW(HKEY_LOCAL_MACHINE, key, 0, KEY_READ, &hKey);  
-	}
-	if (res == ERROR_SUCCESS) {
-	    res = RegQueryValueW(hKey, value, result, &len);
-	    if (res == ERROR_SUCCESS) {
-		This->quickComplete = heap_alloc(len*sizeof(WCHAR));
-		strcpyW(This->quickComplete, result);
-	    }
-	    RegCloseKey(hKey);
-	}
-	heap_free(key);
+    if (pwzsRegKeyPath)
+    {
+        WCHAR *key, *value;
+        DWORD type, sz = MAX_PATH * sizeof(WCHAR);
+        BYTE *qc;
+        HKEY hKey;
+        LSTATUS res;
+        size_t len;
+
+        /* pwszRegKeyPath contains the key as well as the value, so split it */
+        value = strrchrW(pwzsRegKeyPath, '\\');
+        len = value - pwzsRegKeyPath;
+
+        if (value && (key = heap_alloc((len+1) * sizeof(*key))) != NULL) {
+            memcpy(key, pwzsRegKeyPath, len * sizeof(*key));
+            key[len] = '\0';
+            value++;
+
+            res = RegOpenKeyExW(HKEY_CURRENT_USER, key, 0, KEY_READ, &hKey);
+
+            /* if not found, MSDN states we must seek in HKEY_LOCAL_MACHINE */
+            if (res != ERROR_SUCCESS)
+                res = RegOpenKeyExW(HKEY_LOCAL_MACHINE, key, 0, KEY_READ, &hKey);
+
+            if (res == ERROR_SUCCESS && (qc = heap_alloc(sz)) != NULL) {
+                res = RegQueryValueExW(hKey, value, NULL, &type, qc, &sz);
+
+                if ((res != ERROR_SUCCESS && res != ERROR_MORE_DATA) || type != REG_SZ)
+                {
+                    RegCloseKey(hKey);
+                    res = RegOpenKeyExW(HKEY_LOCAL_MACHINE, key, 0, KEY_READ, &hKey);
+                    if (res != ERROR_SUCCESS) {
+                        heap_free(qc);
+                        goto free_key_str;
+                    }
+                    sz = MAX_PATH * sizeof(WCHAR);
+                    res = RegQueryValueExW(hKey, value, NULL, &type, qc, &sz);
+                }
+
+                for (;;) {
+                    if (res == ERROR_SUCCESS && type == REG_SZ) {
+                        This->quickComplete = heap_realloc(qc, sz);
+                        break;
+                    }
+                    heap_free(qc);
+                    if (res != ERROR_MORE_DATA || type != REG_SZ)
+                        break;
+                    if ((qc = heap_alloc(sz)) == NULL)
+                        break;
+                    res = RegQueryValueExW(hKey, value, NULL, &type, qc, &sz);
+                }
+                RegCloseKey(hKey);
+            }
+          free_key_str:
+            heap_free(key);
+        }
     }
 
-    if ((pwszQuickComplete) && (!This->quickComplete)) {
-	This->quickComplete = heap_alloc((lstrlenW(pwszQuickComplete)+1)*sizeof(WCHAR));
-	lstrcpyW(This->quickComplete, pwszQuickComplete);
+    if (!This->quickComplete && pwszQuickComplete)
+    {
+        size_t len = strlenW(pwszQuickComplete)+1;
+        if ((This->quickComplete = heap_alloc(len * sizeof(WCHAR))) != NULL)
+            memcpy(This->quickComplete, pwszQuickComplete, len * sizeof(WCHAR));
     }
 
     return S_OK;
-- 
1.9.1