"""Does the PSP actually fence the TMR against host access?

The TMR is PSP-trusted DRAM (TA sessions, keys, TA working memory) that lives in
host-allocated VRAM (tinygrad puts it at tmr_paddr, SETUP_TMR hands the PSP that
address). After setup the PSP programs a memory-controller fence (read_reg_pair ->
SVC 0x7B register writes) that is SUPPOSED to block CPU/GPU access to the TMR.

If the fence does not actually block the host BAR path, then:
  * READ  -> host can read PSP-trusted state (keys, session material) = disclosure
  * WRITE -> host can corrupt PSP-trusted state = trust violation

This runs on a NORMAL boot (real firmware, real TMR). It reads the TMR region via
the host VRAM BAR, checks whether it contains real (non-zero, non-FF) PSP data,
then attempts a sentinel write + readback to see if the fence blocks writes.

Non-destructive: writes a small sentinel to a few TMR offsets and restores them.
"""
import ctypes, time
from tinygrad import Device
from tinygrad.runtime.autogen.am import am

d = Device["AMD"]; adev = d.iface.dev_impl; psp = adev.psp

def classify(b):
    if all(x == 0 for x in b): return "ZERO"
    if all(x == 0xFF for x in b): return "FF"
    return "DATA"

tp = psp.tmr_paddr
sz = psp.tmr_size
print("TMR: paddr=%#x size=%#x  mc=%#x" % (tp, sz, adev.paddr2mc(tp)))
print("sOS alive:", adev.reg(f"{psp.reg_pref}_81").read() != 0)

# 1) READ the TMR via the host VRAM BAR at several offsets
print("\n--- READ test: is PSP-trusted TMR content visible to the host? ---")
seen = {}
for off in (0x0, 0x1000, 0x10000, 0x100000, 0x300000, 0x800000, sz - 0x1000):
    try:
        b = bytes(adev.vram.view(tp + off, 64)[:])
        k = classify(b)
        seen[k] = seen.get(k, 0) + 1
        print("  +0x%08x : %-5s  %s" % (off, k, b[:32].hex()))
    except Exception as e:
        print("  +0x%08x : read failed %r" % (off, e))
print("  summary:", seen)

# 2) WRITE test: does the fence block host writes into the TMR?
print("\n--- WRITE test: does a host write into the TMR stick? ---")
test_off = 0x2000   # somewhere unlikely to be immediately-critical
orig = bytes(adev.vram.view(tp + test_off, 32)[:])
mark = bytes((0x5A ^ i) & 0xFF for i in range(32))
adev.vram.view(tp + test_off, 32)[:] = mark
adev.gmc.flush_hdp()
time.sleep(0.05)
back = bytes(adev.vram.view(tp + test_off, 32)[:])
print("  wrote :", mark[:16].hex())
print("  read  :", back[:16].hex())
print("  write STUCK (fence does NOT block host):", back == mark)
print("  write BLOCKED (fence works):", back != mark and back == orig)
# restore
adev.vram.view(tp + test_off, 32)[:] = orig
adev.gmc.flush_hdp()
print("  restored:", bytes(adev.vram.view(tp + test_off, 32)[:]) == orig)
print("\nsOS alive at end:", adev.reg(f"{psp.reg_pref}_81").read() != 0)