PSAPI almost finished

svn path=/trunk/; revision=3439

reactos/lib/psapi/enum/module.c

@@ -1,53 +1,76 @@
-/* $Id: module.c,v 1.1 2002/06/18 22:12:51 hyperion Exp $
+/* $Id: module.c,v 1.2 2002/08/29 23:57:53 hyperion Exp $
 */
 /*
  * COPYRIGHT:   See COPYING in the top level directory
+ * LICENSE:     See LGPL.txt in the top level directory
  * PROJECT:     ReactOS system libraries
  * FILE:        reactos/lib/psapi/enum/module.c
- * PURPOSE:     Enumerate system modules
+ * PURPOSE:     Enumerate system and process modules
  * PROGRAMMER:  KJK::Hyperion <noog@libero.it>
  * UPDATE HISTORY:
  *              10/06/2002: Created
+ *              29/08/2002: Generalized the interface to improve reusability,
+ *                          more efficient use of memory operations
  */
 
 #include <ddk/ntddk.h>
+#include <debug.h>
 #include <internal/psapi.h>
-#include <stdlib.h>
+#include <ntdll/ldr.h>
 
 NTSTATUS
 STDCALL
 PsaEnumerateSystemModules
 (
- OUT PVOID * Modules,
+ IN PSYSMOD_ENUM_ROUTINE Callback,
- IN ULONG ModulesLength,
+ IN OUT PVOID CallbackContext,
- OUT ULONG * ReturnLength OPTIONAL
+ IN OUT PVOID AllocatorContext
 )
 {
- NTSTATUS nErrCode = STATUS_SUCCESS;
+ ULONG nSize;
- ULONG nSize = sizeof(ULONG);
+ register NTSTATUS nErrCode = STATUS_SUCCESS;
- PULONG pnModuleCount = NULL;
+ register PULONG pnModuleCount = &nSize;
- PSYSTEM_MODULE_ENTRY psmeCurModule;
+ register PSYSTEM_MODULE_ENTRY psmeCurModule;
- ULONG nBufSize;
+ register ULONG nModuleCount;
 
- /* ignore buffer size if buffer is null */
+ /* initial probe */
- if(Modules == NULL)
+ nErrCode = NtQuerySystemInformation
-  ModulesLength = 0;
+ (
- /* ignore buffer if buffer size is zero */
+  SystemModuleInformation,
- else if(ModulesLength == 0)
+  pnModuleCount,
-  Modules = NULL;
+  sizeof(nSize),
+  NULL
+ );
+ 
+ if(nErrCode != STATUS_INFO_LENGTH_MISMATCH && !NT_SUCCESS(nErrCode))
+ {
+  /* failure */
+  DPRINT(FAILED_WITH_STATUS, "NtQuerySystemInformation", nErrCode);
+  return nErrCode;
+ }
+
+ /* RATIONALE: the loading of a system module is a rare occurrence. To minimize
+    memory operations that could be expensive, or fragment the pool/heap, we try
+    to determine the buffer size in advance, knowing that the number of elements
+    is unlikely to change */
+ nSize = sizeof(ULONG) + nSize * sizeof(SYSTEM_MODULE_ENTRY);
+ pnModuleCount = NULL;
  
  do
  {
-  void * pTmp;
+  register void * pTmp;
 
-  /* resize and/or move the buffer */
+  /* free the buffer, and reallocate it to the new size. RATIONALE: since we
-  pTmp = realloc(pnModuleCount, nSize);
+     ignore the buffer's content at this point, there's no point in a realloc(),
+     that could end up copying a large chunk of data we'd discard anyway */
+  PsaFree(AllocatorContext, pnModuleCount);
+  pTmp = PsaMalloc(AllocatorContext, nSize);
   
   if(pTmp == NULL)
   {
    /* failure */
    nErrCode = STATUS_NO_MEMORY;
-   goto end;
+   goto esm_Finalize;
   }
   
   pnModuleCount = pTmp;
@@ -61,48 +84,192 @@ PsaEnumerateSystemModules
    NULL
   );
 
-  /* while this is less efficient than doubling the buffer size, it should be
+  /* double the buffer for the next loop */
-     executed only once in most cases */
+  nSize += nSize;
-  nSize += sizeof(SYSTEM_MODULE_ENTRY) * (*pnModuleCount);
  }
  /* repeat until the buffer is big enough */
  while(nErrCode == STATUS_INFO_LENGTH_MISMATCH);
 
