#!/usr/bin/env python3
"""Compound test: is SDMA's POLL_REGMEM a register read oracle for an *unprivileged* queue?

This comes from putting two findings together.

  * DMCUB FINDING-10 (`../../dmcub/AUDIT-DMCUB-05-REGSEQ.md`): `REG_REG_WAIT` leaks
    memory because it converts a value the host cannot read into something the host
    *can* measure -- the time the command takes to retire. No memory-safety bug is
    involved; the channel is the wait itself.
  * SDMA (`../AUDIT-SDMA-01.md`): register *access* is gated on `RB_CNTL.RB_PRIV`
    [HW] -- `SRBM_WRITE` (write) and `POLL_REG_WRITE_MEM` (read-into-memory) both
    refuse on an unprivileged queue.

`POLL_REGMEM` (op 8, sub_op 0) is the same shape as REG_REG_WAIT: it polls a register
until `(value & mask) FUNC reference`. Critically it **writes nothing**, so the gating
argument that covers the other two ("don't let an unprivileged queue read registers
into its memory") may simply not have been applied to it. If it is ungated, then the
FINDING-10 technique transplants onto an engine whose queues *are* user-submittable,
and an unprivileged process gets an arbitrary register read oracle at one bit per
packet -- which none of the DMCUB findings can offer, because that inbox is kernel-only.

Packet (navi10_sdma_pkt_open.h):
    dw0 header: op=8, sub_op=0, func<<28 (3 = equal), mem_poll<<31 (0 = register)
    dw1 addr_lo  = register index << 2   (amdgpu's emit_reg_wait convention)
    dw2 addr_hi  = 0
    dw3 value    = reference
    dw4 mask
    dw5 interval/retry count

Oracle: set a scratch register to a known value, then submit a matching poll and a
non-matching poll and compare retire times. A non-match must spin (long); a match must
retire immediately. Run with RB_PRIV=1 (control) and RB_PRIV=0 (the question), using the
live-flip method -- a queue configured unprivileged from scratch does not fetch at all.

usage: PYTHONPATH=/root/psp/tinygrad sdma_pollreg_oracle.py --apply
"""
import argparse
import os
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
FUNC_EQUAL = 3
# DW5 = interval (bits 0..15) | retry_count (bits 16..27). A retry_count of 0 means
# poll forever, which wedges the queue on a mismatch and cost the first run of this
# test its RB_PRIV=0 case. Bound it so a mismatch retires by itself while still taking
# far longer than a match.
INTERVAL = 0x40
RETRY_COUNT = 0x80
DW5 = (INTERVAL & 0xffff) | ((RETRY_COUNT & 0xfff) << 16)


def poll_pkt(reg_index, reference, mask, func=FUNC_EQUAL, mem_poll=0):
    hdr = (SDMA_OP_POLL_REGMEM | (0 << 8) | ((func & 7) << 28)
           | ((mem_poll & 1) << 31))
    return [hdr, (reg_index << 2) & 0xffffffff, 0,
            reference & 0xffffffff, mask & 0xffffffff, DW5]


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)
    MAGIC = 0x5AA50FF0
    print('oracle target %s (host idx 0x%x -> SRC 0x%x), magic 0x%08x' % (
        args.target, tidx, tidx << 2, MAGIC))

    out = {}
    for priv in (1, 0):
        # re-establish a working privileged queue each round, then flip the single
        # bit under test, so a wedge in one round cannot silently void the next
        p.setup(rb_priv=1)
        if priv == 0:
            p.reg('_RB_CNTL').update(rb_priv=0)
        treg.write(MAGIC)
        ok, _ = p.submit([SDMA_OP_NOP])
        if not ok:
            print('  RB_PRIV=%d: queue not running -- inconclusive' % priv)
            out[priv] = None
            continue
        # match: (reg & 0xffffffff) == MAGIC   -> should retire at once
        okm, tm = p.submit(poll_pkt(tidx, MAGIC, 0xffffffff), timeout_s=3.0)
        # mismatch: same register, wrong reference -> must spin
        okx, tx = p.submit(poll_pkt(tidx, MAGIC ^ 0xffffffff, 0xffffffff), timeout_s=3.0)
        sep = (tx > tm * 3) and (tx - tm) > 0.001
        print('  RB_PRIV=%d  match %.2f ms (consumed=%s) | mismatch %.2f ms '
              '(consumed=%s) | separable=%s' % (
                  priv, 1000 * tm, okm, 1000 * tx, okx, sep))
        out[priv] = (okm, tm, okx, tx, sep)

    print('\n=== result')
    if not out.get(1) or not out.get(0):
        print('  inconclusive: a control failed')
    elif out[1][4] and out[0][4]:
        print('  ORACLE ON AN UNPRIVILEGED QUEUE: POLL_REGMEM timing reveals register')
        print('  content with RB_PRIV=0, on the same queue where SRBM_WRITE and')
        print('  POLL_REG_WRITE_MEM are both refused. FINDING-10 transplanted to a')
        print('  user-submittable engine.')
    elif out[1][4] and not out[0][4]:
        print('  gated: the timing difference exists only with RB_PRIV=1, so the poll')
        print('  is covered by the same gate as the other register accesses.')
    else:
        print('  no timing separation even with RB_PRIV=1 -- the packet did not poll a')
        print('  register as intended; no conclusion.')
    print('\nqueue alive: %s' % p.submit([SDMA_OP_NOP])[0])
    return 0


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