#!/usr/bin/env python3
"""Will the PSP write into a mapped region the caller did not designate as its buffer?

Pass 105: cmd_buf = MC 0 hangs the PSP with a flood of VMC faults, i.e. the address is
dereferenced without aperture validation. The follow-up that matters: give a VALID mapped
address that is not the buffer the host "meant", and see whether TA output lands there.

Allocate two separate regions A and B. Declare A as cmd_buf at LOAD_TA, then for the invoke
point the descriptor at B. If output appears in B, the PSP writes wherever it is told.
"""
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)
PATA=0xAAAAAAAA; PATB=0xBBBBBBBB
base=open(glob.glob('/root/psp/contall/ta_16-0-0-9_9e9044cf*')[0],'rb').read()  # AUC: writes output
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
ap_,amc=alloc(0x80000); A,Amc=alloc(0x8000); B,Bmc=alloc(0x8000)
n=(len(base)+3)&~3
adev.vram.view(ap_,n)[:]=base+b'\x00'*(n-len(base))
adev.vram.view(A,0x8000)[:]=bytes(struct.pack('<I',PATA)*(0x8000//4))
adev.vram.view(B,0x8000)[:]=bytes(struct.pack('<I',PATB)*(0x8000//4))
adev.gmc.flush_hdp()
print('A mc=%#x  B mc=%#x'%(Amc,Bmc), flush=True)
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=Amc&0xFFFFFFFF,Amc>>32; t.cmd_buf_len=0x4000
r=submit(c); sid=r.resp.session_id
print('load (cmd_buf=A) %#x sid=%s'%(r.resp.status&0xFFFFFFFF,sid), flush=True)
buf=bytearray(0x4000); struct.pack_into('<I',buf,0,5)
adev.vram.view(A,0x4000)[:]=bytes(buf)   # only A gets the request; B stays pristine
adev.gmc.flush_hdp()
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,5
t0=time.time()
try: s=submit(c).resp.status&0xFFFFFFFF
except Exception as e:
    try: s=int(str(e).strip().split()[-1])&0xFFFFFFFF
    except Exception: s=-1
adev.gmc.flush_hdp()
oa=bytes(adev.vram.view(A,0x8000)[:]); ob=bytes(adev.vram.view(B,0x8000)[:])
ca=sum(1 for o in range(0x10,0x8000,4) if struct.unpack_from('<I',oa,o)[0]!=PATA)
cb=sum(1 for o in range(0,0x8000,4) if struct.unpack_from('<I',ob,o)[0]!=PATB)
print('invoke -> %#x %.1fs ; dwords changed: A=%d (expected) B=%d (must be 0)'%(s,time.time()-t0,ca,cb), flush=True)