+ if(!NT_SUCCESS(nErrCode))
+ {
+  /* failure */
+  DPRINT(FAILED_WITH_STATUS, "NtQuerySystemInformation", nErrCode);
+  goto esm_Finalize;
+ }
+
  /* the array of modules starts right after an ULONG storing their count */
  psmeCurModule = (PSYSTEM_MODULE_ENTRY)(pnModuleCount + 1);
 
- /* element count */
+ nModuleCount = *pnModuleCount;
- nBufSize = ModulesLength / sizeof(*Modules);
 
- /* not enough elements in the buffer */
+ /* repeat until all modules have been returned */
- if((*pnModuleCount) > nBufSize)
+ while(nModuleCount > 0)
-  nErrCode = STATUS_INFO_LENGTH_MISMATCH;
- /* too many elements in the buffer */
- else
-  nBufSize = *pnModuleCount;
- 
- /* return the needed buffer size */
- if(ReturnLength)
-  (*ReturnLength) = (*pnModuleCount) * sizeof(*Modules);
- 
- /* repeat until the buffer is empty or all modules have been returned */
- while(nBufSize > 0)
  {
-  /* return current module base */
+  /* return current module to the callback */
-  (*Modules) = psmeCurModule->BaseAddress;
+  nErrCode = Callback(nModuleCount, psmeCurModule, CallbackContext);
   
-  /* next buffer element */
+  if(!NT_SUCCESS(nErrCode))
-  Modules ++;
+   /* failure */
-  nBufSize --;
+   goto esm_Finalize;
   
   /* next module */
   psmeCurModule ++;
+  nModuleCount --;
  }
 
