aboutsummaryrefslogtreecommitdiffstats
path: root/hw/invoke_newtas_probe.py
blob: 79eb4de80eef9274c51b37102ef1c515ee1c075b (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
#!/usr/bin/env python3
"""Can the host actually INVOKE the twelve newly-discovered TAs, or only load them?

`fw/AUDIT-NEW-TAS.md` establishes that this Navi 31 authenticates twelve TA types the
project never audited, and that all twelve carry class `TA` (the invokable kind) rather
than ASD's `DR`. That is an inference from the image plus a successful load. ASD is the
standing precedent for why that is not enough: it loads cleanly and then refuses **every**
`INVOKE_CMD` with `psp_status 0x4`
(`FINDING-ASD-IS-DRIVER-CLASS-NOT-INVOKABLE.md`).

So this measures it. For each TA: LOAD_TA -> INVOKE_CMD(ta_cmd_id=0) -> UNLOAD_TA.

Command id 0 is deliberate. The goal is to separate "the PSP will not dispatch to this TA
at all" from "the TA received a command and answered", which the status alone shows. An
unimplemented id is the least invasive way to ask; this is not a command-surface fuzz and
it does not try to make any TA do work.

Controls, both required:
  POS  HDCP from this machine's own psp_13_0_0_ta.bin -- must invoke (not 0x4).
  NEG  ASD, driver class -- must refuse with 0x4, reproducing the known behaviour.
       Without it, a run where everything returns 0x4 would look like a finding
       instead of a broken harness.

Successful loads allocate a session and this rig hits a session wall around 45 loads per
boot, so run under AM_RESET=1; the set here is ~15 loads.

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad invoke_newtas_probe.py --dir /root/psp/contall
"""
import argparse
import glob
import os
import struct
import sys

NEW = ['ta_16-0-0-9_9e9044cf', 'ta_42-0-0-6_1c874c3c', 'ta_19-0-0-8_7e8be282',
       'ta_22-0-0-77_7c426e7f', 'ta_24-0-0-70_5cd4afda', 'ta_50-0-0-1_cbd47c59',
       'ta_25-0-0-81_dc596ba6', 'ta_34-0-0-20_e391bdbb', 'ta_7-0-2-19_f368676e',
       'ta_29-0-0-4_ce34f538', 'ta_30-0-0-38_ec5dee6d', 'ta_32-0-0-20_20a3d1a4']
NAMES = dict(zip(NEW, ['AUC', 'CSU', 'FP', 'GCP', 'MERIT', 'OC',
                       'OPM', 'POWERPLAY', 'RAP', 'TMZ', 'WSPE', 'XGMI']))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--dir', default='/root/psp/contall')
    ap.add_argument('--ta-bin', default='/lib/firmware/amdgpu/psp_13_0_0_ta.bin')
    ap.add_argument('--asd', default='/root/psp/asd_full.bin')
    ap.add_argument('--cmd-ids', default='0,1,2',
                    help='ids to try; a TA that answers ANY of them is invokable')
    args = ap.parse_args()

    from tinygrad import Device
    import tinygrad.runtime.autogen.am.am as am

    adev = Device["AMD"].iface.dev_impl
    psp = adev.psp

    def alloc(n):
        pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
        return pa, adev.paddr2mc(pa)

    def load(blob, cmd_id=None):
        ap_, amc = alloc(0x80000)
        cp, cmc = alloc(0x8000)
        n = (len(blob) + 3) & ~3
        adev.vram.view(ap_, n)[:] = blob + b'\x00' * (n - len(blob))
        adev.gmc.flush_hdp()
        c = am.struct_psp_gfx_cmd_resp(cmd_id=cmd_id or am.GFX_CMD_ID_LOAD_TA)
        t = c.cmd.cmd_load_ta
        t.app_phy_addr_lo, t.app_phy_addr_hi = amc & 0xFFFFFFFF, amc >> 32
        t.app_len = len(blob)
        t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = cmc & 0xFFFFFFFF, cmc >> 32
        t.cmd_buf_len = 0x8000
        try:
            r = psp._ring_submit(c)
            return r.resp.status & 0xFFFFFFFF, r.resp.session_id, cp
        except Exception as e:
            try:
                return int(str(e).strip().split()[-1]) & 0xFFFFFFFF, None, cp
            except Exception:
                return -1, None, cp

    def invoke(sid, shm, cmd_id):
        buf = bytearray(0x2000)
        struct.pack_into('<I', buf, 0, cmd_id)          # many TAs read msg[0] too
        adev.vram.view(shm, len(buf))[:] = bytes(buf)
        adev.gmc.flush_hdp()
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
        iv = c.cmd.cmd_invoke_cmd
        iv.session_id, iv.ta_cmd_id = sid, cmd_id
        try:
            r = psp._ring_submit(c)
            st = r.resp.status & 0xFFFFFFFF
        except Exception as e:
            try:
                st = int(str(e).strip().split()[-1]) & 0xFFFFFFFF
            except Exception:
                st = -1
        adev.gmc.flush_hdp()
        out = bytes(adev.vram.view(shm, 16)[:])
        return st, out

    def unload(sid):
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
        c.cmd.cmd_unload_ta.session_id = sid
        try:
            return psp._ring_submit(c).resp.status & 0xFFFFFFFF
        except Exception:
            return -1

    ids = [int(x, 0) for x in args.cmd_ids.split(',')]

    def run(tag, blob, load_cmd=None):
        """Calibrated on this machine's own bundle TAs:
             0x4        -> the PSP will not dispatch to this TA at all (ASD/driver class)
             0xffff0000 -> dispatched; the TA rejected the command id
             0x0        -> dispatched and accepted
           So "invokable" means: at least one id returns something other than 0x4."""
        st, sid, shm = load(blob, load_cmd)
        if st != 0:
            print('  %-12s LOAD failed %#x' % (tag, st), flush=True)
            return None
        res = []
        for cid in ids:
            ist, out = invoke(sid, shm, cid)
            res.append((cid, ist, out[:8].hex()))
        unload(sid)
        invokable = any(r[1] != 4 for r in res)
        print('  %-12s %-26s %s'
              % (tag, 'INVOKABLE' if invokable else 'refused (0x4 on every id)',
                 ' '.join('%d:%#x' % (c, s2) for c, s2, _ in res)), flush=True)
        return 0 if invokable else 4

    # ---- controls -------------------------------------------------------
    print('controls:')
    import subprocess
    tab = args.ta_bin
    if not os.path.exists(tab) and os.path.exists(tab + '.zst'):
        subprocess.run(['zstd', '-d', '-f', '-o', '/tmp/ta.bin', tab + '.zst'],
                       check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        tab = '/tmp/ta.bin'
    pos = None
    data = open(tab, 'rb').read()
    p = 0
    while True:
        m = data.find(b'$PS1', p)
        if m < 0:
            break
        p = m + 4
        cont = m - 0x10
        if cont < 0:
            continue
        n = struct.unpack_from('<I', data, cont + 0x14)[0]
        if not (0 < n < 0x400000) or cont + 0x200 + n > len(data):
            continue
        if data[cont + 0x100:cont + 0x104] != b'\x00\x02\x00\x00':
            continue
        r = run('POS bundle#%d' % (p // max(n, 1) + 1), data[cont:cont + 0x200 + n])
        # bundle TA #1 is ASD (driver class) and refuses everything -- keep going
        # until one is invokable, which is the control we actually need.
        if r is not None and r != 4:
            pos = r
            break
    if pos is None or pos == 4:
        print('  POSITIVE CONTROL FAILED (bundle TA not invokable) -- stopping.')
        return 1
    if os.path.exists(args.asd):
        neg = run('NEG asd(DR)', open(args.asd, 'rb').read(), am.GFX_CMD_ID_LOAD_ASD)
        if neg != 4:
            print('  note: ASD did not return 0x4; the 0x4 signature may not be reliable')

    print('\nnewly discovered TAs:')
    for stem in NEW:
        hits = glob.glob(os.path.join(args.dir, stem + '*.bin'))
        if not hits:
            print('  %-12s (container not found)' % NAMES[stem])
            continue
        run(NAMES[stem], open(hits[0], 'rb').read())
    return 0


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