[PATCH v3 01/17] 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 at gmail.com>
---

Supersedes 150480

v2: Retrieve the registry value without failing in rare race conditions.
v3: Don't open the HKEY_LOCAL_MACHINE twice under some circumstances.

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

diff --git a/dlls/shell32/autocomplete.c b/dlls/shell32/autocomplete.c
index cf8da50..051644a 100644
--- a/dlls/shell32/autocomplete.c
+++ b/dlls/shell32/autocomplete.c
@@ -486,40 +486,66 @@ 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);
+        BOOL try_HKLM;
+        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);
+            try_HKLM = TRUE;
+
+            /* 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);
+                try_HKLM = FALSE;
+            }
+
+            if (res == ERROR_SUCCESS) {
+                while ((qc = heap_alloc(sz)) != NULL) {
+                    res = RegQueryValueExW(hKey, value, NULL, &type, qc, &sz);
+                    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) {
+                        if (!try_HKLM)
+                            break;
+                        RegCloseKey(hKey);
+                        res = RegOpenKeyExW(HKEY_LOCAL_MACHINE, key, 0, KEY_READ, &hKey);
+                        if (res != ERROR_SUCCESS)
+                            goto free_key_str;
+                        sz = MAX_PATH * sizeof(WCHAR);
+                        try_HKLM = FALSE;
+                    }
+                }
+                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