path: root/dmcub/hw/dmcub_regseq_sweep.py

#!/usr/bin/env python3
"""Sweep DMCUB's PSP-owned windows (CW0 code, CW1 data) with the FINDING-10 oracle.

FINDING-10 (`../AUDIT-DMCUB-05-REGSEQ.md`) gives one bit per inbox command, which is
far too slow to dump 283 KB of CW0 bit-by-bit. Two tricks make a full sweep practical:

1. **Verify instead of read.** The oracle tests `(v & mask) == expected` for arbitrary
   mask/expected, so with `mask = 0xffffffff` and `expected =` the word from the signed
   image on disk, **one command covers a whole 32-bit word**: fast means "this word is
   exactly what public linux-firmware contains", slow means "this word differs".
   Only differing words cost 32 probes. At `iters = 300` a match costs ~0.16 ms.
2. **Classify instead of read** for CW1, where there is no ground truth: two commands
   per word test `== 0` and `== 0xffffffff`, which is enough to map the window's
   structure and find the words worth recovering in full.

The point of the sweep is to answer: does the PSP-loaded image in the TMR differ from
the public firmware file, and is there anything in these host-denied windows that is
*not* public?

Every address probed is clamped to the window tops reported by dmcub_bringup, so the
firmware never touches unbacked memory (../../psp_out/CONSTRAINTS.md).

usage:
  python3 dmcub_regseq_sweep.py --state state.json --apply --cw0
  python3 dmcub_regseq_sweep.py --state state.json --apply --cw1 --stride 4
"""
import argparse
import json
import os
import struct
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import dmcub_hwregs as H              # noqa: E402
from dmcub_poc_regseq import RegSeq, addr_to_idx, frame_reg_wait   # noqa: E402

CW0_BASE = 0x60000000
CW1_BASE = 0x61000000


class Sweeper:
    def __init__(self, rs, iters=300):
        self.rs = rs
        self.rs.iters = iters
        self.iters = iters
        self.n_cmd = 0

    def calibrate(self, addr, truth):
        """match/mismatch separation on a word whose value we already know."""
        idx = addr_to_idx(addr)
        _, t_ok = self.rs.submit_timed(frame_reg_wait(idx, 0xffffffff, truth, self.iters))
        _, t_no = self.rs.submit_timed(
            frame_reg_wait(idx, 0xffffffff, truth ^ 0xffffffff, self.iters))
        self.n_cmd += 2
        if not (t_no > t_ok * 2 and (t_no - t_ok) > 0.0002):
            return None
        thr = t_ok + (t_no - t_ok) / 2.0
        print('  calibration: match %.2f ms, mismatch %.2f ms -> threshold %.2f ms'
              % (1000 * t_ok, 1000 * t_no, 1000 * thr))
        return thr

    def equals(self, addr, value, thr, tries=1):
        """True if the word at addr equals value.

        A single command is enough for the common case, but the fast path is only
        ~0.16 ms and host-side scheduling noise occasionally pushes it past the
        threshold. Measured false-positive rate on a full CW0 pass: 371/70979 =
        0.5%% of words looked like mismatches and every one of them turned out to
        equal the image when read bit-by-bit. So a *mismatch* is only believed if it
        reproduces `tries` times; a single fast reading is conclusive for equality.
        """
        for _ in range(max(1, tries)):
            _, t = self.rs.submit_timed(
                frame_reg_wait(addr_to_idx(addr), 0xffffffff, value, self.iters))
            self.n_cmd += 1
            if t < thr:
                return True
        return False

    def recover(self, addr, thr, tries=4):
        """Read a word, dual-polarity with retries. Returns (value, unreliable_bits).

        A single-polarity version of this produced garbage that looked like real data:
        sweeping the CW0 tail it reported words such as 0xffbbdcff / 0x06580007, which
        I nearly wrote up as TMR residue. Re-reading each bit in both polarities with
        retries showed the true contents are a stable 0xffffffff / 0x00000000. One
        jitter-slow reading silently clears a set bit, so cross-checking polarities is
        mandatory, not optional.
        """
        val, bad = 0, []
        for b in range(32):
            m = 1 << b

            def fast(expected):
                for _ in range(tries):
                    _, t = self.rs.submit_timed(
                        frame_reg_wait(addr_to_idx(addr), m, expected, self.iters))
                    self.n_cmd += 1
                    if t < thr:
                        return True
                return False

            one, zero = fast(m), fast(0)
            if one == zero:
                bad.append(b)
            elif one:
                val |= m
        return val, bad


def image_reader(path):
    raw = open(path, 'rb').read()
    ucode_off = struct.unpack_from('<I', raw, 0x18)[0]
    body = ucode_off + 0x100          # nested $PS1 header
    n = len(raw) - body

    def word(va):
        off = va - CW0_BASE + body
        if body <= off <= len(raw) - 4:
            return struct.unpack_from('<I', raw, off)[0]
        return None
    return word, n


