From f1d2a4485975004e5ff1c666ed2d8ebefe5282b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?George=20Bi=C8=99oc?= <george.bisoc@reactos.org>
Date: Sat, 23 Mar 2024 20:28:01 +0100
Subject: [PATCH] [NTOS:CM] Lock the cached KCB before removing it from cache
 entries

- Annotate the CmpEnumerateOpenSubKeys function with SAL2
- When removing an orphaned cached KCB, ensure that it is locked before clearing it from cache table entries
---
 ntoskrnl/config/cmapi.c        | 22 ++++++++++++++++------
 ntoskrnl/config/ntapi.c        |  2 +-
 ntoskrnl/include/internal/cm.h |  7 ++++---
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/ntoskrnl/config/cmapi.c b/ntoskrnl/config/cmapi.c
index 4d4edfde77a..944b469fa17 100644
--- a/ntoskrnl/config/cmapi.c
+++ b/ntoskrnl/config/cmapi.c
@@ -2242,7 +2242,7 @@ CmUnloadKey(IN PCM_KEY_CONTROL_BLOCK Kcb,
     {
         if (Flags != REG_FORCE_UNLOAD)
         {
-            if (CmpEnumerateOpenSubKeys(Kcb, FALSE, FALSE) != 0)
+            if (CmpEnumerateOpenSubKeys(Kcb, FALSE, FALSE, FALSE) != 0)
             {
                 /* There are open subkeys but we don't force hive unloading, fail */
                 Hive->HiveFlags &= ~HIVE_IS_UNLOADING;
@@ -2252,7 +2252,7 @@ CmUnloadKey(IN PCM_KEY_CONTROL_BLOCK Kcb,
         else
         {
             DPRINT1("CmUnloadKey: Force unloading is HALF-IMPLEMENTED, expect dangling KCBs problems!\n");
-            if (CmpEnumerateOpenSubKeys(Kcb, TRUE, TRUE) != 0)
+            if (CmpEnumerateOpenSubKeys(Kcb, TRUE, TRUE, TRUE) != 0)
             {
                 /* There are open subkeys that we cannot force to unload, fail */
                 Hive->HiveFlags &= ~HIVE_IS_UNLOADING;
@@ -2340,9 +2340,10 @@ CmUnloadKey(IN PCM_KEY_CONTROL_BLOCK Kcb,
 ULONG
 NTAPI
 CmpEnumerateOpenSubKeys(
-    IN PCM_KEY_CONTROL_BLOCK RootKcb,
-    IN BOOLEAN RemoveEmptyCacheEntries,
-    IN BOOLEAN DereferenceOpenedEntries)
+    _In_ PCM_KEY_CONTROL_BLOCK RootKcb,
+    _In_ BOOLEAN LockHeldExclusively,
+    _In_ BOOLEAN RemoveEmptyCacheEntries,
+    _In_ BOOLEAN DereferenceOpenedEntries)
 {
     PCM_KEY_HASH Entry;
     PCM_KEY_CONTROL_BLOCK CachedKcb;
@@ -2430,11 +2431,20 @@ CmpEnumerateOpenSubKeys(
                     }
                     else if ((CachedKcb->RefCount == 0) && RemoveEmptyCacheEntries)
                     {
+                        /* Lock the cached KCB of subkey before removing it from cache entries */
+                        if (!LockHeldExclusively)
+                            CmpAcquireKcbLockExclusive(CachedKcb);
+
                         /* Remove the current key from the delayed close list */
                         CmpRemoveFromDelayedClose(CachedKcb);
 
                         /* Remove the current cache entry */
-                        CmpCleanUpKcbCacheWithLock(CachedKcb, TRUE);
+                        // Lock is either held by ourselves or registry is locked exclusively
+                        CmpCleanUpKcbCacheWithLock(CachedKcb, LockHeldExclusively);
+
+                        /* Unlock the cached KCB if it was done by ourselves */
+                        if (!LockHeldExclusively)
+                            CmpReleaseKcbLock(CachedKcb);
 
                         /* Restart, because the hash list has changed */
                         Entry = CmpCacheTable[i].Entry;
diff --git a/ntoskrnl/config/ntapi.c b/ntoskrnl/config/ntapi.c
index 5bf9ec5af0b..6748326a2c8 100644
--- a/ntoskrnl/config/ntapi.c
+++ b/ntoskrnl/config/ntapi.c
@@ -1565,7 +1565,7 @@ NtQueryOpenSubKeys(IN POBJECT_ATTRIBUTES TargetKey,
 
     /* Call the internal API */
     SubKeys = CmpEnumerateOpenSubKeys(KeyBody->KeyControlBlock,
-                                      FALSE, FALSE);
+                                      TRUE, FALSE, FALSE);
 
     /* Unlock the registry */
     CmpUnlockRegistry();
diff --git a/ntoskrnl/include/internal/cm.h b/ntoskrnl/include/internal/cm.h
index 7657fc77fb9..a211862531c 100644
--- a/ntoskrnl/include/internal/cm.h
+++ b/ntoskrnl/include/internal/cm.h
@@ -1333,9 +1333,10 @@ CmUnloadKey(
 ULONG
 NTAPI
 CmpEnumerateOpenSubKeys(
-    IN PCM_KEY_CONTROL_BLOCK RootKcb,
-    IN BOOLEAN RemoveEmptyCacheEntries,
-    IN BOOLEAN DereferenceOpenedEntries
+    _In_ PCM_KEY_CONTROL_BLOCK RootKcb,
+    _In_ BOOLEAN LockHeldExclusively,
+    _In_ BOOLEAN RemoveEmptyCacheEntries,
+    _In_ BOOLEAN DereferenceOpenedEntries
 );
 
 HCELL_INDEX