[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)
{
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 */