#!/usr/bin/env python3
"""Audit the gfx11 SDMA F32 firmware's host-packet surface.

SDMA executes packets the driver writes into a ring in host-visible memory, so its
packet handlers are the same class of surface as the DMCUB inbox (FINDING-1/2) and the
DCN4 LSDMA command (FINDING-9) -- except SDMA is loadable on this card, so findings
here are testable on silicon we actually have.

Two facts make a precise taint model possible:

* **The dispatch table is at the tail of TH0.** The last 0x300 bytes are 191 entries of
  `(key << 16) | target`, where `key == (op << 8) | (sub_op << 4)`. That decoding is
  confirmed against `navi10_sdma_pkt_open.h`: 0x0100 = COPY/LINEAR, 0x0110 =
  COPY/TILED, 0x0130 = COPY/SOA, 0x0140 = COPY/LINEAR_SUB_WIND, 0x0400 = INDIRECT,
  0x0600 = TRAP, and so on.
* **`r1` is the taint source.** Upstream's README: "r1 is a magic register that pulls
  dwords off of the command queue when read". So every `ldw rX, [r0,...]`-free read of
  r1 is literally the next host-supplied dword of the packet.

What this screens for, per handler:
  * `REG_IDX`   a register access whose index is not a constant -- if it derives from
                r1 that is a host-chosen register (the SRBM_WRITE shape)
  * `R1_USE`    how many host dwords the handler consumes
  * `LOOP`      backward branches (copy/fill loops) and whether a host dword reaches
                the loop bound

Screening only: it reports what to read, it does not claim bugs. Handlers are followed
linearly from their table entry to the first `ret`/unconditional exit, which is a
deliberate under-approximation, not a closure.

usage:
  ./sdma_audit.py table  ../psp_out/sdma_6_0_0.bin
  ./sdma_audit.py screen ../psp_out/sdma_6_0_0.bin
  ./sdma_audit.py hand   ../psp_out/sdma_6_0_0.bin --key 0x0e00
"""
import argparse
import os
import re
import collections
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import f32_audit as F                                             # noqa: E402

OPS = {0: 'NOP', 1: 'COPY', 2: 'WRITE', 4: 'INDIRECT', 5: 'FENCE', 6: 'TRAP',
       7: 'SEM', 8: 'POLL_REGMEM', 9: 'COND_EXE', 10: 'ATOMIC', 11: 'CONST_FILL',
       12: 'PTEPDE', 13: 'TIMESTAMP', 14: 'SRBM_WRITE', 15: 'PRE_EXE',
       16: 'GPUVM_INV', 17: 'GCR_REQ', 32: 'DUMMY_TRAP'}
SUBOPS = {
    1: {0: 'LINEAR', 1: 'TILED', 3: 'SOA', 4: 'LINEAR_SUB_WIND', 5: 'TILED_SUB_WIND',
        6: 'T2T_SUB_WIND', 7: 'DIRTY_PAGE', 8: 'LINEAR_PHY'},
    2: {0: 'LINEAR', 1: 'TILED'},
    8: {1: 'POLL_REG_WRITE_MEM', 2: 'POLL_DBIT_WRITE_MEM', 3: 'POLL_MEM_VERIFY'},
    12: {0: 'GEN', 1: 'COPY', 2: 'RMW', 3: 'COPY_BACKWARDS'},
    13: {0: 'SET', 1: 'GET', 2: 'GET_GLOBAL'},
}


def load(path, blob='TH0'):
    d = open(path, 'rb').read()
    for name, off, size in F.split(path)[1]:
        if name == blob:
            body = F.slice_blob(d, off, size)
            return body
    raise SystemExit('no blob %s in %s' % (blob, path))


