#!/usr/bin/env python3
"""Crash-oracle fuzz of the TAs the fault-dump oracle has never been pointed at.

`psp_out/FUZZING-RESULTS.md`'s ~2000-case crash-oracle run was entirely HDCP, and
`hw/gcp_faultoracle_fuzz.py` has since covered GCP. Everything else has only ever been fuzzed
with a status/stall/window oracle, and only over command ids 0..7 at offsets 4..0x80
(`ta_fieldfuzz.py`). This covers the rest with the strongest oracle available.

Command lists are taken from the dispatchers in fw/decomp rather than guessed:

  opm 0..11   merit 1..6   wspe 4,5   csu 1..3   oc 1,2   dtm 2..4
  auc 1..6    fp 1..7      rap 1..7

Commands already recorded as denial-of-service triggers in fw/REPORT-PSP-DOS-TRIGGERS.md are
skipped, since re-triggering them only costs an ATX power cycle: AUC 6, FP 7, FP 2.

Oracle, strongest first:
  1. PSP fault dump C2PMSG 82..88 changes   -> TA data abort, i.e. memory corruption
  2. ring stops responding                  -> hang (the known class)
  3. a write lands past the declared window
  4. stall > 5s, or a status differing from that command's own baseline

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad ta_faultoracle_fuzz.py --apply opm,merit
"""
import argparse
import struct
import sys
import time

DECL = 0x4000      # satisfies every module's minimum-size gate (OPM's 0x3D54 is the largest)
TOTAL = 0x8000
PAT = 0xA5A5A5A5

TAS = {
    # OPM cmd 11 hangs at offset 0x8 with 0x1000 (found by this harness); skipped so the
    # sweep can continue without an ATX power cycle per run.
    'opm':   ('ta_25-0-0-81_dc596ba6', list(range(0, 11)), set()),
    'merit': ('ta_24-0-0-70_5cd4afda', [1, 2, 3, 4, 5, 6], set()),
    # WSPE cmd 5 hangs at offset 0x4 with 0xFFFFFFFF (found by this harness).
    'wspe':  ('ta_30-0-4-27_ba19a7f6', [4], set()),
    'csu':   ('ta_42-0-0-6_1c874c3c', [1, 2, 3], set()),
    'oc':    ('ta_50-0-0-1_cbd47c59', [1, 2], set()),
    # DTM cmd 2 hangs at offset 0x4 with 0x1000 (found by this harness).
    'dtm':   ('ta_18-0-4-19_376639d1', [3, 4], set()),
    'auc':   ('ta_16-0-0-9_9e9044cf', [1, 2, 3, 4, 5], {6}),
    'fp':    ('ta_19-0-0-8_7e8be282', [1, 3, 4, 5, 6], {2, 7}),
    # RAP cmd 3 hangs at offset 0x4 with 0xFFFFFFFF (found by this harness).
    'rap':   ('ta_7-0-2-19_f368676e', [1, 2, 4, 5, 6, 7], set()),
}

OFFS = [0x04, 0x08, 0x0C, 0x10, 0x14, 0x18, 0x1C, 0x20, 0x24, 0x28, 0x2C, 0x30,
        0x40, 0x80, 0x100, 0x200, 0x400, 0x414,
        0x818, 0x81C, 0x820, 0x824, 0x828, 0x82C, 0x830,
        0x1000, 0x2000, 0x3C00, 0x3D00, 0x3D40]
VALUES = [0xFFFFFFFF, 0x80000000, 0x1000]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('tas')
    args = ap.parse_args()
    if not args.apply:
        print('dry run')
        return 0

    import glob
    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:
            return None

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

    def load(blob):
        ap_, amc = alloc(0x80000)
        cp, cmc = alloc(TOTAL)
        n = (len(blob) + 3) & ~3
        adev.vram.view(ap_, 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 = amc & 0xFFFFFFFF, amc >> 32
        t.app_len = len(blob)
        t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = cmc & 0xFFFFFFFF, cmc >> 32
        t.cmd_buf_len = DECL
        try:
            r = psp._ring_submit(c)
            return r.resp.status & 0xFFFFFFFF, r.resp.session_id, cp
        except Exception as e:
            try:
                return int(str(e).strip().split()[-1]) & 0xFFFFFFFF, None, cp
            except Exception:
                return -1, None, cp

    def inv(sid, cp, 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(cp, 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(cp, TOTAL)[:])
        bey = sum(1 for o in range(DECL, TOTAL, 4)
                  if struct.unpack_from('<I', out, o)[0] != PAT)
        return st, dt, bey

    f0 = faults()
    print('baseline fault regs=%s sos_alive=%s' % ([hex(x) for x in f0] if f0 else None, alive()),
          flush=True)
    grand = 0
    for name in args.tas.split(','):
        stem, cmds, skip = TAS[name]
        cands = glob.glob('/root/psp/contall/' + stem + '*')
        if not cands:
            print('%-6s binary not found (%s)' % (name, stem), flush=True)
            continue
        st, sid, cp = load(open(cands[0], 'rb').read())
        if st != 0 or sid is None:
            print('%-6s LOAD %#x' % (name, st), flush=True)
            continue
        found = []
        n = 0
        for cmd in cmds:
            if cmd in skip:
                continue
            base, _, _ = inv(sid, cp, cmd)
            stop = False
            for off in OFFS:
                for v in VALUES:
                    s2, dt, bey = inv(sid, cp, cmd, off, v)
                    n += 1
                    f = faults()
                    if f is not None and f0 is not None and f != f0:
                        found.append(('FAULT-DUMP', cmd, hex(off), hex(v),
                                      [hex(x) for x in f])); stop = True
                    elif s2 == -1:
                        found.append(('RING-DEAD', cmd, hex(off), hex(v))); stop = True
                    elif bey:
                        found.append(('OUT-OF-WINDOW', cmd, hex(off), hex(v), bey)); stop = True
                    elif dt > 5:
                        found.append(('STALL', cmd, hex(off), hex(v), round(dt, 1)))
                    elif s2 != base:
                        found.append(('status', cmd, hex(off), hex(v), hex(base), hex(s2)))
                    if stop:
                        break
                if stop:
                    break
            if stop:
                break
        grand += n
        hard = [x for x in found if x[0] in ('FAULT-DUMP', 'RING-DEAD', 'OUT-OF-WINDOW')]
        print('%-6s %4d cases  %s' % (name, n, hard[:2] if hard else
              ('%d soft diffs' % len(found) if found else 'clean')), flush=True)
        if hard and hard[0][0] in ('FAULT-DUMP', 'RING-DEAD'):
            print('  stopping: %s on %s' % (hard[0][0], name), flush=True)
            break
    print('\ntotal %d cases; sos_alive=%s fault regs=%s'
          % (grand, alive(), [hex(x) for x in (faults() or [])]), flush=True)
    return 0


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