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

#include <cstdio>
#include <string>

#pragma comment(lib, "ntdll")

extern "C" {

NTSTATUS
NTAPI
NtCreateRegistryTransaction(
  PHANDLE,
  DWORD,
  POBJECT_ATTRIBUTES,
  DWORD
);

NTSTATUS
NTAPI
NtCommitRegistryTransaction(
  HANDLE,
  DWORD
);

}  // extern "C"

int main(int argc, char** argv) {
  LSTATUS st;

  //
  // Recursively delete the test subtree to make sure we're starting from
  // a fresh state.
  //

  st = RegDeleteTreeA(HKEY_CURRENT_USER, "Test");

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

  //
  // Create the test tree structure.
  //

  HKEY hTestRoot;
  st = RegCreateKeyExA(HKEY_CURRENT_USER,
                       "Test",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE,
                       KEY_ALL_ACCESS,
                       NULL,
                       &hTestRoot,
                       NULL);

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

  HKEY hTestKey;
  st = RegCreateKeyExA(hTestRoot,
                       "Key1",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE,
                       KEY_ALL_ACCESS,
                       NULL,
                       &hTestKey,
                       NULL);

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

  for (const std::string SubKeyName : { "SubKey1", "SubKey2" }) {
    HKEY hSubKey;
    st = RegCreateKeyExA(hTestKey,
                         SubKeyName.c_str(),
                         0,
                         NULL,
                         REG_OPTION_VOLATILE,
                         KEY_ALL_ACCESS,
                         NULL,
                         &hSubKey,
                         NULL);

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

    //
    // This is left intentionally commented out so that an active KCB for
    // SubKey1 remains in memory. Thanks to this, the RegDeleteKeyTransactedA
    // call later on succeeds instead of failing with ERROR_FILE_NOT_FOUND.
    //
    // RegCloseKey(hSubKey);
  }

  RegCloseKey(hTestKey);

  //
  // Create a transaction.
  //

  HANDLE hTrans;
  NTSTATUS Status = NtCreateRegistryTransaction(&hTrans,
                                                TRANSACTION_ALL_ACCESS,
                                                NULL,
                                                0);

  if (!NT_SUCCESS(Status)) {
    printf("NtCreateRegistryTransaction failed with error %.8x\n", Status);
    return 1;
  }

  //
  // Open HKCU\Test\Key1 twice, and use the second handle to rename it
  // transactionally.
  //

  HKEY hTestKeyOld;
  st = RegOpenKeyTransactedA(hTestRoot,
                             "Key1",
                             0,
                             KEY_ALL_ACCESS,
                             &hTestKeyOld,
                             hTrans,
                             NULL);

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

  HKEY hTestKeyNew;
  st = RegOpenKeyTransactedA(hTestRoot,
                             "Key1",
                             0,
                             KEY_ALL_ACCESS,
                             &hTestKeyNew,
                             hTrans,
                             NULL);

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

  st = RegRenameKey(hTestKeyNew, NULL, L"Key2");

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

  //
  // Delete each of the two subkeys through a different copy of the parent.
  //

  st = RegDeleteKeyTransactedA(hTestKeyOld, "SubKey1", 0, 0, hTrans, NULL);

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

  st = RegDeleteKeyTransactedA(hTestKeyNew, "SubKey2", 0, 0, hTrans, NULL);

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

  //
  // Commit the transaction.
  //

  Status = NtCommitRegistryTransaction(hTrans, 0);
  if (!NT_SUCCESS(Status)) {
    printf("NtCommitRegistryTransaction failed with error %.8x\n", Status);
    return 1;
  }

  //
  // Try to open the deleted "SubKey1", which should trigger a crash.
  // 

  st = RegOpenKeyExA(hTestRoot, "Key2\\SubKey1", 0, KEY_ALL_ACCESS, &hTestKey);

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

  RegCloseKey(hTestKey);
  RegCloseKey(hTestKeyNew);
  RegCloseKey(hTestKeyOld);
  RegCloseKey(hTestRoot);
  CloseHandle(hTrans);

  return 0;
}