#!/usr/bin/env python3
"""Crash-oracle fuzz of GCP's real command space, with the PSP fault dump as the detector.

Two gaps this closes, both established by reading the existing records rather than guessed:

* `psp_out/FUZZING-RESULTS.md`'s ~2000-case crash-oracle run was **all HDCP** (16 top-level
  commands plus the 0x0F subcommand framing). The fault-dump oracle, C2PMSG 82..88 plus an
  sOS-alive check, has never been pointed at any other TA.
* the cross-TA fuzzer that does exist on this box, `ta_fieldfuzz.py`, mutates offsets 4..0x80
  for command ids 0..7 only. GCP's commands are 1..5 and 0xFE02..0xFE0F, and its size/offset
  fields live at 0x414 and 0x818..0x830 and in data20 at +2040..+2100. So GCP's actual
  parameter space has never been fuzzed, with any oracle.

The field list is not arbitrary: those are the words
`fw/FINDING-GCP-FE09-CONFIGINIT-HOST-OFFSET-WRITE.md` showed GCP reads as offsets and limits.

Oracle, in order of strength:
  1. PSP fault dump C2PMSG 82..88 changing  -> a TA data abort, i.e. memory corruption
  2. ring not responding                    -> wedge (already-known class)
  3. a write landing past the declared window
  4. stall > 5s, or a status differing from the command's own baseline

Out-of-window writes are known to wedge the PSP and need an ATX power cycle, so this stops at
the first anomaly per command and aborts the run entirely if the ring dies.

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad gcp_faultoracle_fuzz.py --apply
"""
import argparse
import os
import struct
import subprocess
import time

DECL = 0x4000      # declared cmd_buf_len; GCP needs >= 0x2958
TOTAL = 0x8000     # allocation, so writes past DECL are visible
PAT = 0xA5A5A5A5

CMDS = [1, 2, 3, 4, 5, 0xFE02, 0xFE03, 0xFE07, 0xFE08, 0xFE09,
        0xFE0A, 0xFE0B, 0xFE0C, 0xFE0D, 0xFE0F]

# arg block, the 0x414 limit, the 0x818..0x830 offset/limit cluster, and data20 (+0x20) fields
OFFS = ([0x10, 0x14, 0x18, 0x1C, 0x414]
        + list(range(0x810, 0x844, 4))
        + [0x20 + n for n in (1012, 2040, 2044, 2048, 2052, 2056, 2060, 2064, 2072,
                              2076, 2080, 2088, 2092, 2100)])
VALUES = [0xFFFFFFFF, 0x80000000, 0x1000]