def table(body):
    """(key, target) entries from the tail table; entry 0 is a header, not an entry."""
    n = len(body)
    code = n & ~0xfff
    out = []
    for i in range((n - code) // 4):
        w = struct.unpack_from('<I', body, code + 4 * i)[0]
        key, tgt = w >> 16, w & 0xffff
        if tgt < code // 4:
            out.append((key, tgt))
    return code // 4, out


def name(key):
    op, sub = key >> 8, (key >> 4) & 0xf
    n = OPS.get(op, 'op%d<unpublished>' % op)
    s = SUBOPS.get(op, {}).get(sub)
    if s:
        return '%s/%s' % (n, s)
    return '%s%s' % (n, '' if sub == 0 else '/sub%d' % sub)


def listing(body):
    """address -> instruction text, from the upstream disassembler."""
    out = F.disassemble(body)
    ins = {}
    addr = 0
    for line in out.splitlines():
        if line.startswith('    '):
            ins[addr] = line.strip()
            addr += 1
    return ins


def handler_body(ins, start, limit=400):
    """Linear slice from the entry to the first ret / unconditional b (approximate)."""
    body = []
    a = start
    while a in ins and len(body) < limit:
        t = ins[a]
        body.append((a, t))
        if t == 'ret' or re.match(r'^b\s', t) or t == 'btab':
            break
        a += 1
    return body


def screen_handler(ins, start):
    body = handler_body(ins, start)
    r1 = sum(1 for _, t in body if re.search(r'\br1\b', t))
    # register accesses: upstream prints reg[rX, #imm] / unk[..] / plain [..]
    regs = [t for _, t in body if 'reg[' in t]
    reg_var = [t for t in regs if not re.search(r'reg\[r0,', t)]
    loops = sum(1 for a, t in body
                for m in [re.match(r'^(?:b|cbz|cbnz)\s+\S*?0x([0-9a-f]+)', t)]
                if m and int(m.group(1), 16) <= a)
    unk = sum(1 for _, t in body if t.startswith('dw '))
    return {'len': len(body), 'r1': r1, 'reg': len(regs), 'reg_var': len(reg_var),
            'loops': loops, 'unk': unk, 'body': body}


def cmd_table(path, blob):
    body = load(path, blob)
    ncode, ents = table(body)
    print('%s %s: %d code words, %d dispatch entries' % (
        os.path.basename(path), blob, ncode, len(ents)))
    print('  %-8s %-30s %s' % ('key', 'packet', 'handler'))
    for key, tgt in sorted(ents):
        print('  0x%04x   %-30s 0x%04x' % (key, name(key), tgt))


def cmd_screen(path, blob):
    body = load(path, blob)
    ncode, ents = table(body)
    ins = listing(body)
    print('%s %s: screening %d handlers (linear slice to first ret/branch)' % (
        os.path.basename(path), blob, len(ents)))
    print('  %-8s %-28s %-7s %6s %6s %8s %6s %6s' % (
        'key', 'packet', 'addr', 'insns', 'r1', 'reg-var', 'loops', 'unk'))
    rows = []
    for key, tgt in sorted(ents):
        s = screen_handler(ins, tgt)
        rows.append((key, tgt, s))
        print('  0x%04x   %-28s 0x%04x %6d %6d %8d %6d %6d' % (
            key, name(key), tgt, s['len'], s['r1'], s['reg_var'], s['loops'], s['unk']))
    print('\nhandlers with a non-constant register index (host-chosen register?):')
    hot = [r for r in rows if r[2]['reg_var']]
    for key, tgt, s in hot:
        print('  0x%04x %-28s 0x%04x  %d variable reg access(es)' % (
            key, name(key), tgt, s['reg_var']))
    if not hot:
        print('  none')


def reach(ins, start, limit=250):
    """Worklist walk over a handler, following branch targets.

    Wider than handler_body(), and it deliberately runs into the shared epilogue --
    which is why registers 0x4a14/0x4af8/0x4b58 show up under ~49 of the 52 packets.
    Those are the common path, NOT per-packet gates. Only registers that appear for a
    subset of packets are evidence of anything.
    """
    seen, work, out = set(), [start], []
    while work and len(seen) < limit:
        a = work.pop()
        while a in ins and a not in seen:
            seen.add(a)
            t = ins[a]
            out.append(t)
            m = re.match(r'^(?:cbz|cbnz|bl?)\s+\S*?0x([0-9a-f]+)', t)
            if m:
                work.append(int(m.group(1), 16))
            if t == 'ret' or re.match(r'^b\s', t) or t == 'btab':
                break
            a += 1
    return out


def cmd_gates(path, blob):
    """Which control registers each packet handler consults before acting.

    A binary "is it privilege-gated?" column was the first thing tried here and it was
    wrong: testing only for the SRBM_WRITE idiom (0x4a14 bit 9) reported PTEPDE as
    ungated, when PTEPDE in fact gates on reg 0x5a6c bits 0x101. Different packet
    classes consult different control registers, so the map is reported as-is.
    """
    body = load(path, blob)
    _, ents = table(body)
    ins = listing(body)
    import collections
    allregs = collections.Counter()
    print('  %-8s %-30s %-5s %s' % ('key', 'packet', 'r1', 'control regs consulted'))
    seen = set()
    for key, tgt in sorted(ents):
        if key in seen or key >= 0x1200:
            continue
        seen.add(key)
        txt = reach(ins, tgt)
        j = '\n'.join(txt)
        regs = sorted(set(re.findall(r'ldw r\d+, reg\[r0, #(0x[0-9a-f]+)\]', j)))
        for r in regs:
            allregs[r] += 1
        r1 = sum(1 for t in txt if re.search(r'\br1\b', t))
        shown = ' '.join((reg_name(int(r, 16)) or r).replace('regSDMA0_', '') for r in regs)
        print('  0x%04x %-30s %-5d %s' % (key, name(key), r1, shown or '-'))
    print('\nmost-consulted (>=40 handlers == shared epilogue, not a gate):')
    for r, c in allregs.most_common(10):
        nm = reg_name(int(r, 16)) or '<unpublished, host offset 0x%x>' % (
            (int(r, 16) - F32_REG_BASE) // 4)
        print('   %-8s %-42s %d handlers%s' % (
            r, nm, c, '   <- common path' if c >= 40 else ''))


# --- F32 <-> host register namespace -----------------------------------------
# Solved by structure, not by curve-fitting (an unconstrained delta search scored at
# chance: 25/54 vs a null 95th pct of 24, see AUDIT-SDMA-01). Two anchors fixed it:
#   * TH1 walks the 8 per-queue register blocks with F32 base 0x4b80, stride 0x160
#   * the host per-queue stride is 0x58 dwords (regSDMA0_QUEUE{n}_RB_CNTL)
# 0x160 / 0x58 == 4 exactly, so F32 addresses registers by BYTE offset:
#       F32_index = 4 * host_dword_offset + 0x4980
# Validated: 31/54 observed indices name a regSDMA0_* register against a null model
# whose 95th percentile is 9 and maximum 14, and the top hits are semantically right
# (STATUS_REG x45, CHICKEN_BITS, STATUS2_REG, CNTL, UTCL1_CNTL).
F32_REG_BASE = 0x4980
HOST_OFFSETS = None


def _host_regs(header='/home/user/src/linux-7.2.5/drivers/gpu/drm/amd/include/'
                      'asic_reg/gc/gc_11_0_0_offset.h'):
    global HOST_OFFSETS
    if HOST_OFFSETS is None:
        HOST_OFFSETS = {}
        try:
            off, bidx = {}, {}
            for line in open(header):
                m = re.match(r'#define (reg\w+?)\s+(0x[0-9a-f]+)\s*$', line)
                if m:
                    off[m.group(1)] = int(m.group(2), 16)
                m = re.match(r'#define (reg\w+?)_BASE_IDX\s+(\d+)\s*$', line)
                if m:
                    bidx[m.group(1)] = int(m.group(2))
            for n, o in off.items():
                if n.startswith('regSDMA0') and bidx.get(n) == 0:
                    HOST_OFFSETS.setdefault(o, n)
        except OSError:
            pass
    return HOST_OFFSETS


def reg_name(f32_index):
    """Name an F32 register index, or None. Indices in the 0x3f0..0x43b host-offset
    window resolve to no name: they are inside the SDMA0 aperture but absent from
    AMD's header, i.e. unpublished registers -- the SRBM_WRITE gate lives there."""
    d = f32_index - F32_REG_BASE
    if d < 0 or d % 4:
        return None
    return _host_regs().get(d // 4)


DEST_NONE = ('stw', 'std', 'stb', 'sth', 'b', 'bl', 'cbz', 'cbnz', 'ret', 'nop',
             'btab', 'push', 'pop', 'dw')


def _dest(t):
    """Destination register of an instruction, or None."""
    op = t.split()[0]
    if op in DEST_NONE:
        return None
    m = re.match(r'^\w+\s+(r\d+),', t)
    return m.group(1) if m else None


def _srcs(t):
    """Registers read by an instruction (over-approximate)."""
    regs = re.findall(r'\br\d+\b', t)
    d = _dest(t)
    return [r for r in regs if r != d] if d else regs


def call_target(t):
    m = re.match(r'^bl\s+\S*?0x([0-9a-f]+)', t)
    return int(m.group(1), 16) if m else None


def clobber_sets(ins):
    """Registers each callee may define, transitively. F32 helpers trash scratch
    registers, so a taint that survives a `bl` in a linear scan is an artifact --
    that is exactly how the POLL_MEM_VERIFY false lead was produced."""
    entries = sorted({t for a in ins for t in [call_target(ins[a])] if t is not None})
    own, calls = {}, {}
    for e in entries:
        defs, cs, a = set(), set(), e
        while a in ins and a < e + 600:
            t = ins[a]
            d = _dest(t)
            if d:
                defs.add(d)
            c = call_target(t)
            if c is not None:
                cs.add(c)
            if t == 'ret':
                break
            a += 1
        own[e], calls[e] = defs, cs
    for _ in range(8):                      # fixpoint over the call graph
        changed = False
        for e in entries:
            new = set(own[e])
            for c in calls[e]:
                new |= own.get(c, set())
            if new != own[e]:
                own[e] = new
                changed = True
        if not changed:
            break
    return own


def cmd_taint(path, blob):
    """Track host dwords (reads of r1) to sinks, killing taint across calls."""
    body = load(path, blob)
    _, ents = table(body)
    ins = listing(body)
    clob = clobber_sets(ins)
    print('%s %s: call-clobber-aware taint from r1 (host command queue)' % (
        os.path.basename(path), blob))
    print('  %-8s %-30s %s' % ('key', 'packet', 'sinks reached by host data'))
    total = collections.Counter()
    seen_keys = set()
    for key, tgt in sorted(ents):
        if key in seen_keys or key >= 0x1200:
            continue
        seen_keys.add(key)
        hits = []
        # worklist over (addr, frozenset(tainted regs))
        work, seen = [(tgt, frozenset())], set()
        while work and len(seen) < 3000:
            a, taint = work.pop()
            while a in ins:
                st = (a, taint)
                if st in seen:
                    break
                seen.add(st)
                t = ins[a]
                # sinks: host data used as a register index or a memory base
                m = re.search(r'(reg|mem|unk)?\[(r\d+)', t)
                if m and m.group(2) in taint and m.group(2) != 'r0':
                    hits.append((a, t, 'index/base'))
                c = call_target(t)
                if c is not None:
                    taint = frozenset(taint - clob.get(c, set()))
                else:
                    d = _dest(t)
                    if d:
                        if re.search(r'\br1\b', t) or (set(_srcs(t)) & taint):
                            taint = frozenset(taint | {d})
                        else:
                            taint = frozenset(taint - {d})
                nxt = re.match(r'^(?:cbz|cbnz|b)\s+\S*?0x([0-9a-f]+)', t)
                if nxt:
                    work.append((int(nxt.group(1), 16), taint))
                if t == 'ret' or re.match(r'^b\s', t) or t == 'btab':
                    break
                a += 1
        if hits:
            total[key] = len(hits)
            print('  0x%04x %-30s %d: %s' % (key, name(key), len(hits),
                  '; '.join('0x%04x %s' % (a, t) for a, t, _ in hits[:3])))
    if not total:
        print('  none: no host dword reaches a register index or memory base')
    return total


def cmd_hand(path, blob, key):
    body = load(path, blob)
    _, ents = table(body)
    ins = listing(body)
    for k, tgt in sorted(ents):
        if k == key:
            print('### key 0x%04x %s -> 0x%04x' % (k, name(k), tgt))
            for a, t in handler_body(ins, tgt):
                print('  %04x  %s' % (a, t))
            return
    raise SystemExit('key 0x%x not in the table' % key)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('cmd', choices=['table', 'screen', 'hand', 'gates', 'taint'])
    ap.add_argument('image')
    ap.add_argument('--blob', default='TH0')
    ap.add_argument('--key', type=lambda s: int(s, 0))
    a = ap.parse_args()
    if a.cmd == 'table':
        cmd_table(a.image, a.blob)
    elif a.cmd == 'screen':
        cmd_screen(a.image, a.blob)
    elif a.cmd == 'gates':
        cmd_gates(a.image, a.blob)
    elif a.cmd == 'taint':
        cmd_taint(a.image, a.blob)
    else:
        cmd_hand(a.image, a.blob, a.key)
    return 0


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