def sweep_cw0(sw, thr, image, cw0_top, stride, max_recover):
    word_of, img_len = image_reader(image)
    end = min(CW0_BASE + img_len, cw0_top)
    total = (end - CW0_BASE) // 4
    print('\n=== CW0 verify-sweep: 0x%08x..0x%08x (%d words, stride %d word(s))'
          % (CW0_BASE, end, total, stride))
    print('    one command per word, comparing against the signed image on disk')
    diffs, checked = [], 0
    t0 = time.time()
    for i in range(0, total, stride):
        va = CW0_BASE + 4 * i
        truth = word_of(va)
        if truth is None:
            continue
        checked += 1
        if not sw.equals(va, truth, thr, tries=4):
            diffs.append((va, truth))
        if checked % 5000 == 0:
            print('    %6d/%d words, %d differ, %.0f s elapsed'
                  % (checked, total // stride, len(diffs), time.time() - t0))
    print('    checked %d words in %.0f s: %d match the public image, %d DIFFER'
          % (checked, time.time() - t0, checked - len(diffs), len(diffs)))
    shown = diffs[:max_recover]
    if shown:
        print('\n    recovering the differing words (%d of %d):' % (len(shown), len(diffs)))
        print('    %-12s %-12s %-12s' % ('addr', 'image(public)', 'actual(in TMR)'))
        for va, truth in shown:
            got, bad = sw.recover(va, thr)
            print('    0x%08x   0x%08x   0x%08x%s' % (
                va, truth, got, '' if not bad else '  (unreliable bits %s)' % bad))
    return diffs, checked


def sweep_cw1(sw, thr, cw1_top, stride, max_recover):
    total = (cw1_top - CW1_BASE) // 4
    print('\n=== CW1 classification sweep: 0x%08x..0x%08x (%d words, stride %d)'
          % (CW1_BASE, cw1_top, total, stride))
    print('    two commands per word (==0, ==0xffffffff); no ground truth exists here')
    zero = ones = other = 0
    interesting = []
    t0 = time.time()
    n = 0
    for i in range(0, total, stride):
        va = CW1_BASE + 4 * i
        n += 1
        if sw.equals(va, 0, thr, tries=3):
            zero += 1
        elif sw.equals(va, 0xffffffff, thr, tries=3):
            ones += 1
        else:
            other += 1
            if len(interesting) < 4096:
                interesting.append(va)
        if n % 5000 == 0:
            print('    %6d/%d words: %d zero, %d all-ones, %d other, %.0f s'
                  % (n, total // stride, zero, ones, other, time.time() - t0))
    print('    %d words sampled in %.0f s: zero=%d all-ones=%d other=%d'
          % (n, time.time() - t0, zero, ones, other))
    if interesting:
        lo, hi = interesting[0], interesting[-1]
        print('    non-trivial words span 0x%08x..0x%08x' % (lo, hi))
        shown = interesting[:max_recover]
        print('\n    recovering %d non-trivial words:' % len(shown))
        for va in shown:
            got, bad = sw.recover(va, thr)
            print('    0x%08x = 0x%08x%s' % (
                va, got, '' if not bad else '  (unreliable bits %s)' % bad))
    return zero, ones, other, interesting


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--state', required=True)
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--cw0', action='store_true')
    ap.add_argument('--cw1', action='store_true')
    ap.add_argument('--stride', type=int, default=1, help='in words')
    ap.add_argument('--iters', type=lambda s: int(s, 0), default=300)
    ap.add_argument('--max-recover', type=int, default=16)
    ap.add_argument('--inst', type=int, default=0)
    args = ap.parse_args()

    st = json.load(open(args.state))
    image = st['image']
    cw0_top = st['dmub_regions']['cw0'][1]
    cw1_top = st['dmub_regions']['cw1'][1]

    from tinygrad import Device
    adev = Device["AMD"].iface.dev_impl
    devnull = open(os.devnull, 'w')
    regs = H.Regs(adev=adev, inst=args.inst, dry_run=not args.apply, log=devnull)
    if regs.validate(verbose=False)[0] == 'none':
        raise SystemExit('no usable DCN register path')
    if not (regs.get('DMCUB_CNTL', 'DMCUB_ENABLE')
            and not regs.get('DMCUB_CNTL2', 'DMCUB_SOFT_RESET')):
        raise SystemExit('DMCUB not running -- dmcub_bringup.py --apply first')

    rs = RegSeq(regs, adev, st['paddr']['mailbox'], iters=args.iters,
                dry_run=not args.apply)
    sw = Sweeper(rs, iters=args.iters)
    word_of, _ = image_reader(image)
    thr = sw.calibrate(CW0_BASE, word_of(CW0_BASE))
    if thr is None:
        raise SystemExit('oracle not separable at iters=%d; raise --iters' % args.iters)

    print('  windows from bringup: CW0 top 0x%08x, CW1 top 0x%08x' % (cw0_top, cw1_top))
    if args.cw0:
        sweep_cw0(sw, thr, image, cw0_top, args.stride, args.max_recover)
    if args.cw1:
        sweep_cw1(sw, thr, cw1_top, args.stride, args.max_recover)
    print('\ntotal oracle commands issued: %d' % sw.n_cmd)
    print('dmcub alive: scratch0=0x%x' % regs.read('DMCUB_SCRATCH0'))
    return 0


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