#include <windows.h>
#include <stdio.h>
#include <conio.h>
#include <vector>

CRITICAL_SECTION cs_DB1;
CRITICAL_SECTION cs_DB2;

static DWORD WINAPI ThreadProc1(LPVOID lpParam)
{
    EnterCriticalSection(&cs_DB1);

    //
    // Do work on stack
    //
    wprintf(L"Updating database 1\n");
    Sleep(3000);
    EnterCriticalSection(&cs_DB2);

    wprintf(L"Updating database 2\n");

    LeaveCriticalSection(&cs_DB2);
    LeaveCriticalSection(&cs_DB1);

    return 1;
}

static DWORD WINAPI ThreadProc2(LPVOID lpParam)
{
    EnterCriticalSection(&cs_DB2);

    //
    // Do work on stack
    //
    wprintf(L"Updating database 2\n");
    Sleep(3000);
    EnterCriticalSection(&cs_DB1);

    wprintf(L"Updating database 1\n");

    LeaveCriticalSection(&cs_DB1);
    LeaveCriticalSection(&cs_DB2);

    return 1;
}





typedef struct _RTL_PROCESS_MODULE_INFORMATION
{
    ULONG Section;
    PVOID MappedBase;
    PVOID ImageBase;
    ULONG ImageSize;
    ULONG Flags;
    USHORT LoadOrderIndex;
    USHORT InitOrderIndex;
    USHORT LoadCount;
    USHORT OffsetToFileName;
    CHAR FullPathName[256];
} RTL_PROCESS_MODULE_INFORMATION, *PRTL_PROCESS_MODULE_INFORMATION;

typedef struct _RTL_PROCESS_MODULES
{
    ULONG NumberOfModules;
    RTL_PROCESS_MODULE_INFORMATION Modules[1];
} RTL_PROCESS_MODULES, *PRTL_PROCESS_MODULES;

typedef struct _RTL_PROCESS_MODULE_INFORMATION_EX
{
    ULONG NextOffset;
    RTL_PROCESS_MODULE_INFORMATION BaseInfo;
    ULONG ImageCheckSum;
    ULONG TimeDateStamp;
    PVOID DefaultBase;
} RTL_PROCESS_MODULE_INFORMATION_EX, *PRTL_PROCESS_MODULE_INFORMATION_EX;

typedef struct _RTL_HEAP_TAG_INFO
{
    ULONG NumberOfAllocations;
    ULONG NumberOfFrees;
    SIZE_T BytesAllocated;
} RTL_HEAP_TAG_INFO, *PRTL_HEAP_TAG_INFO;

typedef struct _RTL_HEAP_USAGE_ENTRY
{
    struct _RTL_HEAP_USAGE_ENTRY *Next;
    PVOID Address;
    SIZE_T Size;
    USHORT AllocatorBackTraceIndex;
    USHORT TagIndex;
} RTL_HEAP_USAGE_ENTRY, *PRTL_HEAP_USAGE_ENTRY;

typedef struct _RTL_HEAP_USAGE
{
    ULONG Length;
    SIZE_T BytesAllocated;
    SIZE_T BytesCommitted;
    SIZE_T BytesReserved;
    SIZE_T BytesReservedMaximum;
    PRTL_HEAP_USAGE_ENTRY Entries;
    PRTL_HEAP_USAGE_ENTRY AddedEntries;
    PRTL_HEAP_USAGE_ENTRY RemovedEntries;
    ULONG_PTR Reserved[8];
} RTL_HEAP_USAGE, *PRTL_HEAP_USAGE;

typedef struct _RTL_HEAP_WALK_ENTRY
{
    PVOID DataAddress;
    SIZE_T DataSize;
    UCHAR OverheadBytes;
    UCHAR SegmentIndex;
    USHORT Flags;
    union
    {
        struct
        {
            SIZE_T Settable;
            USHORT TagIndex;
            USHORT AllocatorBackTraceIndex;
            ULONG Reserved[2];
        } Block;
        struct
        {
            ULONG_PTR CommittedSize;
            ULONG_PTR UnCommittedSize;
            PVOID FirstEntry;
            PVOID LastEntry;
        } Segment;
    };
} RTL_HEAP_WALK_ENTRY, *PRTL_HEAP_WALK_ENTRY;

