[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.
This commit is contained in:
Timo Kreuzer 2020-01-03 21:43:44 +01:00
parent 47a350c4db
commit 2c2c570317

View file

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