#!/usr/bin/env python3
"""Confirm: a misaligned ring base is accepted at creation, then wedges on first submit.

Ordering matters. The control runs FIRST in the same boot and must succeed, so the wedge that
follows cannot be an inherited dead PSP -- the mistake that produced phantom hits twice in
this audit. Then the misaligned base is created and one command submitted.

Every other ring-create parameter is rejected with a distinct status (size 0/1/max -> 0x4,
size 2 GB -> 0x5, base inside the trusted region -> 0xa), so this asks whether alignment is
checked too. `cmd_buf` misalignment IS rejected (status 0x6) and image misalignment is
rejected (0x34), which makes an unchecked ring base an inconsistency rather than a policy.
"""
import ctypes, glob, struct, sys, time
sys.path.insert(0, '/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
adev = Device["AMD"].iface.dev_impl; psp = adev.psp
RESP = 864; FRAME = ctypes.sizeof(am.struct_psp_gfx_rb_frame); step = FRAME // 4
P = psp.reg_pref
def alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False); return pa, adev.paddr2mc(pa)
MAGIC = 0x71000002
blob = open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0], 'rb').read()
ap_, amc = alloc(0x80000); cbuf, cbmc = alloc(0x8000); scratch, smc = alloc(0x1000)
n = (len(blob) + 3) & ~3
adev.vram.view(ap_, n)[:] = blob + b'\x00' * (n - len(blob)); adev.gmc.flush_hdp()
ring_mc = adev.paddr2mc(psp.ring_paddr)
def create(base, size, tag):
    if adev.reg(f"{P}_71").read() != 0:
        adev.reg(f"{P}_64").write(am.GFX_CTRL_CMD_ID_DESTROY_RINGS); time.sleep(0.05)
    t0 = time.time()
    while time.time() - t0 < 3:
        if adev.reg(f"{P}_64").read() & 0x80000000: break
    adev.wreg_pair(P, "_69", "_70", base)
    adev.reg(f"{P}_71").write(size)
    adev.reg(f"{P}_64").write(am.PSP_RING_TYPE__KM << 16)
    time.sleep(0.1)
    v = adev.reg(f"{P}_64").read()
    acc = (v & 0x8000FFFF) == 0x80000000
    print('  %-26s base=%#-13x create _64=%#-12x status=%#-6x %s'
          % (tag, base, v, v & 0xFFFF, 'ACCEPTED' if acc else 'rejected'), flush=True)
    return acc
def submit(tag):
    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 = cbmc & 0xFFFFFFFF, cbmc >> 32
    t.cmd_buf_len = 0x4000
    adev.vram.view(psp.cmd_paddr, ctypes.sizeof(c))[:] = memoryview(c).cast('B')
    adev.vram.view(psp.cmd_paddr + RESP, 4, 'I')[0] = 0xA5A5A5A5
    adev.vram.view(scratch, 4, 'I')[0] = 0
    f = am.struct_psp_gfx_rb_frame(fence_value=MAGIC,
        cmd_buf_addr_lo=adev.paddr2mc(psp.cmd_paddr) & 0xFFFFFFFF,
        cmd_buf_addr_hi=adev.paddr2mc(psp.cmd_paddr) >> 32,
        fence_addr_lo=smc & 0xFFFFFFFF, fence_addr_hi=smc >> 32)
    adev.vram.view(psp.ring_paddr, FRAME)[:] = bytes(memoryview(f).cast('B'))
    adev.gmc.flush_hdp()
    h0 = adev.reg(f"{P}_81").read()
    adev.reg(f"{P}_67").write(step)
    t0 = time.time(); rs = None
    while time.time() - t0 < 2.0:
        adev.gmc.flush_hdp()
        v = struct.unpack_from('<I', bytes(adev.vram.view(psp.cmd_paddr + RESP, 4)[:]), 0)[0]
        if v != 0xA5A5A5A5: rs = v; break
    adev.gmc.flush_hdp()
    alive = adev.reg(f"{P}_81").read() != h0
    print('    submit %-18s resp=%-9s hb_advancing=%s'
          % (tag, hex(rs) if rs is not None else 'TIMEOUT', alive), flush=True)
    return rs is not None, alive
print('STEP 1 - control first, in this same boot', flush=True)
if not create(ring_mc, psp.ring_size, 'aligned base (control)'): sys.exit(2)
ok, alive = submit('control')
if not ok:
    print('CONTROL FAILED -> result void'); sys.exit(2)
print('STEP 2 - misaligned base, PSP proven healthy a moment ago', flush=True)
for delta in (1, 2, 4, 0x10, 0x100):
    if not create(ring_mc + delta, psp.ring_size, 'base+%#x' % delta):
        print('    rejected at create -> alignment IS checked for +%#x' % delta, flush=True)
        continue
    ok2, alive2 = submit('base+%#x' % delta)
    print('  VERDICT +%#x: %s' % (delta,
        'ACCEPTED then WEDGED -- missing alignment check' if (not ok2 and not alive2)
        else ('accepted and works' if ok2 else 'accepted, no response but PSP alive')), flush=True)
    if not ok2: break