/* $Id: toolhelp.c 27692 2007-07-16 12:21:38Z weiden $
 *
 * COPYRIGHT:       See COPYING in the top level directory
 * PROJECT:         ReactOS system libraries
 * FILE:            dll/win32/kernel32/misc/toolhelp.c
 * PURPOSE:         Toolhelp functions
 * PROGRAMMER:      Thomas Weidenmueller <w3seek@reactos.com>
 *                  Robert Dickenson (robd@mok.lvcm.com)
 *
 * NOTES:           Do NOT use the heap functions in here because they
 *                  adulterate the heap statistics!
 *
 * UPDATE HISTORY:
 *                  10/30/2004 Implemented some parts (w3)
 *                             Inspired by the book "Windows NT Native API"
 *                  Created 05 January 2003 (robd)
 */

#include <k32.h>

#include <tlhelp32.h>

#define NDEBUG
#include "../include/debug.h"

/* INTERNAL DEFINITIONS *******************************************************/

typedef struct _RTLP_HEAP_ENTRY
{
    ULONG Size;
    USHORT Flags;
    USHORT Unknown1; /* FIXME */
    ULONG Unknown2; /* FIXME */
    PVOID Address;
} RTLP_HEAP_ENTRY, *PRTLP_HEAP_ENTRY;

#define CHECK_PARAM_SIZE(ptr, siz)                                             \
  if((ptr) == NULL || (ptr)->dwSize != (siz))                                  \
  {                                                                            \
    SetLastError(ERROR_INVALID_PARAMETER);                                     \
    return FALSE;                                                              \
  }

/*
 * Tests in win showed that the dwSize field can be greater than the actual size
 * of the structure for the ansi functions. I found this out by accidently
 * forgetting to set the dwSize field in a test application and it just didn't
 * work in ros but in win.
 */

#define CHECK_PARAM_SIZEA(ptr, siz)                                            \
  if((ptr) == NULL || (ptr)->dwSize < (siz))                                   \
  {                                                                            \
    SetLastError(ERROR_INVALID_PARAMETER);                                     \
    return FALSE;                                                              \
  }

#define OffsetToPtr(Snapshot, Offset)                                          \
  ((ULONG_PTR)((Snapshot) + 1) + (ULONG_PTR)(Offset))

typedef struct _TH32SNAPSHOT
{
  /* Heap list */
  ULONG HeapListCount;
  ULONG HeapListIndex;
  ULONG_PTR HeapListOffset;
  /* Module list */
  ULONG ModuleListCount;
  ULONG ModuleListIndex;
  ULONG_PTR ModuleListOffset;
  /* Process list */
  ULONG ProcessListCount;
  ULONG ProcessListIndex;
  ULONG_PTR ProcessListOffset;
  /* Thread list */
  ULONG ThreadListCount;
  ULONG ThreadListIndex;
  ULONG_PTR ThreadListOffset;
} TH32SNAPSHOT, *PTH32SNAPSHOT;

/* INTERNAL FUNCTIONS *********************************************************/

static VOID
TH32FreeAllocatedResources(PRTL_DEBUG_INFORMATION HeapDebug,
                           PRTL_DEBUG_INFORMATION ModuleDebug,
                           PVOID ProcThrdInfo,
                           SIZE_T ProcThrdInfoSize)
{
  if(HeapDebug != NULL)
  {
    RtlDestroyQueryDebugBuffer(HeapDebug);
  }
  if(ModuleDebug != NULL)
  {
    RtlDestroyQueryDebugBuffer(ModuleDebug);
  }

  if(ProcThrdInfo != NULL)
  {
    NtFreeVirtualMemory(NtCurrentProcess(),
                        ProcThrdInfo,
                        &ProcThrdInfoSize,
                        MEM_RELEASE);
  }
}

