#!/usr/bin/env python3
"""Does the PSP validate that cmd_buf lies in an allowed aperture?

Every buffer so far was a legitimate VRAM allocation. If an arbitrary MC address is accepted,
the question is whether TA output can be steered somewhere it should not go. The TMR is
already known to be rejected (pass 87); this asks about everything else.

Cases (control first). Only statuses/timing are interpreted; nothing is read back from
addresses we do not own.
  null (MC 0)
  just past our allocation
  far beyond the VRAM size
  the top of the 48-bit MC space
"""
import glob, struct, sys, time
sys.path.insert(0,'/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
from ring_wrap import install
adev=Device["AMD"].iface.dev_impl; psp=adev.psp
submit=install(psp,adev)
base=open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0],'rb').read()
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
ap_,amc=alloc(0x80000); cp,cmc=alloc(0x8000)
n=(len(base)+3)&~3
adev.vram.view(ap_,n)[:]=base+b'\x00'*(n-len(base)); adev.gmc.flush_hdp()
try: vsz=adev.vram_size
except Exception: vsz=None
print('our cmd_buf mc=%#x ; vram_size=%s'%(cmc, hex(vsz) if vsz else 'unknown'), flush=True)
def run(tag, mc):
    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=mc&0xFFFFFFFF,mc>>32; t.cmd_buf_len=0x4000
    t0=time.time()
    try: r=submit(c); st=r.resp.status&0xFFFFFFFF
    except Exception as e:
        try: st=int(str(e).strip().split()[-1])&0xFFFFFFFF
        except Exception: st=-1
    dt=time.time()-t0
    print('  %-36s mc=%#-14x -> %#-12x %.1fs%s'%(tag,mc,st,dt,'  <<< STALL' if dt>5 else ''), flush=True)
    return dt>5
if run('CONTROL our allocation', cmc): sys.exit(3)
for tag,mc in (('null', 0),
               ('just past our allocation', cmc+0x8000),
               ('far beyond VRAM', (vsz or (1<<34)) + 0x10000000),
               ('top of 48-bit MC space', (1<<48)-0x10000)):
    if run(tag, mc): break