# (command, offset, value) triples already known to hang the ring; skipped so the sweep can
# continue past them. Each costs an ATX power cycle if not skipped.
# FE03 hangs for ANY large value at offset 0x10 (arg10): 0xFFFFFFFF and 0x80000000 both
# confirmed, each costing a power cycle. Skip the whole field.
SKIP_FIELD = {(0xFE03, 0x10)}
SKIP = set()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--from-cmd', type=lambda x: int(x, 0), default=0)
    ap.add_argument('--gcp', default='/root/psp/contall/ta_22-0-0-77_7c426e7f.bin')
    ap.add_argument('--ta-bin', default='/lib/firmware/amdgpu/psp_13_0_0_ta.bin')
    args = ap.parse_args()
    if not args.apply:
        print('dry run')
        return 0

    from tinygrad import Device
    import tinygrad.runtime.autogen.am.am as am

    adev = Device["AMD"].iface.dev_impl
    psp = adev.psp
    FAULT = [getattr(adev, 'mmMP0_SMN_C2PMSG_%d' % n) for n in range(82, 89)]

    def faults():
        try:
            return tuple(r.read() for r in FAULT)
        except Exception:
            return None

    def alive():
        try:
            return bool(psp.is_sos_alive())
        except Exception:
            try:
                return adev.mmMP0_SMN_C2PMSG_58.read() != 0
            except Exception:
                return None

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

    app_pa, app_mc = alloc(0x80000)
    shm_pa, shm_mc = alloc(TOTAL)

    def load(blob):
        n = (len(blob) + 3) & ~3
        adev.vram.view(app_pa, n)[:] = blob + b'\x00' * (n - len(blob))
        adev.gmc.flush_hdp()
        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_lo, t.app_phy_addr_hi = app_mc & 0xFFFFFFFF, app_mc >> 32
        t.app_len = len(blob)
        t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = shm_mc & 0xFFFFFFFF, shm_mc >> 32
        t.cmd_buf_len = DECL
        try:
            r = psp._ring_submit(c)
            return r.resp.status & 0xFFFFFFFF, r.resp.session_id
        except Exception as e:
            try:
                return int(str(e).strip().split()[-1]) & 0xFFFFFFFF, None
            except Exception:
                return -1, None

    def unload(sid):
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
        c.cmd.cmd_unload_ta.session_id = sid
        try:
            return psp._ring_submit(c).resp.status & 0xFFFFFFFF
        except Exception:
            return -1

    def inv(sid, cmd, off=None, val=None):
        buf = bytearray(struct.pack('<I', PAT) * (TOTAL // 4))
        struct.pack_into('<I', buf, 0, cmd)
        for o in OFFS:
            struct.pack_into('<I', buf, o, 0)
        if off is not None:
            struct.pack_into('<I', buf, off, val)
        adev.vram.view(shm_pa, TOTAL)[:] = bytes(buf)
        adev.gmc.flush_hdp()
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
        iv = c.cmd.cmd_invoke_cmd
        iv.session_id, iv.ta_cmd_id = sid, cmd
        t0 = time.time()
        try:
            st = psp._ring_submit(c).resp.status & 0xFFFFFFFF
        except Exception:
            st = -1
        dt = time.time() - t0
        adev.gmc.flush_hdp()
        out = bytes(adev.vram.view(shm_pa, TOTAL)[:])
        bey = sum(1 for o in range(DECL, TOTAL, 4)
                  if struct.unpack_from('<I', out, o)[0] != PAT)
        return st, dt, bey

    tab = args.ta_bin
    if not os.path.exists(tab) and os.path.exists(tab + '.zst'):
        subprocess.run(['zstd', '-d', '-f', '-o', '/tmp/ta.bin', tab + '.zst'],
                       check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        tab = '/tmp/ta.bin'
    d = open(tab, 'rb').read()
    m = d.find(b'$PS1')
    n = struct.unpack_from('<I', d, m - 0x10 + 0x14)[0]
    _, wsid = load(d[m - 0x10:m - 0x10 + 0x200 + n])
    if wsid is not None:
        unload(wsid)

    st, sid = load(open(args.gcp, 'rb').read())
    f0 = faults()
    print('GCP load %#x sid=%s  fault regs=%s  sos_alive=%s'
          % (st, sid, [hex(x) for x in f0] if f0 else None, alive()), flush=True)
    if st != 0:
        return 1

    print('cases: %d commands x %d offsets x %d values = %d\n'
          % (len(CMDS), len(OFFS), len(VALUES), len(CMDS) * len(OFFS) * len(VALUES)), flush=True)

    total = 0
    for cmd in CMDS:
        if cmd < args.from_cmd:
            continue
        base, _, _ = inv(sid, cmd)
        hits = []
        for off in OFFS:
            for v in VALUES:
                if (cmd, off, v) in SKIP or (cmd, off) in SKIP_FIELD:
                    continue
                st2, dt, bey = inv(sid, cmd, off, v)
                total += 1
                f = faults()
                if f is not None and f0 is not None and f != f0:
                    hits.append(('FAULT-DUMP', hex(off), hex(v), [hex(x) for x in f]))
                elif st2 == -1:
                    hits.append(('RING-DEAD', hex(off), hex(v)))
                elif bey:
                    hits.append(('OUT-OF-WINDOW', hex(off), hex(v), bey))
                elif dt > 5:
                    hits.append(('STALL', hex(off), hex(v), round(dt, 1)))
                elif st2 != base:
                    hits.append(('status', hex(off), hex(v), hex(base), hex(st2)))
                if hits and hits[-1][0] in ('FAULT-DUMP', 'RING-DEAD', 'OUT-OF-WINDOW'):
                    break
            if hits and hits[-1][0] in ('FAULT-DUMP', 'RING-DEAD', 'OUT-OF-WINDOW'):
                break
        tag = hits[0][0] if hits else 'clean'
        print('  cmd %#-7x %-14s %s' % (cmd, tag, hits[:4] if hits else ''), flush=True)
        if hits and hits[0][0] in ('FAULT-DUMP', 'RING-DEAD'):
            print('  stopping the run: %s' % hits[0][0], flush=True)
            break
    print('\n%d cases executed; sos_alive=%s fault regs=%s'
          % (total, alive(), [hex(x) for x in (faults() or [])]), flush=True)
    return 0


if __name__ == '__main__':
    raise SystemExit(main())