static NTSTATUS
TH32CreateSnapshot(DWORD dwFlags,
                   DWORD th32ProcessID,
                   PRTL_DEBUG_INFORMATION *HeapDebug,
                   PRTL_DEBUG_INFORMATION *ModuleDebug,
                   PVOID *ProcThrdInfo,
                   SIZE_T *ProcThrdInfoSize)
{
  NTSTATUS Status = STATUS_SUCCESS;

  *HeapDebug = NULL;
  *ModuleDebug = NULL;
  *ProcThrdInfo = NULL;
  *ProcThrdInfoSize = 0;

  /*
   * Allocate the debug information for a heap snapshot
   */
  if(dwFlags & TH32CS_SNAPHEAPLIST)
  {
    *HeapDebug = RtlCreateQueryDebugBuffer(0, FALSE);
    if(*HeapDebug != NULL)
    {
      Status = RtlQueryProcessDebugInformation(th32ProcessID,
                                               RTL_DEBUG_QUERY_HEAPS,
                                               *HeapDebug);
    }
    else
      Status = STATUS_UNSUCCESSFUL;
  }

  /*
   * Allocate the debug information for a module snapshot
   */
  if(dwFlags & TH32CS_SNAPMODULE &&
     NT_SUCCESS(Status))
  {
    *ModuleDebug = RtlCreateQueryDebugBuffer(0, FALSE);
    if(*ModuleDebug != NULL)
    {
      Status = RtlQueryProcessDebugInformation(th32ProcessID,
                                               RTL_DEBUG_QUERY_MODULES,
                                               *ModuleDebug);
    }
    else
      Status = STATUS_UNSUCCESSFUL;
  }

  /*
   * Allocate enough memory for the system's process list
   */

  if(dwFlags & (TH32CS_SNAPPROCESS | TH32CS_SNAPTHREAD) &&
     NT_SUCCESS(Status))
  {
    for(;;)
    {
      (*ProcThrdInfoSize) += 0x10000;
      Status = NtAllocateVirtualMemory(NtCurrentProcess(),
                                       ProcThrdInfo,
                                       0,
                                       ProcThrdInfoSize,
                                       MEM_COMMIT,
                                       PAGE_READWRITE);
      if(!NT_SUCCESS(Status))
      {
        break;
      }

      Status = NtQuerySystemInformation(SystemProcessInformation,
                                        *ProcThrdInfo,
                                        *ProcThrdInfoSize,
                                        NULL);
      if(Status == STATUS_INFO_LENGTH_MISMATCH)
      {
        NtFreeVirtualMemory(NtCurrentProcess(),
                            ProcThrdInfo,
                            ProcThrdInfoSize,
                            MEM_RELEASE);
        *ProcThrdInfo = NULL;
      }
      else
      {
        break;
      }
    }
  }

  /*
   * Free resources in case of failure!
   */

  if(!NT_SUCCESS(Status))
  {
    TH32FreeAllocatedResources(*HeapDebug,
                               *ModuleDebug,
                               *ProcThrdInfo,
                               *ProcThrdInfoSize);
  }

  return Status;
}