typedef struct _RTL_HEAP_ENTRY
{
    SIZE_T Size;
    USHORT Flags;
    USHORT AllocatorBackTraceIndex;
    union
    {
        struct
        {
            SIZE_T Settable;
            ULONG Tag;
        } s1;
        struct
        {
            SIZE_T CommittedSize;
            PVOID FirstBlock;
        } s2;
    } u;
} RTL_HEAP_ENTRY, *PRTL_HEAP_ENTRY;

typedef struct _RTL_HEAP_TAG
{
    ULONG NumberOfAllocations;
    ULONG NumberOfFrees;
    SIZE_T BytesAllocated;
    USHORT TagIndex;
    USHORT CreatorBackTraceIndex;
    WCHAR TagName[24];
} RTL_HEAP_TAG, *PRTL_HEAP_TAG;

typedef struct _RTL_HEAP_INFORMATION
{
    PVOID BaseAddress;
    ULONG Flags;
    USHORT EntryOverhead;
    USHORT CreatorBackTraceIndex;
    SIZE_T BytesAllocated;
    SIZE_T BytesCommitted;
    ULONG NumberOfTags;
    ULONG NumberOfEntries;
    ULONG NumberOfPseudoTags;
    ULONG PseudoTagGranularity;
    ULONG Reserved[4];
    PRTL_HEAP_TAG Tags;
    PRTL_HEAP_ENTRY Entries;
} RTL_HEAP_INFORMATION, *PRTL_HEAP_INFORMATION;

typedef struct _RTL_PROCESS_HEAPS
{
    ULONG NumberOfHeaps;
    RTL_HEAP_INFORMATION Heaps[1];
} RTL_PROCESS_HEAPS, *PRTL_PROCESS_HEAPS;

typedef struct _RTL_PROCESS_LOCK_INFORMATION
{
    PVOID Address;
    USHORT Type;
    USHORT CreatorBackTraceIndex;
    ULONG OwnerThreadId;
    ULONG ActiveCount;
    ULONG ContentionCount;
    ULONG EntryCount;
    ULONG RecursionCount;
    ULONG NumberOfSharedWaiters;
    ULONG NumberOfExclusiveWaiters;
} RTL_PROCESS_LOCK_INFORMATION, *PRTL_PROCESS_LOCK_INFORMATION;

typedef struct _RTL_PROCESS_LOCKS
{
    ULONG NumberOfLocks;
    RTL_PROCESS_LOCK_INFORMATION Locks[1];
} RTL_PROCESS_LOCKS, *PRTL_PROCESS_LOCKS;

typedef struct _RTL_PROCESS_BACKTRACE_INFORMATION
{
    PVOID SymbolicBackTrace;
    ULONG TraceCount;
    USHORT Index;
    USHORT Depth;
    PVOID BackTrace[16];
} RTL_PROCESS_BACKTRACE_INFORMATION, *PRTL_PROCESS_BACKTRACE_INFORMATION;

typedef struct _RTL_PROCESS_BACKTRACES
{
    ULONG CommittedMemory;
    ULONG ReservedMemory;
    ULONG NumberOfBackTraceLookups;
    ULONG NumberOfBackTraces;
    RTL_PROCESS_BACKTRACE_INFORMATION BackTraces[1];
} RTL_PROCESS_BACKTRACES, *PRTL_PROCESS_BACKTRACES;

typedef struct _RTL_PROCESS_VERIFIER_OPTIONS
{
    ULONG SizeStruct;
    ULONG Option;
    UCHAR OptionData[1];
    //
    // Option array continues below
    //
} RTL_PROCESS_VERIFIER_OPTIONS, *PRTL_PROCESS_VERIFIER_OPTIONS;

