|
| 1 | +"""DylibTransport — talk to libkkemu.dylib (or libkkemu.so) over FFI ringbuffers. |
| 2 | +
|
| 3 | +This is the same firmware the standalone ``kkemu`` UDP binary runs, but loaded |
| 4 | +in-process. Two transports cover the two ringbuffer pairs the dylib exposes: |
| 5 | +
|
| 6 | +* iface 0 (main): rb_main_in / rb_main_out — host ↔ firmware protocol |
| 7 | +* iface 1 (debug): rb_debug_in / rb_debug_out — DebugLink |
| 8 | +
|
| 9 | +The vault uses this same FFI surface from Bun. Adding a Python transport that |
| 10 | +mirrors it lets ``python-keepkey`` exercise the firmware contract that the |
| 11 | +dylib path imposes — most importantly, the *caller-driven polling* model: |
| 12 | +nothing happens inside the firmware until the host calls ``kkemu_poll``. UDP |
| 13 | +hides this behind a thread inside ``kkemu``; the dylib does not. |
| 14 | +
|
| 15 | +Usage |
| 16 | +----- |
| 17 | +:: |
| 18 | +
|
| 19 | + from keepkeylib.transport_dylib import DylibState, DylibTransport |
| 20 | +
|
| 21 | + state = DylibState.get_or_init('/path/to/libkkemu.dylib') |
| 22 | + main_transport = DylibTransport(state, iface=0) |
| 23 | + debug_transport = DylibTransport(state, iface=1) |
| 24 | + client = KeepKeyDebugClient(main_transport) |
| 25 | + client.set_debuglink(DebugLink(debug_transport)) |
| 26 | +
|
| 27 | +A *single* ``DylibState`` is shared between the two transports — the dylib's |
| 28 | +``kkemu_init`` may only be called once per process. Re-initialising means |
| 29 | +restarting the test process (or factory-resetting via ``reset_flash``). |
| 30 | +""" |
| 31 | + |
| 32 | +from __future__ import print_function |
| 33 | + |
| 34 | +import ctypes |
| 35 | +import os |
| 36 | +import struct |
| 37 | +import threading |
| 38 | +import time |
| 39 | + |
| 40 | +from .transport import Transport, ConnectionError |
| 41 | + |
| 42 | + |
| 43 | +# ── Dylib singleton ───────────────────────────────────────────────────────── |
| 44 | + |
| 45 | + |
| 46 | +PACKET_SIZE = 64 |
| 47 | +FLASH_SIZE = 1 << 20 # 1 MB |
| 48 | + |
| 49 | +# Max time we'll spin in kkemu_poll() looking for a frame on this iface. |
| 50 | +# Has to cover firmware-internal busy-loops (confirm_helper polls usbPoll |
| 51 | +# in a tight C loop — we just need the next outbound frame to land). |
| 52 | +_POLL_TIMEOUT_S = 30.0 |
| 53 | +_POLL_QUANTUM_S = 0.001 # 1 ms — keep latency low without burning CPU |
| 54 | + |
| 55 | + |
| 56 | +class DylibState(object): |
| 57 | + """Process-wide ``libkkemu.dylib`` handle. |
| 58 | +
|
| 59 | + Holds the ctypes binding and the (locked) flash buffer. Only one instance |
| 60 | + is allowed per process because ``kkemu_init`` is single-shot. Use |
| 61 | + :func:`get_or_init` rather than the constructor. |
| 62 | + """ |
| 63 | + |
| 64 | + _instance = None |
| 65 | + _lock = threading.Lock() |
| 66 | + |
| 67 | + def __init__(self, dylib_path): |
| 68 | + if not os.path.exists(dylib_path): |
| 69 | + raise ConnectionError("dylib not found: %s" % dylib_path) |
| 70 | + |
| 71 | + self.lib = ctypes.CDLL(dylib_path) |
| 72 | + |
| 73 | + self.lib.kkemu_init.argtypes = [ctypes.c_void_p, ctypes.c_size_t] |
| 74 | + self.lib.kkemu_init.restype = ctypes.c_int |
| 75 | + |
| 76 | + self.lib.kkemu_shutdown.argtypes = [] |
| 77 | + self.lib.kkemu_shutdown.restype = None |
| 78 | + |
| 79 | + self.lib.kkemu_poll.argtypes = [] |
| 80 | + self.lib.kkemu_poll.restype = ctypes.c_int |
| 81 | + |
| 82 | + self.lib.kkemu_is_running.argtypes = [] |
| 83 | + self.lib.kkemu_is_running.restype = ctypes.c_int |
| 84 | + |
| 85 | + self.lib.kkemu_write.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] |
| 86 | + self.lib.kkemu_write.restype = ctypes.c_int |
| 87 | + |
| 88 | + self.lib.kkemu_read.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] |
| 89 | + self.lib.kkemu_read.restype = ctypes.c_int |
| 90 | + |
| 91 | + self.lib.kkemu_get_display.argtypes = [ |
| 92 | + ctypes.POINTER(ctypes.c_int), |
| 93 | + ctypes.POINTER(ctypes.c_int), |
| 94 | + ] |
| 95 | + self.lib.kkemu_get_display.restype = ctypes.c_void_p |
| 96 | + |
| 97 | + # Allocate flash as 0xFF (erased NOR state). Held by the singleton so |
| 98 | + # GC doesn't free it underneath the firmware's still-live mlock. |
| 99 | + self.flash = (ctypes.c_uint8 * FLASH_SIZE)(*([0xFF] * FLASH_SIZE)) |
| 100 | + |
| 101 | + rc = self.lib.kkemu_init(ctypes.cast(self.flash, ctypes.c_void_p), FLASH_SIZE) |
| 102 | + if rc != 0: |
| 103 | + raise ConnectionError("kkemu_init failed: %d" % rc) |
| 104 | + |
| 105 | + # Single mutex around every FFI call. The dylib's internals aren't |
| 106 | + # thread-safe; main + debug transport may both poll/read concurrently. |
| 107 | + self.io_lock = threading.Lock() |
| 108 | + |
| 109 | + # Pump a few ticks so the firmware finishes its boot sequence (loads |
| 110 | + # storage, draws home screen) before the first test touches it. |
| 111 | + with self.io_lock: |
| 112 | + for _ in range(8): |
| 113 | + self.lib.kkemu_poll() |
| 114 | + |
| 115 | + @classmethod |
| 116 | + def get_or_init(cls, dylib_path): |
| 117 | + """Return the per-process singleton, creating it on first call. |
| 118 | +
|
| 119 | + Subsequent calls ignore ``dylib_path`` — the dylib is single-shot and |
| 120 | + re-loading risks UB (mlock'd flash buffer would dangle). |
| 121 | + """ |
| 122 | + with cls._lock: |
| 123 | + if cls._instance is None: |
| 124 | + cls._instance = cls(dylib_path) |
| 125 | + return cls._instance |
| 126 | + |
| 127 | + def shutdown(self): |
| 128 | + """Tear down the firmware. Used by tests; not safe to re-init after.""" |
| 129 | + with self.io_lock: |
| 130 | + self.lib.kkemu_shutdown() |
| 131 | + |
| 132 | + |
| 133 | +# ── Transport ─────────────────────────────────────────────────────────────── |
| 134 | + |
| 135 | + |
| 136 | +class DylibTransport(Transport): |
| 137 | + """One transport per (DylibState, iface) pair. |
| 138 | +
|
| 139 | + ``iface=0`` is the main protocol channel, ``iface=1`` is DebugLink. |
| 140 | + """ |
| 141 | + |
| 142 | + def __init__(self, state, iface=0, *args, **kwargs): |
| 143 | + if not isinstance(state, DylibState): |
| 144 | + raise TypeError("state must be a DylibState") |
| 145 | + if iface not in (0, 1): |
| 146 | + raise ValueError("iface must be 0 (main) or 1 (debug)") |
| 147 | + |
| 148 | + self.state = state |
| 149 | + self.iface = iface |
| 150 | + self.read_buffer = b"" |
| 151 | + |
| 152 | + # Transport.__init__ calls self._open(); device arg is just metadata. |
| 153 | + super(DylibTransport, self).__init__("dylib:iface=%d" % iface, *args, **kwargs) |
| 154 | + |
| 155 | + # ── Transport hooks ───────────────────────────────────────────────── |
| 156 | + |
| 157 | + def _open(self): |
| 158 | + # Nothing to do — the dylib was opened when DylibState was created. |
| 159 | + pass |
| 160 | + |
| 161 | + def _close(self): |
| 162 | + # Don't shut the dylib down on close; the singleton outlives us. |
| 163 | + self.read_buffer = b"" |
| 164 | + |
| 165 | + def ready_to_read(self): |
| 166 | + # Drive the firmware once so any pending outbound frame surfaces in |
| 167 | + # the ring. When a frame arrives, stash through the SAME path |
| 168 | + # _pump_one uses (strip the leading '?' HID marker before |
| 169 | + # appending). Mixing stripped + unstripped frames in one buffer |
| 170 | + # corrupts multi-frame reassembly: _read_headers would see a stray |
| 171 | + # '?' from one chunk in the middle of contiguous payload bytes |
| 172 | + # from another, and decode the wrong message-type / length. |
| 173 | + self._poll_and_stash() |
| 174 | + return bool(self.read_buffer) |
| 175 | + |
| 176 | + # ── Wire protocol ─────────────────────────────────────────────────── |
| 177 | + |
| 178 | + def _write(self, msg, protobuf_msg): |
| 179 | + """Chunk ``msg`` into 64-byte HID frames and shove them at the firmware. |
| 180 | +
|
| 181 | + ``msg`` already starts with ``"##"`` + msg-type + length (see |
| 182 | + ``Transport.write``). The first chunk needs a leading ``"?"`` marker; |
| 183 | + continuation chunks just get their leading ``"?"`` to round out |
| 184 | + the 64-byte HID report. |
| 185 | + """ |
| 186 | + # 63 bytes per chunk + leading '?' = 64 bytes per HID frame |
| 187 | + for chunk in [msg[i : i + 63] for i in range(0, len(msg), 63)]: |
| 188 | + chunk = chunk + b"\0" * (63 - len(chunk)) |
| 189 | + frame = b"?" + chunk |
| 190 | + assert len(frame) == PACKET_SIZE |
| 191 | + with self.state.io_lock: |
| 192 | + rc = self.state.lib.kkemu_write(frame, PACKET_SIZE, self.iface) |
| 193 | + if rc != 0: |
| 194 | + raise ConnectionError( |
| 195 | + "kkemu_write failed (iface=%d, rc=%d)" % (self.iface, rc) |
| 196 | + ) |
| 197 | + # Pump immediately so the firmware can start consuming this |
| 198 | + # chunk before the next one arrives. Required because the |
| 199 | + # caller (not a daemon) is the only thing driving the FSM. |
| 200 | + self.state.lib.kkemu_poll() |
| 201 | + |
| 202 | + def _read(self): |
| 203 | + """Read one full message — header parse drives chunk reassembly.""" |
| 204 | + try: |
| 205 | + (msg_type, datalen) = self._read_headers(_FrameStream(self)) |
| 206 | + payload = self._read_bytes(datalen) |
| 207 | + return (msg_type, payload) |
| 208 | + except Exception as exc: |
| 209 | + print("DylibTransport._read failed: %s" % exc) |
| 210 | + raise |
| 211 | + |
| 212 | + # ── Internals ─────────────────────────────────────────────────────── |
| 213 | + |
| 214 | + def _read_bytes(self, length): |
| 215 | + """Block until ``length`` payload bytes have been gathered.""" |
| 216 | + deadline = time.time() + _POLL_TIMEOUT_S |
| 217 | + while len(self.read_buffer) < length: |
| 218 | + if time.time() > deadline: |
| 219 | + raise ConnectionError( |
| 220 | + "Timed out reading %d bytes from iface %d" % (length, self.iface) |
| 221 | + ) |
| 222 | + self._pump_one() |
| 223 | + out = self.read_buffer[:length] |
| 224 | + self.read_buffer = self.read_buffer[length:] |
| 225 | + return out |
| 226 | + |
| 227 | + def _pump_one(self): |
| 228 | + """Run one poll/read cycle and back off briefly if no frame arrived. |
| 229 | +
|
| 230 | + Used inside the _read_bytes deadline loop. Sleeps so we don't spin |
| 231 | + a hot CPU loop while waiting on the firmware. |
| 232 | + """ |
| 233 | + if not self._poll_and_stash(): |
| 234 | + time.sleep(_POLL_QUANTUM_S) |
| 235 | + |
| 236 | + def _poll_and_stash(self): |
| 237 | + """Single poll + read; append any frame to read_buffer with '?' |
| 238 | + marker stripped. Returns True if a frame was consumed. |
| 239 | +
|
| 240 | + Shared by ``ready_to_read`` (no sleep) and ``_pump_one`` |
| 241 | + (sleeps on miss). Centralises the strip-the-leading-'?' rule so |
| 242 | + the buffer always contains continuation+payload bytes only. |
| 243 | + """ |
| 244 | + with self.state.io_lock: |
| 245 | + self.state.lib.kkemu_poll() |
| 246 | + buf = (ctypes.c_uint8 * PACKET_SIZE)() |
| 247 | + n = self.state.lib.kkemu_read(buf, PACKET_SIZE, self.iface) |
| 248 | + if n > 0: |
| 249 | + # Drop the leading '?' marker; rest is payload (and HID |
| 250 | + # padding zeros at the tail of the last frame of a short |
| 251 | + # message — _read_headers' magic-character search skips |
| 252 | + # those harmlessly on the next message). |
| 253 | + self.read_buffer += bytes(buf[1:n]) |
| 254 | + return True |
| 255 | + return False |
| 256 | + |
| 257 | + |
| 258 | +class _FrameStream(object): |
| 259 | + """File-like adapter so Transport._read_headers can drive _pump_one.""" |
| 260 | + |
| 261 | + def __init__(self, transport): |
| 262 | + self.transport = transport |
| 263 | + |
| 264 | + def read(self, n): |
| 265 | + return self.transport._read_bytes(n) |
0 commit comments