static NTSTATUS
TH32CreateSnapshotSectionInitialize(DWORD dwFlags,
                                    DWORD th32ProcessID,
                                    PRTL_DEBUG_INFORMATION HeapDebug,
                                    PRTL_DEBUG_INFORMATION ModuleDebug,
                                    PVOID ProcThrdInfo,
                                    HANDLE *SectionHandle)
{
  PSYSTEM_PROCESS_INFORMATION ProcessInfo;
  LPHEAPLIST32 HeapListEntry;
  LPMODULEENTRY32W ModuleListEntry;
  LPPROCESSENTRY32W ProcessListEntry;
  LPTHREADENTRY32 ThreadListEntry;
  OBJECT_ATTRIBUTES ObjectAttributes;
  LARGE_INTEGER SSize, SOffset;
  HANDLE hSection;
  PTH32SNAPSHOT Snapshot;
  ULONG_PTR DataOffset;
  ULONG ViewSize, i;
  ULONG nProcesses = 0, nThreads = 0, nHeaps = 0, nModules = 0;
  ULONG RequiredSnapshotSize = sizeof(TH32SNAPSHOT);
  PRTL_PROCESS_HEAPS hi = NULL;
  PRTL_PROCESS_MODULES mi = NULL;
  NTSTATUS Status = STATUS_SUCCESS;

  /*
   * Determine the required size for the heap snapshot
   */
  if(dwFlags & TH32CS_SNAPHEAPLIST)
  {
    hi = (PRTL_PROCESS_HEAPS)HeapDebug->Heaps;
    nHeaps = hi->NumberOfHeaps;
    RequiredSnapshotSize += nHeaps * sizeof(HEAPLIST32);
  }

  /*
   * Determine the required size for the module snapshot
   */
  if(dwFlags & TH32CS_SNAPMODULE)
  {
    mi = (PRTL_PROCESS_MODULES)ModuleDebug->Modules;
    nModules = mi->NumberOfModules;
    RequiredSnapshotSize += nModules * sizeof(MODULEENTRY32W);
  }

  /*
   * Determine the required size for the processes and threads snapshot
   */
  if(dwFlags & (TH32CS_SNAPPROCESS | TH32CS_SNAPTHREAD))
  {
    ULONG ProcOffset = 0;

    ProcessInfo = (PSYSTEM_PROCESS_INFORMATION)ProcThrdInfo;
    do
    {
      ProcessInfo = (PSYSTEM_PROCESS_INFORMATION)((ULONG_PTR)ProcessInfo + ProcOffset);
      nProcesses++;
      nThreads += ProcessInfo->NumberOfThreads;
      ProcOffset = ProcessInfo->NextEntryOffset;
    } while(ProcOffset != 0);

    if(dwFlags & TH32CS_SNAPPROCESS)
    {
      RequiredSnapshotSize += nProcesses * sizeof(PROCESSENTRY32W);
    }
    if(dwFlags & TH32CS_SNAPTHREAD)
    {
      RequiredSnapshotSize += nThreads * sizeof(THREADENTRY32);
    }
  }

  /*
   * Create and map the section
   */

  SSize.QuadPart = RequiredSnapshotSize;

  InitializeObjectAttributes(&ObjectAttributes,
                             NULL,
                             ((dwFlags & TH32CS_INHERIT) ? OBJ_INHERIT : 0),
			     NULL,
			     NULL);

  Status = NtCreateSection(&hSection,
                           SECTION_ALL_ACCESS,
                           &ObjectAttributes,
                           &SSize,
                           PAGE_READWRITE,
                           SEC_COMMIT,
                           NULL);
  if(!NT_SUCCESS(Status))
  {
    return Status;
  }

  SOffset.QuadPart = 0;
  ViewSize = 0;
  Snapshot = NULL;

  Status = NtMapViewOfSection(hSection,
                              NtCurrentProcess(),
                              (PVOID*)&Snapshot,
                              0,
                              0,
                              &SOffset,
                              &ViewSize,
                              ViewShare,
                              0,
                              PAGE_READWRITE);
  if(!NT_SUCCESS(Status))
  {
    NtClose(hSection);
    return Status;
  }

  RtlZeroMemory(Snapshot, sizeof(TH32SNAPSHOT));
  DataOffset = 0;

  /*
   * Initialize the section data and fill it with all the data we collected
   */

  /* initialize the heap list */
  if(dwFlags & TH32CS_SNAPHEAPLIST)
  {
    Snapshot->HeapListCount = nHeaps;
    Snapshot->HeapListOffset = DataOffset;
    HeapListEntry = (LPHEAPLIST32)OffsetToPtr(Snapshot, DataOffset);
    for(i = 0; i < nHeaps; i++)
    {
      HeapListEntry->dwSize = sizeof(HEAPLIST32);
      HeapListEntry->th32ProcessID = th32ProcessID;
      HeapListEntry->th32HeapID = (ULONG_PTR)hi->Heaps[i].BaseAddress;
      HeapListEntry->dwFlags = hi->Heaps[i].Flags;

      HeapListEntry++;
    }

    DataOffset += hi->NumberOfHeaps * sizeof(HEAPLIST32);
  }

  /* initialize the module list */
  if(dwFlags & TH32CS_SNAPMODULE)
  {
    Snapshot->ModuleListCount = nModules;
    Snapshot->ModuleListOffset = DataOffset;
    ModuleListEntry = (LPMODULEENTRY32W)OffsetToPtr(Snapshot, DataOffset);
    for(i = 0; i < nModules; i++)
    {
      ModuleListEntry->dwSize = sizeof(MODULEENTRY32W);
      ModuleListEntry->th32ModuleID = 1; /* no longer used, always set to one! */
      ModuleListEntry->th32ProcessID = th32ProcessID;
      ModuleListEntry->GlblcntUsage = mi->Modules[i].LoadCount;
      ModuleListEntry->ProccntUsage = mi->Modules[i].LoadCount;
      ModuleListEntry->modBaseAddr = (BYTE*)mi->Modules[i].ImageBase;
      ModuleListEntry->modBaseSize = mi->Modules[i].ImageSize;
      ModuleListEntry->hModule = (HMODULE)mi->Modules[i].ImageBase;

      MultiByteToWideChar(CP_ACP,
                          0,
                          &mi->Modules[i].FullPathName[mi->Modules[i].OffsetToFileName],
                          -1,
                          ModuleListEntry->szModule,
                          sizeof(ModuleListEntry->szModule) / sizeof(ModuleListEntry->szModule[0]));

      MultiByteToWideChar(CP_ACP,
                          0,
                          mi->Modules[i].FullPathName,
                          -1,
                          ModuleListEntry->szExePath,
                          sizeof(ModuleListEntry->szExePath) / sizeof(ModuleListEntry->szExePath[0]));

      ModuleListEntry++;
    }

    DataOffset += mi->NumberOfModules * sizeof(MODULEENTRY32W);
  }

  /* initialize the process list */
  if(dwFlags & TH32CS_SNAPPROCESS)
  {
    ULONG ProcOffset = 0;

    Snapshot->ProcessListCount = nProcesses;
    Snapshot->ProcessListOffset = DataOffset;
    ProcessListEntry = (LPPROCESSENTRY32W)OffsetToPtr(Snapshot, DataOffset);
    ProcessInfo = (PSYSTEM_PROCESS_INFORMATION)ProcThrdInfo;
    do
    {
      ProcessInfo = (PSYSTEM_PROCESS_INFORMATION)((ULONG_PTR)ProcessInfo + ProcOffset);

      ProcessListEntry->dwSize = sizeof(PROCESSENTRY32W);
      ProcessListEntry->cntUsage = 0; /* no longer used */
      ProcessListEntry->th32ProcessID = (ULONG)ProcessInfo->UniqueProcessId;
      ProcessListEntry->th32DefaultHeapID = 0; /* no longer used */
      ProcessListEntry->th32ModuleID = 0; /* no longer used */
      ProcessListEntry->cntThreads = ProcessInfo->NumberOfThreads;
      ProcessListEntry->th32ParentProcessID = (ULONG)ProcessInfo->InheritedFromUniqueProcessId;
      ProcessListEntry->pcPriClassBase = ProcessInfo->BasePriority;
      ProcessListEntry->dwFlags = 0; /* no longer used */
      if(ProcessInfo->ImageName.Buffer != NULL)
      {
        lstrcpynW(ProcessListEntry->szExeFile,
                  ProcessInfo->ImageName.Buffer,
                  min(ProcessInfo->ImageName.Length / sizeof(WCHAR), sizeof(ProcessListEntry->szExeFile) / sizeof(ProcessListEntry->szExeFile[0])));
      }
      else
      {
        lstrcpyW(ProcessListEntry->szExeFile, L"[System Process]");
      }

      ProcessListEntry++;

      ProcOffset = ProcessInfo->NextEntryOffset;
    } while(ProcOffset != 0);

    DataOffset += nProcesses * sizeof(PROCESSENTRY32W);
  }

  /* initialize the thread list */
  if(dwFlags & TH32CS_SNAPTHREAD)
  {
    ULONG ProcOffset = 0;

    Snapshot->ThreadListCount = nThreads;
    Snapshot->ThreadListOffset = DataOffset;
    ThreadListEntry = (LPTHREADENTRY32)OffsetToPtr(Snapshot, DataOffset);
    ProcessInfo = (PSYSTEM_PROCESS_INFORMATION)ProcThrdInfo;