feat(transport): add DylibTransport for in-process libkkemu testing

BitHighlander · BitHighlander · commit ebc0db71507a · 2026-04-25T21:13:46.000-05:00
Same firmware as the standalone UDP kkemu binary, loaded in-process via
ctypes. Lets python-keepkey exercise the firmware contract that the
keepkey-vault FFI path imposes — most importantly, the caller-driven
polling model (no daemon thread to call kkemu_poll for you).

- keepkeylib/transport_dylib.py: DylibState (process-wide singleton over
  ctypes-loaded libkkemu) + DylibTransport (one per iface 0/1).
  Pumps kkemu_poll on every read/write so the firmware actually makes
  forward progress on caller turns.
- tests/config.py: KK_TRANSPORT=dylib KK_DYLIB=/path/to/libkkemu.dylib
  routes the same fixture to the FFI transport instead of UDP.
- tests/test_dylib_confirm_flow.py: regression for the confirm-flow
  contract (Initialize, WipeDevice, LoadDevice, GetAddress). Skipped
  unless KK_TRANSPORT=dylib so it won't break the default UDP run.

Reproduces the keepkey-vault hang deterministically: Initialize round-
trips fine, wipe_device hangs because confirm_helper busy-loops on a
ButtonAck the dylib silently consumed but never delivered. Caught in
~10s, no electrobun / bun stack required.

Run:
  cd tests &amp;&amp; KK_TRANSPORT=dylib KK_DYLIB=.../libkkemu.dylib \
    PYTHONPATH=..:../keepkeylib python3 -m pytest \
    test_dylib_confirm_flow.py -v