-end:
+esm_Finalize:
- free(pnModuleCount);
+ /* free the buffer */
+ PsaFree(AllocatorContext, pnModuleCount);
+
+ return (nErrCode);
+}
+
+NTSTATUS
+STDCALL
+PsaEnumerateProcessModules
+(
+ IN HANDLE ProcessHandle,
+ IN PPROCMOD_ENUM_ROUTINE Callback,
+ IN OUT PVOID CallbackContext
+)
+{
+ register NTSTATUS nErrCode;
+
+ /* current process - use direct memory copy */
+ if(ProcessHandle == NtCurrentProcess())
+ {
+  register PLIST_ENTRY pleListHead;
+  register PLIST_ENTRY pleCurEntry;
+
+#if 0
+  /* FIXME: activate this when GCC supports SEH */
+  __try
+  {
+#endif
+   pleListHead = &(NtCurrentPeb()->Ldr->InLoadOrderModuleList);
+   pleCurEntry = pleListHead->Flink;
+ 
+   while(pleCurEntry != pleListHead)
+   {
+    register PLDR_MODULE plmModule = CONTAINING_RECORD
+    (
+     pleCurEntry,
+     LDR_MODULE,
+     InLoadOrderModuleList
+    );
+   
+    /* return the current module to the callback */
+    nErrCode = Callback(ProcessHandle, plmModule, CallbackContext);
+    
+    if(!NT_SUCCESS(nErrCode))
+     /* failure */
+     goto epm_Failure;
+    
+    pleCurEntry = plmModule->InLoadOrderModuleList.Flink;
+   }
+#if 0
+  /* FIXME: activate this when GCC supports SEH */
+  }
+  __except(EXCEPTION_EXECUTE_HANDLER)
+  {
+   return GetExceptionCode();
+  }
+#endif
+ }
+ /* another process */
+ else
+ {
+  PROCESS_BASIC_INFORMATION pbiInfo;
+  PPEB_LDR_DATA ppldLdrData;
+  LDR_MODULE lmModule;
+  PLIST_ENTRY pleListHead;
+  PLIST_ENTRY pleCurEntry;
+ 
+  /* query the process basic information (includes the PEB address) */
+  nErrCode = NtQueryInformationProcess
+  (
+   ProcessHandle,
+   ProcessBasicInformation,
+   &pbiInfo,
+   sizeof(pbiInfo),
+   NULL
+  );
+ 
+  if(!NT_SUCCESS(nErrCode))
+  {
+   /* failure */
+   DPRINT(FAILED_WITH_STATUS, "NtQueryInformationProcess", nErrCode);
+   goto epm_Failure;
+  }
+ 
+  /* get the address of the PE Loader data */
+  nErrCode = NtReadVirtualMemory
+  (
+   ProcessHandle,
+   &(pbiInfo.PebBaseAddress->Ldr),
+   &ppldLdrData,
+   sizeof(ppldLdrData),
+   NULL
+  );
+ 
+  if(!NT_SUCCESS(nErrCode))
+  {
+   /* failure */
+   DPRINT(FAILED_WITH_STATUS, "NtReadVirtualMemory", nErrCode);
+   goto epm_Failure;
+  }
+ 
+  /* head of the module list: the last element in the list will point to this */
+  pleListHead = &ppldLdrData->InLoadOrderModuleList;
+ 
+  /* get the address of the first element in the list */
+  nErrCode = NtReadVirtualMemory
+  (
+   ProcessHandle,
+   &(ppldLdrData->InLoadOrderModuleList.Flink),
+   &pleCurEntry,
+   sizeof(pleCurEntry),
+   NULL
+  );
+ 
+  while(pleCurEntry != pleListHead)
+  {
+   /* read the current module */
+   nErrCode = NtReadVirtualMemory
+   (
+    ProcessHandle,
+    CONTAINING_RECORD(pleCurEntry, LDR_MODULE, InLoadOrderModuleList),
+    &lmModule,
+    sizeof(lmModule),
+    NULL
+   );
+ 
+   if(!NT_SUCCESS(nErrCode))
+   {
+    /* failure */
+    DPRINT(FAILED_WITH_STATUS, "NtReadVirtualMemory", nErrCode);
+    goto epm_Failure;
+   }
+
+   /* return the current module to the callback */
+   nErrCode = Callback(ProcessHandle, &lmModule, CallbackContext);
+   
+   if(!NT_SUCCESS(nErrCode))
+    /* failure */
+    goto epm_Failure;
+    
+   /* address of the next module in the list */
+   pleCurEntry = lmModule.InLoadOrderModuleList.Flink;
+  }
+ 
+ }
+
+ /* success */
+ return (STATUS_SUCCESS);
+
+epm_Failure:
+ /* failure */
  return (nErrCode);
 }
 
 /* EOF */
-

reactos/lib/psapi/enum/process.c

@@ -1,54 +1,58 @@
-/* $Id: process.c,v 1.1 2002/06/18 22:12:51 hyperion Exp $
+/* $Id: process.c,v 1.2 2002/08/29 23:57:53 hyperion Exp $
 */
 /*
  * COPYRIGHT:   See COPYING in the top level directory
+ * LICENSE:     See LGPL.txt in the top level directory
  * PROJECT:     ReactOS system libraries
  * FILE:        reactos/lib/psapi/enum/process.c
- * PURPOSE:     Enumerate process ids
+ * PURPOSE:     Enumerate processes
  * PROGRAMMER:  KJK::Hyperion <noog@libero.it>
  * UPDATE HISTORY:
  *              10/06/2002: Created
+ *              29/08/2002: Generalized the interface to improve reusability,
+ *                          more efficient use of memory operations
  */
 
-#include <ddk/ntddk.h>
-#include <internal/psapi.h>
 #include <stdlib.h>
+#include <ddk/ntddk.h>
+#include <debug.h>
+#include <internal/psapi.h>
 
 NTSTATUS
 STDCALL
