#!/usr/bin/env python3
"""Can an unprivileged SDMA queue read an arbitrary register into memory?

`SDMA_OP_POLL_REGMEM / POLL_REG_WRITE_MEM` (key 0x0810, entry 0x0653) is the read
counterpart of SRBM_WRITE. The call-clobber-aware taint pass
(`sdma_audit.py taint`) flags it as the *only* place in TH0 where a host dword reaches
a load address:

    0654  mov r5, r1              ; dword 1 = SRC_ADDR (ABI: addr_31_2)
    0655  mov r8, r1              ; dword 2 = DST_ADDR_LO
    0656  mov r9, r1              ; dword 3 = DST_ADDR_HI
    0657  ldw r10, unk[r5, #0x0]  ; read the register the packet names
    0663  stw r10, mem[r0, #0x48] ; value ...
    0666  stw #0x4, [r0, #0x81]   ; ... DMA'd out to DST_ADDR

SRBM_WRITE is gated on `RB_CNTL.RB_PRIV` ([HW], sdma_srbm_probe.py). This handler shows
no such gate, so the question is whether the *read* side is ungated -- which would let
an unprivileged queue pull privileged register state into its own memory.

Address convention is known, not guessed: amdgpu's `sdma_v6_0_ring_emit_reg_wait`
writes `reg << 2` for POLL_REGMEM packets, i.e. the host register index shifted left by
2. So the test can predict the exact expected value: the host writes a magic into
GRBM_SCRATCH_REG0 first, and the DMA'd result must equal it.

Method as in the other probes: bring the queue up privileged, then clear only RB_PRIV
on the live queue (a queue configured unprivileged from scratch does not fetch at all).

usage: PYTHONPATH=/root/psp/tinygrad sdma_regread_probe.py --apply
"""
import argparse
import os
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from sdma_srbm_probe import Probe, SDMA_OP_NOP                    # noqa: E402

SDMA_OP_POLL_REGMEM = 8
SUBOP_POLL_REG_WRITE_MEM = 1
MAGIC = 0xFEEDFACE


def pkt(reg_index, dst_mc):
    hdr = SDMA_OP_POLL_REGMEM | (SUBOP_POLL_REG_WRITE_MEM << 8)
    return [hdr, (reg_index << 2) & 0xffffffff,
            dst_mc & 0xffffffff, (dst_mc >> 32) & 0xffffffff]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--target', default='regGRBM_SCRATCH_REG0')
    args = ap.parse_args()
    if not args.apply:
        print('dry run')
        return 0

    from tinygrad import Device
    adev = Device["AMD"].iface.dev_impl
    treg = adev.reg(args.target)
    tidx = treg.addr[0]
    p = Probe(adev)
    dst = adev.mm.palloc(0x1000, zero=True, boot=False)
    dst_mc = adev.paddr2mc(dst)
    print('read %s (host idx 0x%x -> packet SRC_ADDR 0x%x) into mc 0x%x' % (
        args.target, tidx, tidx << 2, dst_mc))

    results = {}
    for priv in (1, 0):
        label = 'RB_PRIV=%d' % priv
        if priv == 1:
            p.setup(rb_priv=1)
        else:
            p.reg('_RB_CNTL').update(rb_priv=0)
        magic = MAGIC if priv else (MAGIC ^ 0xffff)
        treg.write(magic)
        adev.vram.view(dst, 0x10)[:] = bytes(0x10)
        adev.gmc.flush_hdp()

        ok, _ = p.submit([SDMA_OP_NOP])
        if not ok:
            print('  %-12s queue not running -- inconclusive' % label)
            results[priv] = None
            continue
        ok2, dt = p.submit(pkt(tidx, dst_mc))
        got = struct.unpack('<I', bytes(adev.vram.view(dst, 4)[:]))[0]
        hit = (got == magic)
        print('  %-12s consumed=%-5s (%.1f ms)  host wrote 0x%08x -> memory got '
              '0x%08x  %s' % (label, ok2, 1000 * dt, magic, got,
                              'LEAKED' if hit else ('changed' if got else 'nothing')))
        results[priv] = hit

    print('\n=== result')
    if results.get(1) is None or results.get(0) is None:
        print('  inconclusive: a control failed')
    elif results[1] and results[0]:
        print('  UNGATED REGISTER READ: an unprivileged queue (RB_PRIV=0) DMA\'d a')
        print('  register value into its own memory, on the same queue where')
        print('  SRBM_WRITE is refused. Reads are not gated; writes are.')
    elif results[1] and not results[0]:
        print('  gated: the read works only with RB_PRIV=1, same trust model as')
        print('  SRBM_WRITE. Not a primitive.')
    else:
        print('  the read did not land even with RB_PRIV=1 -- addressing or packet')
        print('  layout is wrong; no conclusion.')
    print('\nqueue alive at end: %s' % p.submit([SDMA_OP_NOP])[0])
    return 0


if __name__ == '__main__':
    raise SystemExit(main())