diff --git a/keepkeylib/transport_dylib.py b/keepkeylib/transport_dylib.py
@@ -0,0 +1,249 @@
+"""DylibTransport — talk to libkkemu.dylib (or libkkemu.so) over FFI ringbuffers.
+
+This is the same firmware the standalone ``kkemu`` UDP binary runs, but loaded
+in-process. Two transports cover the two ringbuffer pairs the dylib exposes:
+
+* iface 0 (main):  rb_main_in  / rb_main_out   — host ↔ firmware protocol
+* iface 1 (debug): rb_debug_in / rb_debug_out  — DebugLink
+
+The vault uses this same FFI surface from Bun. Adding a Python transport that
+mirrors it lets ``python-keepkey`` exercise the firmware contract that the
+dylib path imposes — most importantly, the *caller-driven polling* model:
+nothing happens inside the firmware until the host calls ``kkemu_poll``. UDP
+hides this behind a thread inside ``kkemu``; the dylib does not.
+
+Usage
+-----
+::
+
+    from keepkeylib.transport_dylib import DylibState, DylibTransport
+
+    state = DylibState.get_or_init('/path/to/libkkemu.dylib')
+    main_transport = DylibTransport(state, iface=0)
+    debug_transport = DylibTransport(state, iface=1)
+    client = KeepKeyDebugClient(main_transport)
+    client.set_debuglink(DebugLink(debug_transport))
+
+A *single* ``DylibState`` is shared between the two transports — the dylib's
+``kkemu_init`` may only be called once per process. Re-initialising means
+restarting the test process (or factory-resetting via ``reset_flash``).
+"""
+
+from __future__ import print_function
+
+import ctypes
+import os
+import struct
+import threading
+import time
+
+from .transport import Transport, ConnectionError
+
+
+# ── Dylib singleton ─────────────────────────────────────────────────────────
+
+
+PACKET_SIZE = 64
+FLASH_SIZE = 1 << 20  # 1 MB
+
+# Max time we'll spin in kkemu_poll() looking for a frame on this iface.
+# Has to cover firmware-internal busy-loops (confirm_helper polls usbPoll
+# in a tight C loop — we just need the next outbound frame to land).
+_POLL_TIMEOUT_S = 30.0
+_POLL_QUANTUM_S = 0.001  # 1 ms — keep latency low without burning CPU
+
+
+class DylibState(object):
+    """Process-wide ``libkkemu.dylib`` handle.
+
+    Holds the ctypes binding and the (locked) flash buffer. Only one instance
+    is allowed per process because ``kkemu_init`` is single-shot. Use
+    :func:`get_or_init` rather than the constructor.
+    """
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __init__(self, dylib_path):
+        if not os.path.exists(dylib_path):
+            raise ConnectionError("dylib not found: %s" % dylib_path)
+
+        self.lib = ctypes.CDLL(dylib_path)
+
+        self.lib.kkemu_init.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+        self.lib.kkemu_init.restype = ctypes.c_int
+
+        self.lib.kkemu_shutdown.argtypes = []
+        self.lib.kkemu_shutdown.restype = None
+
+        self.lib.kkemu_poll.argtypes = []
+        self.lib.kkemu_poll.restype = ctypes.c_int
+
+        self.lib.kkemu_is_running.argtypes = []
+        self.lib.kkemu_is_running.restype = ctypes.c_int
+
+        self.lib.kkemu_write.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+        self.lib.kkemu_write.restype = ctypes.c_int
+
+        self.lib.kkemu_read.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+        self.lib.kkemu_read.restype = ctypes.c_int
+
+        self.lib.kkemu_get_display.argtypes = [
+            ctypes.POINTER(ctypes.c_int),
+            ctypes.POINTER(ctypes.c_int),
+        ]
+        self.lib.kkemu_get_display.restype = ctypes.c_void_p
+
+        # Allocate flash as 0xFF (erased NOR state). Held by the singleton so
+        # GC doesn't free it underneath the firmware's still-live mlock.
+        self.flash = (ctypes.c_uint8 * FLASH_SIZE)(*([0xFF] * FLASH_SIZE))
+
+        rc = self.lib.kkemu_init(ctypes.cast(self.flash, ctypes.c_void_p), FLASH_SIZE)
+        if rc != 0:
+            raise ConnectionError("kkemu_init failed: %d" % rc)
+
+        # Single mutex around every FFI call. The dylib's internals aren't
+        # thread-safe; main + debug transport may both poll/read concurrently.
+        self.io_lock = threading.Lock()
+
+        # Pump a few ticks so the firmware finishes its boot sequence (loads
+        # storage, draws home screen) before the first test touches it.
+        with self.io_lock:
+            for _ in range(8):
+                self.lib.kkemu_poll()
+
+    @classmethod
+    def get_or_init(cls, dylib_path):
+        """Return the per-process singleton, creating it on first call.
+
+        Subsequent calls ignore ``dylib_path`` — the dylib is single-shot and
+        re-loading risks UB (mlock'd flash buffer would dangle).
+        """
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = cls(dylib_path)
+            return cls._instance
+
+    def shutdown(self):
+        """Tear down the firmware. Used by tests; not safe to re-init after."""
+        with self.io_lock:
+            self.lib.kkemu_shutdown()
+
+
+# ── Transport ───────────────────────────────────────────────────────────────
+
+
+class DylibTransport(Transport):
+    """One transport per (DylibState, iface) pair.
+
+    ``iface=0`` is the main protocol channel, ``iface=1`` is DebugLink.
+    """
+
+    def __init__(self, state, iface=0, *args, **kwargs):
+        if not isinstance(state, DylibState):
+            raise TypeError("state must be a DylibState")
+        if iface not in (0, 1):
+            raise ValueError("iface must be 0 (main) or 1 (debug)")
+
+        self.state = state
+        self.iface = iface
+        self.read_buffer = b""
+
+        # Transport.__init__ calls self._open(); device arg is just metadata.
+        super(DylibTransport, self).__init__("dylib:iface=%d" % iface, *args, **kwargs)
+
+    # ── Transport hooks ─────────────────────────────────────────────────
+
+    def _open(self):
+        # Nothing to do — the dylib was opened when DylibState was created.
+        pass
+
+    def _close(self):
+        # Don't shut the dylib down on close; the singleton outlives us.
+        self.read_buffer = b""
+
+    def ready_to_read(self):
+        # Drive the firmware once so any pending outbound frame surfaces
+        # in the ringbuffer. Without this, nothing ever appears to be ready.
+        with self.state.io_lock:
+            self.state.lib.kkemu_poll()
+            buf = (ctypes.c_uint8 * PACKET_SIZE)()
+            n = self.state.lib.kkemu_read(buf, PACKET_SIZE, self.iface)
+            if n > 0:
+                # Stash the frame so the next _read sees it without losing data.
+                self.read_buffer += bytes(buf[:n])
+            return bool(self.read_buffer)
+
+    # ── Wire protocol ───────────────────────────────────────────────────
+
+    def _write(self, msg, protobuf_msg):
+        """Chunk ``msg`` into 64-byte HID frames and shove them at the firmware.
+
+        ``msg`` already starts with ``"##"`` + msg-type + length (see
+        ``Transport.write``). The first chunk needs a leading ``"?"`` marker;
+        continuation chunks just get their leading ``"?"`` to round out
+        the 64-byte HID report.
+        """
+        # 63 bytes per chunk + leading '?' = 64 bytes per HID frame
+        for chunk in [msg[i : i + 63] for i in range(0, len(msg), 63)]:
+            chunk = chunk + b"\0" * (63 - len(chunk))
+            frame = b"?" + chunk
+            assert len(frame) == PACKET_SIZE
+            with self.state.io_lock:
+                rc = self.state.lib.kkemu_write(frame, PACKET_SIZE, self.iface)
+                if rc != 0:
+                    raise ConnectionError(
+                        "kkemu_write failed (iface=%d, rc=%d)" % (self.iface, rc)
+                    )
+                # Pump immediately so the firmware can start consuming this
+                # chunk before the next one arrives. Required because the
+                # caller (not a daemon) is the only thing driving the FSM.
+                self.state.lib.kkemu_poll()
+
+    def _read(self):
+        """Read one full message — header parse drives chunk reassembly."""
+        try:
+            (msg_type, datalen) = self._read_headers(_FrameStream(self))
+            payload = self._read_bytes(datalen)
+            return (msg_type, payload)
+        except Exception as exc:
+            print("DylibTransport._read failed: %s" % exc)
+            raise
+
+    # ── Internals ───────────────────────────────────────────────────────
+
+    def _read_bytes(self, length):
+        """Block until ``length`` payload bytes have been gathered."""
+        deadline = time.time() + _POLL_TIMEOUT_S
+        while len(self.read_buffer) < length:
+            if time.time() > deadline:
+                raise ConnectionError(
+                    "Timed out reading %d bytes from iface %d" % (length, self.iface)
+                )
+            self._pump_one()
+        out = self.read_buffer[:length]
+        self.read_buffer = self.read_buffer[length:]
+        return out
+
+    def _pump_one(self):
+        """Run one poll/read cycle. Strips the leading '?' HID marker."""
+        with self.state.io_lock:
+            self.state.lib.kkemu_poll()
+            buf = (ctypes.c_uint8 * PACKET_SIZE)()
+            n = self.state.lib.kkemu_read(buf, PACKET_SIZE, self.iface)
+            if n > 0:
+                # Drop the leading '?' marker; rest is payload.
+                self.read_buffer += bytes(buf[1:n])
+                return
+        # No frame available — back off briefly so we don't spin a hot loop.
+        time.sleep(_POLL_QUANTUM_S)
+
+
+class _FrameStream(object):
+    """File-like adapter so Transport._read_headers can drive _pump_one."""
+
+    def __init__(self, transport):
+        self.transport = transport
+
+    def read(self, n):
+        return self.transport._read_bytes(n)
diff --git a/tests/config.py b/tests/config.py
@@ -73,6 +73,25 @@
     DEBUG_TRANSPORT = WebUsbTransport
     DEBUG_TRANSPORT_ARGS = (webusb_devices[0],)
     DEBUG_TRANSPORT_KWARGS = {'debug_link': True}
