"""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))