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

[Huw Davies]

On Wed, Sep 05, 2018 at 07:13:03PM +0300, Gabriel Ivăncescu wrote:
> 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;

This try_HKLM thing is ugly.  I'd suggest having a outer loop over the two registry
root keys (HKCU and HKLM) which tries to access the supplied key / value.  If the
access succeeds then break out of the loop.

Also, be consistent about your brace style (I know the function is already mixed).
I'd be happy if you moved all of the opening braces, in code you're rewriting, to a
new line.


> +
> +            /* 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
> 
> 
>