[reactos] 01/02: [NTOS:PNP] Handle arbitrary subkey name lengths in EnumerateDevices. CORE-15882

[Thomas Faber]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=b8e091a44c9852fb759e9c...

commit b8e091a44c9852fb759e9cb7f503bab74146a2d1
Author:     Thomas Faber [e-mail]
AuthorDate: Sat Apr 20 10:48:43 2019 +0200
Commit:     Thomas Faber [e-mail]
CommitDate: Wed Jul 10 10:58:45 2019 +0200

[NTOS:PNP] Handle arbitrary subkey name lengths in EnumerateDevices. CORE-15882

This also makes it necessary to fix a bug in the previous code:
    ZwEnumerateKey will not account for space for a null terminator, so to
    ensure we have space, we must allocate the additional WCHAR, but not
    include it in the buffer size passed to the function.
---
 ntoskrnl/io/pnpmgr/pnproot.c | 53 +++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 47 insertions(+), 6 deletions(-)

diff --git a/ntoskrnl/io/pnpmgr/pnproot.c b/ntoskrnl/io/pnpmgr/pnproot.c
index 34b529ad509..65ea5cc937b 100644
--- a/ntoskrnl/io/pnpmgr/pnproot.c
+++ b/ntoskrnl/io/pnpmgr/pnproot.c
@@ -450,7 +450,7 @@ EnumerateDevices(
     HANDLE KeyHandle = NULL;
     HANDLE SubKeyHandle = NULL;
     HANDLE DeviceKeyHandle = NULL;
-    ULONG BufferSize;
+    ULONG KeyInfoSize, SubKeyInfoSize;
     ULONG ResultSize;
     ULONG Index1, Index2;
     BUFFER Buffer1, Buffer2;
@@ -461,15 +461,20 @@ EnumerateDevices(
     DeviceExtension = (PPNPROOT_FDO_DEVICE_EXTENSION)DeviceObject->DeviceExtension;
     KeAcquireGuardedMutex(&DeviceExtension->DeviceListLock);

-    BufferSize = sizeof(KEY_BASIC_INFORMATION) + (MAX_PATH + 1) * sizeof(WCHAR);
-    KeyInfo = ExAllocatePoolWithTag(PagedPool, BufferSize, TAG_PNP_ROOT);
+    KeyInfoSize = sizeof(KEY_BASIC_INFORMATION) + (MAX_PATH + 1) * sizeof(WCHAR);
+    KeyInfo = ExAllocatePoolWithTag(PagedPool,
+                                    KeyInfoSize + sizeof(UNICODE_NULL),
+                                    TAG_PNP_ROOT);
     if (!KeyInfo)
     {
         DPRINT("ExAllocatePoolWithTag() failed\n");
         Status = STATUS_NO_MEMORY;
         goto cleanup;
     }
-    SubKeyInfo = ExAllocatePoolWithTag(PagedPool, BufferSize, TAG_PNP_ROOT);
+    SubKeyInfoSize = KeyInfoSize;
+    SubKeyInfo = ExAllocatePoolWithTag(PagedPool,
+                                       SubKeyInfoSize + sizeof(UNICODE_NULL),
+                                       TAG_PNP_ROOT);
     if (!SubKeyInfo)
     {
         DPRINT("ExAllocatePoolWithTag() failed\n");
@@ -496,13 +501,30 @@ EnumerateDevices(
             Index1,
             KeyBasicInformation,
             KeyInfo,
-            BufferSize,
+            KeyInfoSize,
             &ResultSize);
         if (Status == STATUS_NO_MORE_ENTRIES)
         {
             Status = STATUS_SUCCESS;
             break;
         }
+        else if (Status == STATUS_BUFFER_OVERFLOW ||
+                 Status == STATUS_BUFFER_TOO_SMALL)
+        {
+            ASSERT(KeyInfoSize < ResultSize);
+            KeyInfoSize = ResultSize;
+            ExFreePoolWithTag(KeyInfo, TAG_PNP_ROOT);
+            KeyInfo = ExAllocatePoolWithTag(PagedPool,
+                                            KeyInfoSize + sizeof(UNICODE_NULL),
+                                            TAG_PNP_ROOT);
+            if (!KeyInfo)
+            {
+                DPRINT1("ExAllocatePoolWithTag(%lu) failed\n", KeyInfoSize);
+                Status = STATUS_NO_MEMORY;
+                goto cleanup;
+            }
+            continue;
+        }
         else if (!NT_SUCCESS(Status))
         {
             DPRINT("ZwEnumerateKey() failed with status 0x%08lx\n", Status);
@@ -539,10 +561,29 @@ EnumerateDevices(
                 Index2,
                 KeyBasicInformation,
                 SubKeyInfo,
-                BufferSize,
+                SubKeyInfoSize,
                 &ResultSize);
             if (Status == STATUS_NO_MORE_ENTRIES)
+            {
                 break;
+            }
+            else if (Status == STATUS_BUFFER_OVERFLOW ||
+                     Status == STATUS_BUFFER_TOO_SMALL)
+            {
+                ASSERT(SubKeyInfoSize < ResultSize);
+                SubKeyInfoSize = ResultSize;
+                ExFreePoolWithTag(SubKeyInfo, TAG_PNP_ROOT);
+                SubKeyInfo = ExAllocatePoolWithTag(PagedPool,
+                                                   SubKeyInfoSize + sizeof(UNICODE_NULL),
+                                                   TAG_PNP_ROOT);
+                if (!SubKeyInfo)
+                {
+                    DPRINT1("ExAllocatePoolWithTag(%lu) failed\n", SubKeyInfoSize);
+                    Status = STATUS_NO_MEMORY;
+                    goto cleanup;
+                }
+                continue;
+            }
             else if (!NT_SUCCESS(Status))
             {
                 DPRINT("ZwEnumerateKey() failed with status 0x%08lx\n", Status);