#!/usr/bin/env python3
"""Does GBR_IH_SET apply the trusted-region check that ring creation does?

Ring creation rejects a base inside the TMR with a distinct status (`0xa`). GBR_IH_SET
(`0x80000`) is the sibling path that points the PSP at a host-supplied interrupt ring; in the
control-space sweep it was inert because no parameters were supplied. It takes its address
from the same `C2PMSG_69/70` pair (and size in `_71`), so the check can be compared directly.

An IH ring is a structure hardware WRITES into, so a missing range check here would be
trusted firmware programming a writer at a host-chosen address -- the same shape as the
unvalidated `fence_addr`, but with the check demonstrably present on the neighbouring path.

Cases, each with the PSP proven healthy first in the same boot:
  A control  normal address                -> baseline status
  B          address inside the TMR         -> rejected like ring-create (0xa) or accepted?
  C          address 0                      -> rejected?
  D          misaligned address             -> rejected, or accepted like the ring base?
"""
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 = 0x18000001
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)
ih, ihmc = alloc(0x10000)
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); tmr_mc = adev.paddr2mc(psp.tmr_paddr)
def create():
    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", ring_mc)
    adev.reg(f"{P}_71").write(psp.ring_size)
    adev.reg(f"{P}_64").write(am.PSP_RING_TYPE__KM << 16)
    time.sleep(0.1)
    return (adev.reg(f"{P}_64").read() & 0x8000FFFF) == 0x80000000
def submit():
    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
    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()
    adev.reg(f"{P}_67").write(step)
    t0 = time.time()
    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: return v
    return None
def ih_set(addr, tag):
    h0 = adev.reg(f"{P}_81").read()
    adev.wreg_pair(P, "_69", "_70", addr)
    adev.reg(f"{P}_71").write(0x10000)
    adev.reg(f"{P}_64").write(am.GFX_CTRL_CMD_ID_GBR_IH_SET)
    time.sleep(0.2)
    v = adev.reg(f"{P}_64").read()
    alive = adev.reg(f"{P}_81").read() != h0
    ok = create() and submit() is not None
    print('  %-26s addr=%#-13x _64=%#-12x status=%#-6x hb=%-5s psp_usable_after=%s'
          % (tag, addr, v, v & 0xFFFF, alive, ok), flush=True)
    return v & 0xFFFF, ok
if not create() or submit() is None:
    print('CONTROL FAILED -> void'); sys.exit(2)
print('control passed first in this boot: ring creates, command completes', flush=True)
case = sys.argv[1]
CASES = {'normal': ihmc, 'tmr': tmr_mc + 0x4000, 'zero': 0, 'unalign': ihmc + 1}
ih_set(CASES[case], case)