typedef struct _RTL_DEBUG_INFORMATION
{
    HANDLE SectionHandleClient;
    PVOID ViewBaseClient;
    PVOID ViewBaseTarget;
    ULONG ViewBaseDelta;
    HANDLE EventPairClient;
    PVOID EventPairTarget;
    HANDLE TargetProcessId;
    HANDLE TargetThreadHandle;
    ULONG Flags;
    ULONG OffsetFree;
    ULONG CommitSize;
    ULONG ViewSize;
    union
    {
        PRTL_PROCESS_MODULES Modules;
        PRTL_PROCESS_MODULE_INFORMATION_EX ModulesEx;
    };
    PRTL_PROCESS_BACKTRACES BackTraces;
    PRTL_PROCESS_HEAPS Heaps;
    PRTL_PROCESS_LOCKS Locks;
    HANDLE SpecificHeap;
    HANDLE TargetProcessHandle;
    RTL_PROCESS_VERIFIER_OPTIONS VerifierOptions;
    HANDLE ProcessHeap;
    HANDLE CriticalSectionHandle;
    HANDLE CriticalSectionOwnerThread;
    PVOID Reserved[4];
} RTL_DEBUG_INFORMATION, *PRTL_DEBUG_INFORMATION;





typedef struct _DEBUGMODULEINFO {
    DWORD	ImageBase;
    DWORD	ImageSize;
    DWORD	Unknown01;			// possibly some kind of version info. nthandle doesnt use it
    USHORT	DllSequenceNum;  // if 0 then EXE
    USHORT	NumDlls;			// only know if seqnum is 0
    DWORD	GrantedAccess;
    CHAR	Name[MAX_PATH];
    DWORD	Unknown02;
}DEBUGMODULEINFO, *PDEBUGMODULEINFO;

typedef struct _QUERYDEBUGBUFFER_HEADER {
    DWORD	Unkown12[12];
    DWORD*	ModArrayHeader;
    DWORD	Unkown11[11];
    DWORD	NumNames;         // total entries including the EXE
    DWORD	Reserved[2];
    DEBUGMODULEINFO ModInfo[1];
}QUERYDEBUGBUFFER, *PQUERYDEBUGBUFFER;

#define RTL_DEBUG_QUERY_MODULES                             0x01
#define RTL_DEBUG_QUERY_BACKTRACES                          0x02
#define RTL_DEBUG_QUERY_HEAPS                               0x04
#define RTL_DEBUG_QUERY_HEAP_TAGS                           0x08
#define RTL_DEBUG_QUERY_HEAP_BLOCKS                         0x10
#define RTL_DEBUG_QUERY_LOCKS                               0x20

typedef DWORD(WINAPI *PRtlQueryProcessDebugInformation)(HANDLE, DWORD, PVOID);
typedef DWORD* (WINAPI *PRtlCreateQueryDebugBuffer)(DWORD, DWORD);
typedef DWORD(WINAPI *PRtlDestroyQueryDebugBuffer)(PDWORD);

static bool EnumerateHeaps()
{
    // Get addresses of NTDLL functions necessary for the heap enumeration.
    HMODULE hNtDll = GetModuleHandle(L"ntdll.dll");
    PRtlCreateQueryDebugBuffer pfnRtlCreateQueryDebugBuffer = (PRtlCreateQueryDebugBuffer)GetProcAddress(hNtDll, "RtlCreateQueryDebugBuffer");
    PRtlDestroyQueryDebugBuffer pfnRtlDestroyQueryDebugBuffer = (PRtlDestroyQueryDebugBuffer)GetProcAddress(hNtDll, "RtlDestroyQueryDebugBuffer");
    PRtlQueryProcessDebugInformation pfnRtlQueryProcessDebugInformation = (PRtlQueryProcessDebugInformation)GetProcAddress(hNtDll, "RtlQueryProcessDebugInformation");

    if (!pfnRtlCreateQueryDebugBuffer || !pfnRtlDestroyQueryDebugBuffer || !pfnRtlQueryProcessDebugInformation)
    {
        return false;
    }

    // Create debug buffer to hold heap information.
    PRTL_DEBUG_INFORMATION db = (PRTL_DEBUG_INFORMATION)pfnRtlCreateQueryDebugBuffer(0, FALSE);
    if (!db)
    {
        return false;
    }

    HANDLE procHdl = OpenProcess(PROCESS_ALL_ACCESS, FALSE, 2580);

    // Get heap information and put it inside the debug buffer.
    NTSTATUS result = pfnRtlQueryProcessDebugInformation(procHdl, RTL_DEBUG_QUERY_LOCKS, db);
    if (result != 0)
    {
        pfnRtlDestroyQueryDebugBuffer((PDWORD)db);
        return false;
    }

    // Walk the heaps and add them to the input vector.
    for (unsigned int i = 0; i < db->Heaps->NumberOfHeaps; ++i)
    {
        PRTL_HEAP_INFORMATION curHeap = &db->Heaps->Heaps[i];
        //Win32HeapInformation& heap = heapInfoList.Add();
        //heap.VirtualAddress = (LONG_PTR)curHeap->BaseAddress;
        //heap.BlockCount = curHeap->NumberOfEntries;
        //heap.CommittedSize = (LONG)curHeap->BytesCommitted;
        //heap.AllocatedSize = (LONG)curHeap->BytesAllocated;
        //heap.Flags = curHeap->Flags;
    }

    // Clean up buffer and return.
    pfnRtlDestroyQueryDebugBuffer((PDWORD)db);
    CloseHandle(procHdl);

    return true;
}

