From 2c2c5703172b68c1d253065cb4e76cfc575cfd3b Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Fri, 3 Jan 2020 21:43:44 +0100
Subject: [PATCH] [PSAPI] Fix and simplify FindDeviceDriver

The previous version was miscalculating the buffer size on x64 (due to alignment) and always using a too small buffer.
The new version removes the size calculation entirely and uses the required size returned by NtQuerySystemInformation.
---
 dll/win32/psapi/psapi.c | 40 ++++++++++++----------------------------
 1 file changed, 12 insertions(+), 28 deletions(-)

diff --git a/dll/win32/psapi/psapi.c b/dll/win32/psapi/psapi.c
index 6eae8375f71..79eb5610eff 100644
--- a/dll/win32/psapi/psapi.c
+++ b/dll/win32/psapi/psapi.c
@@ -43,13 +43,13 @@ FindDeviceDriver(IN PVOID ImageBase,
                  OUT PRTL_PROCESS_MODULE_INFORMATION MatchingModule)
 {
     NTSTATUS Status;
-    DWORD NewSize, Count;
+    DWORD i, RequiredSize;
     PRTL_PROCESS_MODULES Information;
     RTL_PROCESS_MODULE_INFORMATION Module;
     /* By default, to prevent too many reallocations, we already make room for 4 modules */
     DWORD Size = sizeof(RTL_PROCESS_MODULES) + 3 * sizeof(RTL_PROCESS_MODULE_INFORMATION);
 
-    do
+    while (TRUE)
     {
         /* Allocate a buffer to hold modules information */
         Information = LocalAlloc(LMEM_FIXED, Size);
@@ -60,13 +60,10 @@ FindDeviceDriver(IN PVOID ImageBase,
         }
 
         /* Query information */
-        Status = NtQuerySystemInformation(SystemModuleInformation, Information, Size, &Count);
-        /* In case of an error */
+        Status = NtQuerySystemInformation(SystemModuleInformation, Information, Size, &RequiredSize);
         if (!NT_SUCCESS(Status))
         {
-            /* Save the amount of output modules */
-            NewSize = Information->NumberOfModules;
-            /* And free buffer */
+            /* Free the current buffer */
             LocalFree(Information);
 
             /* If it was not a length mismatch (ie, buffer too small), just leave */
@@ -76,21 +73,8 @@ FindDeviceDriver(IN PVOID ImageBase,
                 return FALSE;
             }
 
-            /* Compute new size length */
-            ASSERT(Size >= sizeof(RTL_PROCESS_MODULES));
-            NewSize *= sizeof(RTL_PROCESS_MODULE_INFORMATION);
-            NewSize += sizeof(ULONG);
-            ASSERT(NewSize >= sizeof(RTL_PROCESS_MODULES));
-            /* Check whether it is really bigger - otherwise, leave */
-            if (NewSize < Size)
-            {
-                ASSERT(NewSize > Size);
-                SetLastError(RtlNtStatusToDosError(STATUS_INFO_LENGTH_MISMATCH));
-                return FALSE;
-            }
-
-            /* Loop again with that new buffer */
-            Size = NewSize;
+            /* Try again with the required size */
+            Size = RequiredSize;
             continue;
         }
 
@@ -101,9 +85,9 @@ FindDeviceDriver(IN PVOID ImageBase,
         }
 
         /* Try to find which module matches the base address given */
-        for (Count = 0; Count < Information->NumberOfModules; ++Count)
+        for (i = 0; i < Information->NumberOfModules; ++i)
         {
-            Module = Information->Modules[Count];
+            Module = Information->Modules[i];
             if (Module.ImageBase == ImageBase)
             {
                 /* Copy the matching module and leave */
@@ -115,7 +99,7 @@ FindDeviceDriver(IN PVOID ImageBase,
 
         /* If we arrive here, it means we were not able to find matching base address */
         break;
-    } while (TRUE);
+    }
 
     /* Release and leave */
     LocalFree(Information);
@@ -231,7 +215,7 @@ CallBackConvertToAscii(LPVOID pContext,
                        LPCWSTR lpFilename)
 {
     BOOL Ret;
-    DWORD Len;
+    SIZE_T Len;
     LPSTR AnsiFileName;
     PINTERNAL_ENUM_PAGE_FILES_CONTEXT Context = (PINTERNAL_ENUM_PAGE_FILES_CONTEXT)pContext;
 
@@ -646,7 +630,7 @@ GetDeviceDriverBaseNameA(LPVOID ImageBase,
                          LPSTR lpBaseName,
                          DWORD nSize)
 {
-    DWORD Len, LenWithNull;
+    SIZE_T Len, LenWithNull;
     RTL_PROCESS_MODULE_INFORMATION Module;
 
     /* Get the associated device driver to the base address */
@@ -685,7 +669,7 @@ GetDeviceDriverFileNameA(LPVOID ImageBase,
                          LPSTR lpFilename,
                          DWORD nSize)
 {
-    DWORD Len, LenWithNull;
+    SIZE_T Len, LenWithNull;
     RTL_PROCESS_MODULE_INFORMATION Module;
 
     /* Get the associated device driver to the base address */