[PATCH 3/3] 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 properly
read the registry value with arbitrary size, even REG_SZ without a NUL
terminator (MSDN states they are possible).

Signed-off-by: Gabriel Ivăncescu <gabrielopcode at gmail.com>
---
 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..97ce61a 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;
+        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) {
+                WCHAR *qc;
+                DWORD type, sz;
+                res = RegQueryValueExW(hKey, value, NULL, &type, NULL, &sz);
+
+                if (res != ERROR_SUCCESS || type != REG_SZ) {
+                    RegCloseKey(hKey);
+                    res = RegOpenKeyExW(HKEY_LOCAL_MACHINE, key, 0, KEY_READ, &hKey);
+                    if (res != ERROR_SUCCESS)
+                        goto free_key_str;
+                    res = RegQueryValueExW(hKey, value, NULL, &type, NULL, &sz);
+                }
+
+                /* the size does include the NUL terminator, however, strings
+                   can be stored without it, so make room for NUL just in case */
+                if (res == ERROR_SUCCESS && type == REG_SZ &&
+                    (qc = heap_alloc(sz + sizeof(*qc))) != NULL)
+                {
+                    DWORD old_sz = sz;
+                    res = RegQueryValueExW(hKey, value, NULL, &type, (BYTE*)qc, &sz);
+
+                    /* make sure the value wasn't messed with */
+                    if (res == ERROR_SUCCESS && sz == old_sz && type == REG_SZ) {
+                        qc[sz / sizeof(*qc)] = '\0';
+                        This->quickComplete = qc;
+                    }
+                    else
+                        heap_free(qc);
+                }
+                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