+elif os.getenv('KK_TRANSPORT') == 'dylib':
+    # In-process FFI transport against libkkemu.dylib (or libkkemu.so).
+    # Same firmware as UDP, different transport — exposes caller-driven
+    # polling bugs that the UDP daemon hides behind its own poll thread.
+    print('Using Emulator (dylib FFI)')
+    from keepkeylib.transport_dylib import DylibState, DylibTransport
+    _dylib_path = os.getenv('KK_DYLIB')
+    if not _dylib_path:
+        raise RuntimeError(
+            "KK_TRANSPORT=dylib requires KK_DYLIB=/path/to/libkkemu.dylib"
+        )
+    _dylib_state = DylibState.get_or_init(_dylib_path)
+    TRANSPORT = DylibTransport
+    TRANSPORT_ARGS = (_dylib_state, 0)
+    TRANSPORT_KWARGS = {}
+    DEBUG_TRANSPORT = DylibTransport
+    DEBUG_TRANSPORT_ARGS = (_dylib_state, 1)
+    DEBUG_TRANSPORT_KWARGS = {}
+
 else:
     print('Using Emulator')
     TRANSPORT = UDPTransport
diff --git a/tests/test_dylib_confirm_flow.py b/tests/test_dylib_confirm_flow.py
@@ -0,0 +1,74 @@
+"""Regression test for the dylib confirm-flow contract.
+
+Exercises the exact sequence the keepkey-vault FFI path runs:
+
+  1. Initialize       — Features round-trip (no confirm)
+  2. WipeDevice       — needs one confirm (BA on iface 0 + DLD on iface 1)
+  3. LoadDevice       — needs one confirm
+  4. GetAddress       — Features cache + xpub derivation, no confirm
+
+Each step calls into ``confirm_helper`` inside the firmware while the
+caller (this test process) is the only thing driving ``kkemu_poll``. The
+exact same firmware passes the UDP-transport tests because the standalone
+``kkemu`` binary has its own poll thread; the dylib path doesn't, so any
+busy-loop in confirm_helper that waits on a frame the dylib silently
+dropped will hang here.
+
+Skips automatically when ``KK_TRANSPORT != 'dylib'`` so the file is safe
+to keep in the regular pytest run.
+"""
+
+import os
+import unittest
+
+import config
+
+
+@unittest.skipUnless(
+    os.environ.get("KK_TRANSPORT") == "dylib",
+    "dylib confirm-flow regression — set KK_TRANSPORT=dylib KK_DYLIB=...",
+)
+class TestDylibConfirmFlow(unittest.TestCase):
+    """Skipped under the default UDP transport; the UDP daemon hides the
+    polling contract that this test specifically validates."""
+
+    # We import lazily so the module loads even when KK_TRANSPORT != 'dylib'
+    # (config.py only constructs the dylib state on demand in that branch).
+    def setUp(self):
+        # Late import — `common` is heavy (it eagerly wipes the device on
+        # construction) and would defeat the skip above.
+        import common  # noqa: WPS433
+
+        self._common = common
+        self.test = common.KeepKeyTest("setUp")
+        self.test.setUp()
+        self.client = self.test.client
+
+    def tearDown(self):
+        self.test.tearDown()
+
+    def test_features_round_trip(self):
+        """The connection itself works; Features should have firmware fields."""
+        self.client.init_device()
+        f = self.client.features
+        self.assertGreaterEqual(f.major_version, 7)
+
+    def test_load_device_with_auto_confirm(self):
+        """The full LoadDevice flow — confirm_helper must exit cleanly.
+
+        This is the exact path the vault hangs on. If the dylib's tiny-msg
+        dispatch is broken, this test hangs (eventually pytest's timeout
+        kills it) instead of returning.
+        """
+        # KeepKeyTest.setUp already wipes; load a known mnemonic on top.
+        self.test.setup_mnemonic_nopin_nopassphrase()
+        # Round-trip something that requires the seed — confirms LoadDevice
+        # actually committed instead of bouncing off a confirm timeout.
+        addr = self.client.get_address("Bitcoin", [])
+        # Valid mainnet P2PKH addresses start with '1' and are 26-35 chars.
+        self.assertTrue(addr.startswith("1"))
+        self.assertGreaterEqual(len(addr), 26)
+
+
+if __name__ == "__main__":
+    unittest.main()
