From 413b5a0827af51e102f99f3b2cb59867cdbacb6b Mon Sep 17 00:00:00 2001
From: Whindmar Saksit <whindsaks@proton.me>
Date: Mon, 13 Nov 2023 18:02:41 +0100
Subject: [PATCH] [ADVAPI32] Handle HKEY_CLASSES_ROOT in RegQueryInfoKeyW
 (#5870)

CORE-8582 , CORE-14676

This fixes the bug where Regedit is unable to show all the keys in HKCR when a key exists in both HKCU and HKLM.
---
 dll/win32/advapi32/reg/hkcr.c | 65 +++++++++++++++++++++++++++++++++++
 dll/win32/advapi32/reg/reg.c  | 10 ++++++
 dll/win32/advapi32/reg/reg.h  | 15 ++++++++
 3 files changed, 90 insertions(+)

diff --git a/dll/win32/advapi32/reg/hkcr.c b/dll/win32/advapi32/reg/hkcr.c
index f55a783eb5d..21921ebb42b 100644
--- a/dll/win32/advapi32/reg/hkcr.c
+++ b/dll/win32/advapi32/reg/hkcr.c
@@ -1014,3 +1014,68 @@ Exit:
 
     return ErrorCode;
 }
+
+/* HKCR version of RegQueryInfoKeyW */
+LONG
+WINAPI
+QueryInfoHKCRKey(
+    _In_ HKEY hKey,
+    _Out_writes_to_opt_(*lpcchClass, *lpcchClass + 1) LPWSTR lpClass,
+    _Inout_opt_ LPDWORD lpcchClass,
+    _Reserved_ LPDWORD lpReserved,
+    _Out_opt_ LPDWORD lpcSubKeys,
+    _Out_opt_ LPDWORD lpcbMaxSubKeyLen,
+    _Out_opt_ LPDWORD lpcbMaxClassLen,
+    _Out_opt_ LPDWORD lpcValues,
+    _Out_opt_ LPDWORD lpcbMaxValueNameLen,
+    _Out_opt_ LPDWORD lpcbMaxValueLen,
+    _Out_opt_ LPDWORD lpcbSecurityDescriptor,
+    _Out_opt_ PFILETIME lpftLastWriteTime)
+{
+    HKEY PreferredKey, FallbackKey;
+    LONG retval, err;
+    DWORD OtherSubKeys = 0, OtherMaxSub = 0, OtherMaxClass = 0;
+    DWORD OtherValues = 0, OtherMaxName = 0, OtherMaxVal = 0;
+
+    ASSERT(IsHKCRKey(hKey));
+
+    /* Remove the HKCR flag while we're working */
+    hKey = (HKEY)(((ULONG_PTR)hKey) & ~0x2);
+
+    retval = GetFallbackHKCRKey(hKey, &FallbackKey, FALSE);
+    if (retval == ERROR_SUCCESS)
+    {
+        retval = RegQueryInfoKeyW(FallbackKey, lpClass, lpcchClass, lpReserved,
+                                  &OtherSubKeys, &OtherMaxSub, &OtherMaxClass,
+                                  &OtherValues, &OtherMaxName, &OtherMaxVal,
+                                  lpcbSecurityDescriptor, lpftLastWriteTime);
+        if (FallbackKey != hKey)
+            RegCloseKey(FallbackKey);
+    }
+
+    err = GetPreferredHKCRKey(hKey, &PreferredKey);
+    if (err == ERROR_SUCCESS)
+    {
+        err = RegQueryInfoKeyW(PreferredKey, lpClass, lpcchClass, lpReserved,
+                               lpcSubKeys, lpcbMaxSubKeyLen, lpcbMaxClassLen,
+                               lpcValues, lpcbMaxValueNameLen, lpcbMaxValueLen,
+                               lpcbSecurityDescriptor, lpftLastWriteTime);
+        if (PreferredKey != hKey)
+            RegCloseKey(PreferredKey);
+    }
+
+    if (lpcSubKeys)
+        *lpcSubKeys = (err ? 0 : *lpcSubKeys) + OtherSubKeys;
+    if (lpcValues)
+        *lpcValues = (err ? 0 : *lpcValues) + OtherValues;
+    if (lpcbMaxSubKeyLen)
+        *lpcbMaxSubKeyLen = max((err ? 0 : *lpcbMaxSubKeyLen), OtherMaxSub);
+    if (lpcbMaxClassLen)
+        *lpcbMaxClassLen = max((err ? 0 : *lpcbMaxClassLen), OtherMaxClass);
+    if (lpcbMaxValueNameLen)
+        *lpcbMaxValueNameLen = max((err ? 0 : *lpcbMaxValueNameLen), OtherMaxName);
+    if (lpcbMaxValueLen)
+        *lpcbMaxValueLen = max((err ? 0 : *lpcbMaxValueLen), OtherMaxVal);
+
+    return (err == ERROR_SUCCESS) ? ERROR_SUCCESS : retval;
+}
diff --git a/dll/win32/advapi32/reg/reg.c b/dll/win32/advapi32/reg/reg.c
index 58515793c4e..1ded4b54222 100644
--- a/dll/win32/advapi32/reg/reg.c
+++ b/dll/win32/advapi32/reg/reg.c
@@ -3718,6 +3718,16 @@ RegQueryInfoKeyW(HKEY hKey,
         return RtlNtStatusToDosError(Status);
     }
 
+    if (IsHKCRKey(KeyHandle))
+    {
+        ErrorCode = QueryInfoHKCRKey(KeyHandle, lpClass, lpcClass, lpReserved,
+                                     lpcSubKeys, lpcMaxSubKeyLen, lpcMaxClassLen,
+                                     lpcValues, lpcMaxValueNameLen, lpcMaxValueLen,
+                                     lpcbSecurityDescriptor, lpftLastWriteTime);
+        ClosePredefKey(KeyHandle);
+        return ErrorCode;
+    }
+
     if (lpClass != NULL)
     {
         if (*lpcClass > 0)
diff --git a/dll/win32/advapi32/reg/reg.h b/dll/win32/advapi32/reg/reg.h
index 2c5c128d38c..469f2b6fbcd 100644
--- a/dll/win32/advapi32/reg/reg.h
+++ b/dll/win32/advapi32/reg/reg.h
@@ -96,3 +96,18 @@ EnumHKCRValue(
     _Out_opt_ LPBYTE data,
     _Inout_opt_ PDWORD count);
 
+LONG
+WINAPI
+QueryInfoHKCRKey(
+    _In_ HKEY hKey,
+    _Out_writes_to_opt_(*lpcchClass, *lpcchClass + 1) LPWSTR lpClass,
+    _Inout_opt_ LPDWORD lpcchClass,
+    _Reserved_ LPDWORD lpReserved,
+    _Out_opt_ LPDWORD lpcSubKeys,
+    _Out_opt_ LPDWORD lpcbMaxSubKeyLen,
+    _Out_opt_ LPDWORD lpcbMaxClassLen,
+    _Out_opt_ LPDWORD lpcValues,
+    _Out_opt_ LPDWORD lpcbMaxValueNameLen,
+    _Out_opt_ LPDWORD lpcbMaxValueLen,
+    _Out_opt_ LPDWORD lpcbSecurityDescriptor,
+    _Out_opt_ PFILETIME lpftLastWriteTime);