#!/usr/bin/env python3
"""Audit MEC's PM4 packet surface — the only user-submittable firmware surface left.

Why MEC: `../psp_out/SYNTHESIS-CROSS-FINDING-COMPOUNDS.md` reordered the targets.
DMUB's inbox is kernel-only, so none of its ten findings is user-reachable; SDMA's
queues are user-submittable but every register path is gated on `RB_PRIV` [HW]. PM4 is
submitted directly by unprivileged user queues and parsed by MEC, and the CP enforces a
privileged/unprivileged packet distinction in firmware — the same structural shape as
`RB_PRIV`, on the one engine where a failure is a genuine user->kernel escalation.

Builds on `../psp_out/MEC-PM4-FOOTHOLD.md`, which established:
  * MEC is plaintext RV64 (RS64), two `$PS1` sub-images (the second at 0x41300);
  * the **PM4 type-3 dispatch table is at file 0x414d8, stride 0x18**, indexed by
    opcode, cross-validated against five consecutive opcodes;
  * table entries are *case labels* inside large dispatchers, not function starts;
  * **every host-controlled byte enters through the packet data port `0x6F8(tp)`** —
    MEC's equivalent of SDMA's `r1`, and therefore the taint source.

This tool does what the foothold doc listed as next steps:
  table   dump every populated slot, named against the kernel's PACKET3_* enum, and
          report the per-entry second dword (the suspected length bound)
  bounds  test the "second dword is a length bound" hypothesis across all slots
          rather than the four that were eyeballed

Naming is the validation: if the table base/stride were wrong, opcodes would not land
on their published names.

usage:
  ./mec_audit.py table  ../psp_out/mec_11_0_0.bin
  ./mec_audit.py bounds ../psp_out/mec_11_0_0.bin
"""
import argparse
import collections
import os
import re
import struct
import sys

PM4_TABLE_FILE_OFF = 0x415d8      # corrected: MEC-PM4-FOOTHOLD.md said 0x414d8, which
                                  # is 0x100 low (the nested $PS1 offset class of error
                                  # this tree keeps hitting). All five anchor opcodes
                                  # from that doc verify exactly at 0x415d8.
PM4_STRIDE = 0x18
KERNEL_PM4 = ('/home/user/src/linux-7.2.5/drivers/gpu/drm/amd/amdgpu/nvd.h',
              '/home/user/src/linux-7.2.5/drivers/gpu/drm/amd/amdgpu/soc15d.h')


def pm4_names():
    names = {}
    for path in KERNEL_PM4:
        try:
            for line in open(path):
                m = re.match(r'#define\s+PACKET3_([A-Z0-9_]+)\s+(0x[0-9A-Fa-f]+)\s*$',
                             line)
                if m:
                    names.setdefault(int(m.group(2), 16), m.group(1))
        except OSError:
            pass
    return names


def load(path):
    d = open(path, 'rb').read()
    ucode_size, ucode_off = struct.unpack_from('<II', d, 0x14)
    return d, ucode_off, ucode_size


def entries(d, lo=0, hi=0x100):
    """(opcode, handler, second_dword, rest) for slots that look populated."""
    out = []
    for op in range(lo, hi):
        off = PM4_TABLE_FILE_OFF + op * PM4_STRIDE
        if off + PM4_STRIDE > len(d):
            break
        w = struct.unpack_from('<6I', d, off)
        out.append((op, w[0], w[1], w[2:]))
    return out


def in_code(d, va):
    return 0x1000 <= va < len(d)


def cmd_table(path):
    d, ucode_off, ucode_size = load(path)
    names = pm4_names()
    print('%s: ucode 0x%x bytes at 0x%x; PM4 table at 0x%x stride 0x%x' % (
        os.path.basename(path), ucode_size, ucode_off, PM4_TABLE_FILE_OFF, PM4_STRIDE))
    print('  %-6s %-34s %-10s %-10s %s' % ('opcode', 'packet', 'handler', 'dword1',
                                           'rest'))
    named = pop = 0
    for op, handler, second, rest in entries(d):
        if not in_code(d, handler):
            continue
        pop += 1
        nm = names.get(op)
        if nm:
            named += 1
        print('  0x%02x   %-34s 0x%08x 0x%08x %s' % (
            op, nm or '<not in kernel enum>', handler, second,
            ' '.join('%08x' % r for r in rest)))
    print('\n%d populated slots, %d carry a published PACKET3_* name (%.0f%%)' % (
        pop, named, 100.0 * named / max(1, pop)))
    return pop, named


