[ros-diffs] [jmorlan] 34112: Make EnumDeviceDrivers and EnumProcessModules return the total size, even if the buffer is smaller. (from bug 3319)

[jmorlan at svn.reactos.org]

Author: jmorlan
Date: Thu Jun 26 14:52:05 2008
New Revision: 34112

URL: http://svn.reactos.org/svn/reactos?rev=34112&view=rev
Log:
Make EnumDeviceDrivers and EnumProcessModules return the total size, even if the buffer is smaller. (from bug 3319)

Modified:
    trunk/reactos/dll/win32/psapi/psapi.c

Modified: trunk/reactos/dll/win32/psapi/psapi.c
URL: http://svn.reactos.org/svn/reactos/trunk/reactos/dll/win32/psapi/psapi.c?rev=34112&r1=34111&r2=34112&view=diff
==============================================================================
--- trunk/reactos/dll/win32/psapi/psapi.c [iso-8859-1] (original)
+++ trunk/reactos/dll/win32/psapi/psapi.c [iso-8859-1] Thu Jun 26 14:52:05 2008
@@ -39,6 +39,7 @@
 {
   LPVOID *lpImageBase;
   DWORD nCount;
+  DWORD nTotal;
 } ENUM_DEVICE_DRIVERS_CONTEXT, *PENUM_DEVICE_DRIVERS_CONTEXT;
 
 NTSTATUS STDCALL
@@ -47,19 +48,18 @@
 {
   PENUM_DEVICE_DRIVERS_CONTEXT Context = (PENUM_DEVICE_DRIVERS_CONTEXT)CallbackContext;
 
-  /* no more buffer space */
-  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--;
-
+  /* check if more buffer space is available */
+  if(Context->nCount != 0)
+  {
+    /* return current module */
+    *Context->lpImageBase = CurrentModule->ImageBase;
+
+    /* go to next array slot */
+    Context->lpImageBase++;
+    Context->nCount--;
+  }
+
+  Context->nTotal++;
   return STATUS_SUCCESS;
 }
 
@@ -97,6 +97,7 @@
 {
   HMODULE *lphModule;
   DWORD nCount;
+  DWORD nTotal;
 } ENUM_PROCESS_MODULES_CONTEXT, *PENUM_PROCESS_MODULES_CONTEXT;
 
 NTSTATUS STDCALL
@@ -106,19 +107,18 @@
 {
   PENUM_PROCESS_MODULES_CONTEXT Context = (PENUM_PROCESS_MODULES_CONTEXT)CallbackContext;
 
-  /* no more buffer space */
-  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--;
-
+  /* check if more buffer space is available */
+  if(Context->nCount != 0)
+  {
+    /* return current process */
+    *Context->lphModule = CurrentModule->DllBase;
+
+    /* go to next array slot */
+    Context->lphModule++;
+    Context->nCount--;
+  }
+
+  Context->nTotal++;
   return STATUS_SUCCESS;
 }
 
@@ -678,23 +678,18 @@
   ENUM_DEVICE_DRIVERS_CONTEXT Context;
   NTSTATUS Status;
 
-  if(cb == 0 || lpImageBase == NULL)
-  {
-    *lpcbNeeded = 0;
-    return TRUE;
-  }
-
   cb /= sizeof(PVOID);
 
   Context.lpImageBase = lpImageBase;
   Context.nCount = cb;
+  Context.nTotal = 0;
 
   Status = PsaEnumerateSystemModules(EnumDeviceDriversCallback, &Context);
 
-  /* return the count of bytes returned */
-  *lpcbNeeded = (cb - Context.nCount) * sizeof(PVOID);
-
-  if(!NT_SUCCESS(Status) && (Status != STATUS_INFO_LENGTH_MISMATCH))
+  /* return the count of bytes that would be needed for a complete enumeration */
+  *lpcbNeeded = Context.nTotal * sizeof(PVOID);
+
+  if(!NT_SUCCESS(Status))
   {
     SetLastErrorByStatus(Status);
     return FALSE;
@@ -757,21 +752,16 @@
 
   cb /= sizeof(HMODULE);
 
-  if(cb == 0 || lphModule == NULL)
-  {
-    *lpcbNeeded = 0;
-    return TRUE;
-  }
-
   Context.lphModule = lphModule;
   Context.nCount = cb;
+  Context.nTotal = 0;
 
   /* enumerate the process modules */
   Status = PsaEnumerateProcessModules(hProcess, EnumProcessModulesCallback, &Context);
 
-  *lpcbNeeded = (cb - Context.nCount) * sizeof(DWORD);
-
-  if(!NT_SUCCESS(Status) && (Status != STATUS_INFO_LENGTH_MISMATCH))
+  *lpcbNeeded = Context.nTotal * sizeof(HMODULE);
+
+  if(!NT_SUCCESS(Status))
   {
     SetLastErrorByStatus(Status);
     return FALSE;