Make EnumDeviceDrivers and EnumProcessModules return the total size, even if the buffer is smaller. (from bug 3319)

svn path=/trunk/; revision=34112
This commit is contained in:
Jeffrey Morlan 2008-06-26 19:52:05 +00:00
parent 558c11db56
commit 7816c09fea

View file

@ -39,6 +39,7 @@ typedef struct _ENUM_DEVICE_DRIVERS_CONTEXT
{ {
LPVOID *lpImageBase; LPVOID *lpImageBase;
DWORD nCount; DWORD nCount;
DWORD nTotal;
} ENUM_DEVICE_DRIVERS_CONTEXT, *PENUM_DEVICE_DRIVERS_CONTEXT; } ENUM_DEVICE_DRIVERS_CONTEXT, *PENUM_DEVICE_DRIVERS_CONTEXT;
NTSTATUS STDCALL NTSTATUS STDCALL
@ -47,19 +48,18 @@ EnumDeviceDriversCallback(IN PRTL_PROCESS_MODULE_INFORMATION CurrentModule,
{ {
PENUM_DEVICE_DRIVERS_CONTEXT Context = (PENUM_DEVICE_DRIVERS_CONTEXT)CallbackContext; PENUM_DEVICE_DRIVERS_CONTEXT Context = (PENUM_DEVICE_DRIVERS_CONTEXT)CallbackContext;
/* no more buffer space */ /* check if more buffer space is available */
if(Context->nCount == 0) if(Context->nCount != 0)
{ {
return STATUS_INFO_LENGTH_MISMATCH; /* return current module */
*Context->lpImageBase = CurrentModule->ImageBase;
/* go to next array slot */
Context->lpImageBase++;
Context->nCount--;
} }
/* return current module */ Context->nTotal++;
*Context->lpImageBase = CurrentModule->ImageBase;
/* go to next array slot */
Context->lpImageBase++;
Context->nCount--;
return STATUS_SUCCESS; return STATUS_SUCCESS;
} }
@ -97,6 +97,7 @@ typedef struct _ENUM_PROCESS_MODULES_CONTEXT
{ {
HMODULE *lphModule; HMODULE *lphModule;
DWORD nCount; DWORD nCount;
DWORD nTotal;
} ENUM_PROCESS_MODULES_CONTEXT, *PENUM_PROCESS_MODULES_CONTEXT; } ENUM_PROCESS_MODULES_CONTEXT, *PENUM_PROCESS_MODULES_CONTEXT;
NTSTATUS STDCALL NTSTATUS STDCALL
@ -106,19 +107,18 @@ EnumProcessModulesCallback(IN HANDLE ProcessHandle,
{ {
PENUM_PROCESS_MODULES_CONTEXT Context = (PENUM_PROCESS_MODULES_CONTEXT)CallbackContext; PENUM_PROCESS_MODULES_CONTEXT Context = (PENUM_PROCESS_MODULES_CONTEXT)CallbackContext;
/* no more buffer space */ /* check if more buffer space is available */
if(Context->nCount == 0) if(Context->nCount != 0)
{ {
return STATUS_INFO_LENGTH_MISMATCH; /* return current process */
*Context->lphModule = CurrentModule->DllBase;
/* go to next array slot */
Context->lphModule++;
Context->nCount--;
} }
/* return current process */ Context->nTotal++;
*Context->lphModule = CurrentModule->DllBase;
/* go to next array slot */
Context->lphModule++;
Context->nCount--;
return STATUS_SUCCESS; return STATUS_SUCCESS;
} }
@ -678,23 +678,18 @@ EnumDeviceDrivers(LPVOID *lpImageBase,
ENUM_DEVICE_DRIVERS_CONTEXT Context; ENUM_DEVICE_DRIVERS_CONTEXT Context;
NTSTATUS Status; NTSTATUS Status;
if(cb == 0 || lpImageBase == NULL)
{
*lpcbNeeded = 0;
return TRUE;
}
cb /= sizeof(PVOID); cb /= sizeof(PVOID);
Context.lpImageBase = lpImageBase; Context.lpImageBase = lpImageBase;
Context.nCount = cb; Context.nCount = cb;
Context.nTotal = 0;
Status = PsaEnumerateSystemModules(EnumDeviceDriversCallback, &Context); Status = PsaEnumerateSystemModules(EnumDeviceDriversCallback, &Context);
/* return the count of bytes returned */ /* return the count of bytes that would be needed for a complete enumeration */
*lpcbNeeded = (cb - Context.nCount) * sizeof(PVOID); *lpcbNeeded = Context.nTotal * sizeof(PVOID);
if(!NT_SUCCESS(Status) && (Status != STATUS_INFO_LENGTH_MISMATCH)) if(!NT_SUCCESS(Status))
{ {
SetLastErrorByStatus(Status); SetLastErrorByStatus(Status);
return FALSE; return FALSE;
@ -757,21 +752,16 @@ EnumProcessModules(HANDLE hProcess,
cb /= sizeof(HMODULE); cb /= sizeof(HMODULE);
if(cb == 0 || lphModule == NULL)
{
*lpcbNeeded = 0;
return TRUE;
}
Context.lphModule = lphModule; Context.lphModule = lphModule;
Context.nCount = cb; Context.nCount = cb;
Context.nTotal = 0;
/* enumerate the process modules */ /* enumerate the process modules */
Status = PsaEnumerateProcessModules(hProcess, EnumProcessModulesCallback, &Context); Status = PsaEnumerateProcessModules(hProcess, EnumProcessModulesCallback, &Context);
*lpcbNeeded = (cb - Context.nCount) * sizeof(DWORD); *lpcbNeeded = Context.nTotal * sizeof(HMODULE);
if(!NT_SUCCESS(Status) && (Status != STATUS_INFO_LENGTH_MISMATCH)) if(!NT_SUCCESS(Status))
{ {
SetLastErrorByStatus(Status); SetLastErrorByStatus(Status);
return FALSE; return FALSE;