def cmd_bounds(path):
    """Is the second dword a per-packet length bound?

    MEC-PM4-FOOTHOLD.md noticed the shape (0x01 << 16) | N on four entries and
    suggested a table-driven length check. Four samples is not a hypothesis test, so
    look at the distribution across every populated slot and compare against the PM4
    packets whose length is actually fixed by the ABI.
    """
    d, _, _ = load(path)
    names = pm4_names()
    rows = [(op, h, s, r) for op, h, s, r in entries(d) if in_code(d, h)]
    hi = collections.Counter(s >> 16 for _, _, s, _ in rows)
    print('second-dword high half (bits 31:16) across %d populated slots:' % len(rows))
    for k, c in hi.most_common():
        print('   0x%04x  x%d' % (k, c))
    print('\nlow half (bits 15:0) — candidate dword counts:')
    lo = collections.Counter(s & 0xffff for _, _, s, _ in rows)
    for k, c in lo.most_common(12):
        print('   0x%04x  x%d' % (k, c))
    print('\nper-packet, for packets with an ABI-fixed length:')
    for op, h, s, _ in rows:
        nm = names.get(op)
        if nm in ('NOP', 'SET_BASE', 'CLEAR_STATE', 'DISPATCH_DIRECT', 'MAP_QUEUES',
                  'UNMAP_QUEUES', 'QUERY_STATUS', 'RUN_LIST', 'MAP_PROCESS_VM',
                  'INDIRECT_BUFFER', 'WRITE_DATA', 'ATOMIC_MEM', 'RELEASE_MEM'):
            print('   0x%02x %-22s dword1 = 0x%08x  (hi 0x%04x, lo %d)' % (
                op, nm, s, s >> 16, s & 0xffff))
    return rows


PORT_RE = re.compile(r'^(?:ld|lwu|lw|lhu|lbu)\s+(\w+),\s*0x6f8\(tp\)')


def decode(d):
    """RV64 decode of the first sub-image, addressed by file offset (handler values
    in the PM4 table are file offsets, so this keeps one address space)."""
    import capstone as cs
    md = cs.Cs(cs.CS_ARCH_RISCV, cs.CS_MODE_RISCV64 | cs.CS_MODE_RISCVC)
    md.skipdata = True
    ins = {}
    for x in md.disasm(d[0x200:0x41300], 0x200):
        ins[x.address] = '%s %s' % (x.mnemonic, x.op_str)
    return ins


def cmd_taint(path):
    """Track host packet data from the PM4 data port to address-forming sinks.

    `MEC-PM4-FOOTHOLD.md` identified `0x6F8(tp)` as the packet data port: repeated
    reads of the same address return successive dwords, i.e. a FIFO pop. That is MEC's
    equivalent of SDMA's `r1`, so it is the taint source.

    Clobber-aware by construction: a `jal` kills every caller-saved register, because
    the opposite assumption is what manufactured a false positive in the SDMA audit
    (`../f32/AUDIT-SDMA-01.md`) -- two calls sat between a load and a store and the
    value did not survive them.
    """
    d, _, _ = load(path)
    ins = decode(d)
    addrs = sorted(ins)
    caller_saved = {'ra', 't0', 't1', 't2', 't3', 't4', 't5', 't6',
                    'a0', 'a1', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7'}
    ports = [a for a in addrs if PORT_RE.match(ins[a])]
    print('%s: %d instructions decoded, %d reads of the PM4 data port 0x6f8(tp)' % (
        os.path.basename(path), len(ins), len(ports)))
    sinks = []
    for start in ports:
        taint = {PORT_RE.match(ins[start]).group(1)}
        i = addrs.index(start)
        for a in addrs[i + 1:i + 120]:
            t = ins[a]
            op = t.split()[0]
            parts = t.split(None, 1)
            ops = [o.strip() for o in parts[1].split(',')] if len(parts) > 1 else []
            if op in ('jal', 'jalr'):
                taint -= caller_saved
                if not taint:
                    break
                continue
            if op in ('ret', 'ecall', 'ebreak', 'j', 'jr', 'tail'):
                # Stop at unconditional transfers. Walking through a `j` makes the
                # next basic block look like a continuation of this one, which
                # manufactured a false "packet dword -> MMIO address" lead at
                # 0x3af08 (the block after the jump gets its base from elsewhere).
                break
            m = re.search(r'(-?(?:0x)?[0-9a-f]*)\((\w+)\)$', t)
            if m and m.group(2) in taint and op[0] in 'ls':
                sinks.append((start, a, t))
            if op in ('beq', 'bne', 'blt', 'bge', 'bltu', 'bgeu', 'beqz', 'bnez'):
                continue
            if ops:
                dst = ops[0]
                if len(ops) > 1 and any(r in taint for r in ops[1:]):
                    taint.add(dst)
                elif dst in taint and not (m and m.group(2) in taint):
                    taint.discard(dst)
    print('\nhost packet data reaching an address-forming sink: %d' % len(sinks))
    seen = set()
    for src, a, t in sinks:
        if (a, t) in seen:
            continue
        seen.add((a, t))
        print('   port 0x%05x -> 0x%05x  %s' % (src, a, t))
    if not sinks:
        print('   none within 120 instructions of a port read')
    return sinks


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('cmd', choices=['table', 'bounds', 'taint'])
    ap.add_argument('image')
    a = ap.parse_args()
    if a.cmd == 'table':
        cmd_table(a.image)
    elif a.cmd == 'bounds':
        cmd_bounds(a.image)
    else:
        cmd_taint(a.image)
    return 0


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