#include "ReflectiveFreeAndExitThread.h"

typedef NTSTATUS
(*NtQueueApcThread)(
    HANDLE    ThreadHandle,
    PVOID     ApcRoutine,
    ULONG_PTR SystemArgument1,
    ULONG_PTR SystemArgument2,
    ULONG_PTR SystemArgument3
);

VOID ReflectiveFreeAndExitThread(HINSTANCE hAppInstance, DWORD dwExitCode) {
    NtQueueApcThread pNtQueueApcThread = (NtQueueApcThread)GetProcAddress(GetModuleHandle(TEXT("ntdll")), "NtQueueApcThread");
    HANDLE hThread = NULL;
    HANDLE hThisThread = NULL;

    do {
        if (!pNtQueueApcThread)
            break;

        // create a suspended thread that will just exit once the APCs have executed
        hThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)ExitThread, 0, CREATE_SUSPENDED, NULL);
        if (!hThread)
            break;

        // open a real handle to this thread to pass in the APC so it operates on this thread and not itself
        hThisThread = OpenThread(THREAD_QUERY_INFORMATION | SYNCHRONIZE, FALSE, GetCurrentThreadId());
        if (!hThisThread)
            break;

        // tell that thread to wait on this thread, ensures VirtualFree isn't called until this thread has exited
        pNtQueueApcThread(hThread, WaitForSingleObjectEx, (ULONG_PTR)hThisThread, INFINITE, FALSE);

        // then close the handle so it's not leaked
        QueueUserAPC((PAPCFUNC)CloseHandle, hThread, (ULONG_PTR)hThisThread);

        // then free the memory
        pNtQueueApcThread(hThread, VirtualFree, (ULONG_PTR)hAppInstance, 0, MEM_RELEASE);
        
        ResumeThread(hThread);
    } while (FALSE);

    if (hThread)
        CloseHandle(hThread);

    ExitThread(dwExitCode);
    return;
}