-PsaEnumerateProcessIds
+PsaEnumerateProcesses
 (
- OUT ULONG * ProcessIds,
+ IN PPROC_ENUM_ROUTINE Callback,
- IN ULONG ProcessIdsLength,
+ IN OUT PVOID CallbackContext,
- OUT ULONG * ReturnLength OPTIONAL
+ IN OUT PVOID AllocatorContext
 )
 {
- NTSTATUS nErrCode = STATUS_SUCCESS;
+ register NTSTATUS nErrCode = STATUS_SUCCESS;
+ PSYSTEM_PROCESS_INFORMATION pInfoBuffer = NULL;
+ PSYSTEM_PROCESS_INFORMATION pInfoHead = NULL;
  ULONG nSize = 32768;
- SYSTEM_PROCESS_INFORMATION * pInfoBuffer = NULL;
- SYSTEM_PROCESS_INFORMATION * pInfoHead = NULL;
- ULONG nBufSize;
- ULONG nRetLen = 0;
-
- /* ignore buffer size if buffer is null */
- if(ProcessIds == NULL)
-  ProcessIdsLength = 0;
- /* ignore buffer if buffer size is zero */
- else if(ProcessIdsLength == 0)
-  ProcessIds = NULL;
 
+ /* FIXME: if the system has loaded several processes and threads, the buffer
+    could get really big. But if there's several processes and threads, the
+    system is already under stress, and a huge buffer could only make things
+    worse. The function should be profiled to see what's the average minimum
+    buffer size, to succeed on the first shot */
  do
  {
   void * pTmp;
   
-  /* resize and/or move the buffer */
+  /* free the buffer, and reallocate it to the new size. RATIONALE: since we
-  pTmp = realloc(pInfoBuffer, nSize);
+     ignore the buffer's contents at this point, there's no point in a realloc()
+     that could end up copying a large chunk of data we'd discard anyway */
+  PsaFree(AllocatorContext, pInfoBuffer);
+  pTmp = PsaMalloc(AllocatorContext, nSize);
   
   if(pTmp == NULL)
   {
    /* failure */
+   DPRINT(FAILED_WITH_STATUS, "malloc", STATUS_NO_MEMORY);
    nErrCode = STATUS_NO_MEMORY;
-   goto end;
+   goto esp_Finalize;
   }
   
   pInfoBuffer = pTmp;
@@ -70,26 +74,23 @@ PsaEnumerateProcessIds
  
  /* failure */
  if(!NT_SUCCESS(nErrCode))
-  goto end;
+ {
+  DPRINT(FAILED_WITH_STATUS, "NtQuerySystemInformation", nErrCode);
+  goto esp_Finalize;
+ }
  
- /* size of ProcessIds in elements */
- nBufSize = ProcessIdsLength / sizeof(*ProcessIds);
  /* list head */
  pInfoHead = pInfoBuffer;
 
- /* repeat until the buffer is empty */
+ /* scan the list */
- while(nBufSize > 0)
+ while(1)
  {
-  /* return the current process id */
+  /* notify the callback */
-  (*ProcessIds) = pInfoHead->ProcessId;
+  nErrCode = Callback(pInfoHead, CallbackContext);
 
-  /* move to the next buffer entry */
+  /* if the callback returned an error or this is the end of the process list,
-  ProcessIds ++;
+     break out */
-  nBufSize --;
+  if(!NT_SUCCESS(nErrCode) || pInfoHead->RelativeOffset == 0)
-  nRetLen ++;
-  
-  /* end of process list */
-  if(pInfoHead->RelativeOffset == 0)
    break;
   
   /* move to the next process */
@@ -98,37 +99,9 @@ PsaEnumerateProcessIds
    ((ULONG)pInfoHead + pInfoHead->RelativeOffset);
  }
 
- if(pInfoHead->RelativeOffset == 0)
+esp_Finalize:
-  /* all process ids were returned */
-  nErrCode = STATUS_SUCCESS;
- else
- {
-  /* insufficient buffer */
-  nErrCode = STATUS_INFO_LENGTH_MISMATCH;
-  
-  /* caller doesn't need buffer size */
-  if(ReturnLength == NULL)
-   goto end;
-  
-  /* repeat while there are still processes */
-  while(pInfoHead->RelativeOffset != 0)
-  {
-   pInfoHead =
-    (SYSTEM_PROCESS_INFORMATION*)
-    ((ULONG)pInfoHead + pInfoHead->RelativeOffset);
-
-   nRetLen++;
-  }
-  
- }
-
- /* used buffer size */
- if(ReturnLength)
-  (*ReturnLength) = nRetLen * sizeof(DWORD);
- 
-end:
  /* free the buffer */
- free(pInfoBuffer);
+ PsaFree(AllocatorContext, pInfoBuffer);
  
  /* return the last status */
  return (nErrCode);

reactos/lib/psapi/include/internal/psapi.h

@@ -1,9 +1,9 @@
-/* $Id: psapi.h,v 1.1 2002/06/18 22:15:57 hyperion Exp $
+/* $Id: psapi.h,v 1.2 2002/08/29 23:57:54 hyperion Exp $
 */
 /*
  * internal/psapi.h
  *
- * Process Status Helper API
+ * Process Status Helper API, native interface
  *
  * This file is part of the ReactOS Operating System.
  *
@@ -27,33 +27,69 @@
 
 /* INCLUDES */
 #include <ddk/ntddk.h>
+#include <ntdll/ldr.h>
 
 /* OBJECTS */
 
 /* TYPES */
+typedef NTSTATUS STDCALL (*PPROC_ENUM_ROUTINE)
+(
+ IN PSYSTEM_PROCESS_INFORMATION CurrentProcess,
+ IN OUT PVOID CallbackContext
+);
+
+typedef NTSTATUS STDCALL (*PSYSMOD_ENUM_ROUTINE)
+(
+ IN ULONG ModuleCount,
+ IN PSYSTEM_MODULE_ENTRY CurrentModule,
+ IN OUT PVOID CallbackContext
+);
+
+typedef NTSTATUS STDCALL (*PPROCMOD_ENUM_ROUTINE)
+(
+ IN HANDLE ProcessHandle,
+ IN PLDR_MODULE CurrentModule,
+ IN OUT PVOID CallbackContext
+);
 
 /* CONSTANTS */
+#define FAILED_WITH_STATUS DEFINE_DBG_MSG("%s() failed, status 0x%08X")
 
 /* PROTOTYPES */
 NTSTATUS
 STDCALL
-PsaEnumerateProcessIds
+PsaEnumerateProcesses
 (
- OUT ULONG * ProcessIds,
+ IN PPROC_ENUM_ROUTINE Callback,
- IN ULONG ProcessIdsLength,
+ IN OUT PVOID CallbackContext,
- OUT ULONG * ReturnLength OPTIONAL
+ IN OUT PVOID AllocatorContext
 );
 
 NTSTATUS
 STDCALL
 PsaEnumerateSystemModules
 (
- OUT PVOID * Modules,
+ IN PSYSMOD_ENUM_ROUTINE Callback,
- IN ULONG ModulesLength,
+ IN OUT PVOID CallbackContext,
- OUT ULONG * ReturnLength OPTIONAL
+ IN OUT PVOID AllocatorContext
 );
 
+NTSTATUS
+STDCALL
+PsaEnumerateProcessModules
+(
+ IN HANDLE ProcessHandle,
+ IN PPROCMOD_ENUM_ROUTINE Callback,
+ IN OUT PVOID CallbackContext
+);
+
+/* the user must provide these */
+PVOID STDCALL PsaMalloc(IN OUT PVOID Context, IN ULONG Size);
+PVOID STDCALL PsaRealloc(IN OUT PVOID Context, IN PVOID Addr, IN ULONG Size);
+VOID STDCALL PsaFree(IN OUT PVOID Context, IN PVOID Addr);
+
 /* MACROS */
+#define DEFINE_DBG_MSG(__str__) "PSAPI: " __str__ "\n"
 
 #endif /* __INTERNAL_PSAPI_H_INCLUDED__ */
 

reactos/lib/psapi/misc/dllmain.c

@@ -1,8 +1,9 @@
-/* $Id: dllmain.c,v 1.2 2002/06/18 22:15:58 hyperion Exp $
+/* $Id: dllmain.c,v 1.3 2002/08/29 23:57:54 hyperion Exp $
  * 
  * ReactOS PSAPI.DLL
  */
 #include <windows.h>
+#include <ntdll/ldr.h>
 
 BOOLEAN STDCALL DllMain
 (
@@ -11,6 +12,12 @@ BOOLEAN STDCALL DllMain
  PVOID reserved
 )
 {
+ if(dwReason == DLL_PROCESS_ATTACH)
+  /* don't bother calling the entry point on thread startup - PSAPI.DLL doesn't
+     store any per-thread data */
+  LdrDisableThreadCalloutsForDll(hinstDll);
+
  return (TRUE);
 }
+
 /* EOF */

reactos/lib/psapi/misc/malloc.c

@@ -1,3 +1,5 @@
+/* $Id: malloc.c,v 1.2 2002/08/29 23:57:54 hyperion Exp $
+ */
 /*
  * COPYRIGHT:   See COPYING in the top level directory
  * PROJECT:     ReactOS system libraries

reactos/lib/psapi/misc/stubs.c

@@ -1,4 +1,4 @@
-/* $Id: stubs.c,v 1.2 2002/06/18 22:15:58 hyperion Exp $ */
+/* $Id: stubs.c,v 1.3 2002/08/29 23:57:54 hyperion Exp $ */
 #include <windows.h>
 #include <psapi.h>
 
@@ -12,17 +12,6 @@ BOOL STDCALL EnumPageFiles(
  return FALSE;
 }
 #endif
-BOOL STDCALL EnumProcessModules(
-  HANDLE hProcess,      // handle to process
-  HMODULE *lphModule,   // array of module handles
-  DWORD cb,             // size of array
-  LPDWORD lpcbNeeded    // number of bytes required
-)
-{
- SetLastError(ERROR_INVALID_FUNCTION);
- return FALSE;
-}
-
 DWORD STDCALL GetDeviceDriverBaseNameA(
   LPVOID ImageBase,  // driver load address
   LPSTR lpBaseName,  // driver base name buffer

reactos/lib/psapi/misc/win32.c

@@ -1,19 +1,38 @@
-/* $Id: win32.c,v 1.1 2002/06/18 22:14:07 hyperion Exp $
+/* $Id: win32.c,v 1.2 2002/08/29 23:57:54 hyperion Exp $
  */
 /*
  * COPYRIGHT:   See COPYING in the top level directory
  * PROJECT:     ReactOS system libraries
  * FILE:        reactos/lib/psapi/misc/win32.c
- * PURPOSE:     Win32 stubs for PSAPI
+ * PURPOSE:     Win32 interfaces for PSAPI
  * PROGRAMMER:  KJK::Hyperion <noog@libero.it>
  * UPDATE HISTORY:
  *              10/06/2002: Created
  */
 
 #include <windows.h>
-#include <internal/psapi.h>
 #include <psapi.h>
+#include <stdlib.h>
+#include <ddk/ntddk.h>
+#include <internal/psapi.h>
 
+/* Memory allocators for PSAPI */
+PVOID STDCALL PsaMalloc(IN OUT PVOID Context, IN ULONG Size)
+{
+ return malloc(Size);
+}
+
+PVOID STDCALL PsaRealloc(IN OUT PVOID Context, IN PVOID Addr, IN ULONG Size)
+{
+ return realloc(Addr, Size);
+}
+
+VOID STDCALL PsaFree(IN OUT PVOID Context, IN PVOID Addr)
+{
+ free(Addr);
+}
+
+/* EmptyWorkingSet */
 BOOL STDCALL EmptyWorkingSet(HANDLE hProcess)
 {
  NTSTATUS nErrCode;
@@ -56,6 +75,40 @@ fail:
  return (FALSE);
 }
 
+/* EnumDeviceDrivers */
+/* callback context */
+typedef struct _ENUM_DEVICE_DRIVERS_CONTEXT
+{
+ LPVOID *lpImageBase;
+ DWORD nCount;
+} ENUM_DEVICE_DRIVERS_CONTEXT, *PENUM_DEVICE_DRIVERS_CONTEXT;
+
+/* callback routine */
+NTSTATUS STDCALL EnumDeviceDriversCallback
+(
+ IN ULONG ModuleCount,
+ IN PSYSTEM_MODULE_ENTRY CurrentModule,
+ IN OUT PVOID CallbackContext
+)
+{
+ register PENUM_DEVICE_DRIVERS_CONTEXT peddcContext =
+  (PENUM_DEVICE_DRIVERS_CONTEXT)CallbackContext;
+
+ /* no more buffer space */
+ if(peddcContext->nCount == 0)
+  return STATUS_INFO_LENGTH_MISMATCH;
+
+ /* return current module */
+ *(peddcContext->lpImageBase) = CurrentModule->BaseAddress;
+ 
+ /* go to next array slot */
+ (peddcContext->lpImageBase) ++;
+ (peddcContext->nCount) --;
+ 
+ return STATUS_SUCCESS;
+}
+
+/* exported interface */
 BOOL STDCALL EnumDeviceDrivers
 (
  LPVOID *lpImageBase, 
@@ -63,29 +116,32 @@ BOOL STDCALL EnumDeviceDrivers
  LPDWORD lpcbNeeded   
 )
 {
- NTSTATUS nErrCode;
+ register NTSTATUS nErrCode;
+ ENUM_DEVICE_DRIVERS_CONTEXT eddcContext = {lpImageBase, cb / sizeof(PVOID)};
+
+ cb /= sizeof(PVOID);
 
  /* do nothing if the buffer is empty */
- if(cb < sizeof(DWORD) || lpImageBase == NULL)
+ if(cb == 0 || lpImageBase == NULL)
  {
   *lpcbNeeded = 0;
   return (TRUE);
  }
  
  /* enumerate the system modules */
- nErrCode = PsaEnumerateSystemModules(lpImageBase, cb, lpcbNeeded);
+ nErrCode = PsaEnumerateSystemModules
+ (
+  &EnumDeviceDriversCallback,
+  &eddcContext,
+  NULL
+ );
+ 
+ /* return the count of bytes returned */
+ *lpcbNeeded = (cb - eddcContext.nCount) * sizeof(PVOID);
 
  /* success */
- if(NT_SUCCESS(nErrCode))
+ if(NT_SUCCESS(nErrCode) || nErrCode == STATUS_INFO_LENGTH_MISMATCH)
   return (TRUE);
- 
- /* failure or partial success */
- if(nErrCode == STATUS_INFO_LENGTH_MISMATCH)
- {
-  /* insufficient buffer: ignore this error */
-  *lpcbNeeded = cb;
-  return (TRUE);
- }
  else
  {
   /* failure */
@@ -94,6 +150,39 @@ BOOL STDCALL EnumDeviceDrivers
  }
 }
 
+/* EnumProcesses */
+/* callback context */
+typedef struct _ENUM_PROCESSES_CONTEXT
+{
+ DWORD *lpidProcess;
+ DWORD nCount;
+} ENUM_PROCESSES_CONTEXT, *PENUM_PROCESSES_CONTEXT;
+
+/* callback routine */
+NTSTATUS STDCALL EnumProcessesCallback
+(
+ IN PSYSTEM_PROCESS_INFORMATION CurrentProcess,
+ IN OUT PVOID CallbackContext
+)
+{
+ register PENUM_PROCESSES_CONTEXT pepcContext =
+  (PENUM_PROCESSES_CONTEXT)CallbackContext;
+
+ /* no more buffer space */
+ if(pepcContext->nCount == 0)
+  return STATUS_INFO_LENGTH_MISMATCH;
+
+ /* return current process */
+ *(pepcContext->lpidProcess) = CurrentProcess->ProcessId;
+ 
+ /* go to next array slot */
+ (pepcContext->lpidProcess) ++;
+ (pepcContext->nCount) --;
+ 
+ return STATUS_SUCCESS;
+}
+
+/* exported interface */
 BOOL STDCALL EnumProcesses
 (
  DWORD *lpidProcess,
@@ -101,29 +190,100 @@ BOOL STDCALL EnumProcesses
  LPDWORD lpcbNeeded
 )
 {
- NTSTATUS nErrCode;
+ register NTSTATUS nErrCode;
+ ENUM_PROCESSES_CONTEXT epcContext = {lpidProcess, cb / sizeof(DWORD)};
+ 
+ cb /= sizeof(DWORD);
 
  /* do nothing if the buffer is empty */
- if(cb < sizeof(DWORD) || lpidProcess == NULL)
+ if(cb == 0 || lpidProcess == NULL)
  {
   *lpcbNeeded = 0;
   return (TRUE);
  }
  
  /* enumerate the process ids */
- nErrCode = PsaEnumerateProcessIds(lpidProcess, cb, lpcbNeeded);
+ nErrCode = PsaEnumerateProcesses(&EnumProcessesCallback, &epcContext, NULL);
+ 
+ *lpcbNeeded = (cb - epcContext.nCount) * sizeof(DWORD);
 
  /* success */
- if(NT_SUCCESS(nErrCode))
+ if(NT_SUCCESS(nErrCode) || nErrCode == STATUS_INFO_LENGTH_MISMATCH)
   return (TRUE);
- 
+ else
- /* failure or partial success */
- if(nErrCode == STATUS_INFO_LENGTH_MISMATCH)
  {
-  /* insufficient buffer: ignore this error */
+  /* failure */
-  *lpcbNeeded = cb;
+  SetLastError(RtlNtStatusToDosError(nErrCode));
+  return (FALSE);
+ }
+}
+
+/* EnumProcessModules */
+/* callback context */
+typedef struct _ENUM_PROCESS_MODULES_CONTEXT
+{
+ HMODULE *lphModule;
+ DWORD nCount;
+} ENUM_PROCESS_MODULES_CONTEXT, *PENUM_PROCESS_MODULES_CONTEXT;
+
+/* callback routine */
+NTSTATUS STDCALL EnumProcessModulesCallback
+(
+ IN HANDLE ProcessHandle,
+ IN PLDR_MODULE CurrentModule,
+ IN OUT PVOID CallbackContext
+)
+{
+ register PENUM_PROCESS_MODULES_CONTEXT pepmcContext = 
+  (PENUM_PROCESS_MODULES_CONTEXT)CallbackContext;
+
+ /* no more buffer space */
+ if(pepmcContext->nCount == 0)
+  return STATUS_INFO_LENGTH_MISMATCH;
+
+ /* return current process */
+ *(pepmcContext->lphModule) = CurrentModule->BaseAddress;
+ 
+ /* go to next array slot */
+ (pepmcContext->lphModule) ++;
+ (pepmcContext->nCount) --;
+ 
+ return STATUS_SUCCESS;
+}
+
+/* exported interface */
+BOOL STDCALL EnumProcessModules(
+  HANDLE hProcess,      // handle to process
+  HMODULE *lphModule,   // array of module handles
+  DWORD cb,             // size of array
+  LPDWORD lpcbNeeded    // number of bytes required
+)
+{
+ register NTSTATUS nErrCode;
+ ENUM_PROCESS_MODULES_CONTEXT epmcContext = {lphModule, cb / sizeof(HMODULE)};
+ 
+ cb /= sizeof(DWORD);
+
+ /* do nothing if the buffer is empty */
+ if(cb == 0 || lphModule == NULL)
+ {
+  *lpcbNeeded = 0;
   return (TRUE);
  }
+ 
+ /* enumerate the process modules */
+ nErrCode = PsaEnumerateProcessModules
+ (
+  hProcess,
+  &EnumProcessModulesCallback,
+  &epmcContext
+ );
+ 
+ *lpcbNeeded = (cb - epmcContext.nCount) * sizeof(DWORD);
+
+ /* success */
+ if(NT_SUCCESS(nErrCode) || nErrCode == STATUS_INFO_LENGTH_MISMATCH)
+  return (TRUE);
  else
  {
   /* failure */