void __cdecl main()
{
    EnumerateHeaps();

    if (!InitializeCriticalSectionAndSpinCount(&cs_DB1, 2000))
    {
        return;
    }
    if (!InitializeCriticalSectionAndSpinCount(&cs_DB2, 2000))
    {
        DeleteCriticalSection(&cs_DB1);
        return;
    }

    size_t size1 = sizeof(RTL_DEBUG_INFORMATION);
    size_t size2 = sizeof(QUERYDEBUGBUFFER);

    static PRtlQueryProcessDebugInformation RtlQueryProcessDebugInformation;
    static PRtlCreateQueryDebugBuffer RtlCreateQueryDebugBuffer;
    static PRtlDestroyQueryDebugBuffer RtlDestroyQueryDebugBuffer;

    HMODULE hNtDll = GetModuleHandle(L"ntdll.dll");
    RtlQueryProcessDebugInformation = (PRtlQueryProcessDebugInformation)
        GetProcAddress(hNtDll,
            "RtlQueryProcessDebugInformation");

    RtlCreateQueryDebugBuffer = (PRtlCreateQueryDebugBuffer)
        GetProcAddress(hNtDll,
            "RtlCreateQueryDebugBuffer");

    RtlDestroyQueryDebugBuffer = (PRtlDestroyQueryDebugBuffer)
        GetProcAddress(hNtDll,
            "RtlDestroyQueryDebugBuffer");

    HANDLE procHdl = OpenProcess(PROCESS_ALL_ACCESS, FALSE, GetCurrentProcessId());
    if (procHdl != NULL && procHdl != INVALID_HANDLE_VALUE)
    {
        DWORD *pRtlBuffer = RtlCreateQueryDebugBuffer(NULL, true);
        DWORD error = RtlQueryProcessDebugInformation(procHdl, RTL_DEBUG_QUERY_LOCKS, pRtlBuffer);
        RTL_DEBUG_INFORMATION *pDebugInfo = (RTL_DEBUG_INFORMATION*)pRtlBuffer;
        RtlDestroyQueryDebugBuffer(pRtlBuffer);
    }
    CloseHandle(procHdl);


    return;

    DWORD dwId = 0;
    HANDLE hThread1 = CreateThread(NULL, 0, ThreadProc1, NULL, 0, &dwId);
    HANDLE hThread2 = CreateThread(NULL, 0, ThreadProc2, NULL, 0, &dwId);

    HANDLE threads[] = { hThread1,hThread2 };

    Sleep(10000);


    wprintf(L"Debugging information (critical sections):\n");
    std::vector<PRTL_CRITICAL_SECTION> criticalSections;
    PRTL_CRITICAL_SECTION_DEBUG dbgInfo = cs_DB1.DebugInfo;
    LIST_ENTRY lockList = dbgInfo->ProcessLocksList;



    DWORD res = WaitForMultipleObjectsEx(2, threads, true, INFINITE, false);


    DeleteCriticalSection(&cs_DB1);
    DeleteCriticalSection(&cs_DB2);
}