From e8fa6bda6b7046668827f980824dc5ba3ee82ea4 Mon Sep 17 00:00:00 2001 From: Martin Milata Date: Fri, 21 Feb 2025 00:43:26 +0100 Subject: [PATCH] fix(python): transport handling with sessions [no changelog] --- python/src/trezorlib/_internal/emulator.py | 27 ++++--- python/src/trezorlib/cli/__init__.py | 16 +++- python/src/trezorlib/cli/debug.py | 3 +- python/src/trezorlib/cli/trezorctl.py | 3 + python/src/trezorlib/client.py | 13 ++- python/src/trezorlib/debuglink.py | 79 ++++++++----------- python/src/trezorlib/device.py | 2 - python/src/trezorlib/transport/hid.py | 3 + .../trezorlib/transport/thp/protocol_v1.py | 1 - python/src/trezorlib/transport/udp.py | 4 - python/src/trezorlib/transport/webusb.py | 5 +- tests/click_tests/test_recovery.py | 7 +- tests/click_tests/test_repeated_backup.py | 2 +- tests/device_handler.py | 7 +- .../bitcoin/test_authorize_coinjoin.py | 7 +- tests/upgrade_tests/test_firmware_upgrades.py | 1 + 16 files changed, 93 insertions(+), 87 deletions(-) diff --git a/python/src/trezorlib/_internal/emulator.py b/python/src/trezorlib/_internal/emulator.py index 8772770b40..f3a59f8b25 100644 --- a/python/src/trezorlib/_internal/emulator.py +++ b/python/src/trezorlib/_internal/emulator.py @@ -23,6 +23,7 @@ from pathlib import Path from typing import Any, Dict, Iterable, List, Optional, Sequence, TextIO, Union, cast from ..debuglink import TrezorClientDebugLink +from ..transport import Transport from ..transport.udp import UdpTransport LOG = logging.getLogger(__name__) @@ -118,13 +119,12 @@ class Emulator: def wait_until_ready(self, timeout: float = EMULATOR_WAIT_TIME) -> None: assert self.process is not None, "Emulator not started" - transport = self._get_transport() - transport.open() + self.transport.open() LOG.info("Waiting for emulator to come up...") start = time.monotonic() try: while True: - if transport.ping(): + if self.transport.ping(): break if self.process.poll() is not None: raise RuntimeError("Emulator process died") @@ -135,7 +135,7 @@ class Emulator: time.sleep(0.1) finally: - transport.close() + self.transport.close() LOG.info(f"Emulator ready after {time.monotonic() - start:.3f} seconds") @@ -166,7 +166,11 @@ class Emulator: env=env, ) - def start(self) -> None: + def start( + self, + transport: Optional[UdpTransport] = None, + debug_transport: Optional[Transport] = None, + ) -> None: if self.process: if self.process.poll() is not None: # process has died, stop and start again @@ -176,6 +180,7 @@ class Emulator: # process is running, no need to start again return + self.transport = transport or self._get_transport() self.process = self.launch_process() _RUNNING_PIDS.add(self.process) try: @@ -189,15 +194,16 @@ class Emulator: (self.profile_dir / "trezor.pid").write_text(str(self.process.pid) + "\n") (self.profile_dir / "trezor.port").write_text(str(self.port) + "\n") - transport = self._get_transport() self._client = TrezorClientDebugLink( - transport, auto_interact=self.auto_interact + self.transport, + auto_interact=self.auto_interact, + open_transport=True, + debug_transport=debug_transport, ) - self._client.open() def stop(self) -> None: if self._client: - self._client.close() + self._client.close_transport() self._client = None if self.process: @@ -221,8 +227,9 @@ class Emulator: # preserving the recording directory between restarts self.restart_amount += 1 prev_screenshot_dir = self.client.debug.screenshot_recording_dir + debug_transport = self.client.debug.transport self.stop() - self.start() + self.start(transport=self.transport, debug_transport=debug_transport) if prev_screenshot_dir: self.client.debug.start_recording( prev_screenshot_dir, refresh_index=self.restart_amount diff --git a/python/src/trezorlib/cli/__init__.py b/python/src/trezorlib/cli/__init__.py index f8a45d3612..74b6c11d91 100644 --- a/python/src/trezorlib/cli/__init__.py +++ b/python/src/trezorlib/cli/__init__.py @@ -16,6 +16,7 @@ from __future__ import annotations +import atexit import functools import logging import os @@ -33,6 +34,8 @@ from ..transport.session import Session, SessionV1 LOG = logging.getLogger(__name__) +_TRANSPORT: Transport | None = None + if t.TYPE_CHECKING: # Needed to enforce a return value from decorators # More details: https://www.python.org/dev/peps/pep-0612/ @@ -167,16 +170,25 @@ class TrezorConnection: return session def get_transport(self) -> "Transport": + global _TRANSPORT + if _TRANSPORT is not None: + return _TRANSPORT + try: # look for transport without prefix search - return transport.get_transport(self.path, prefix_search=False) + _TRANSPORT = transport.get_transport(self.path, prefix_search=False) except Exception: # most likely not found. try again below. pass # look for transport with prefix search # if this fails, we want the exception to bubble up to the caller - return transport.get_transport(self.path, prefix_search=True) + if not _TRANSPORT: + _TRANSPORT = transport.get_transport(self.path, prefix_search=True) + + _TRANSPORT.open() + atexit.register(_TRANSPORT.close) + return _TRANSPORT def get_client(self) -> TrezorClient: client = get_client(self.get_transport()) diff --git a/python/src/trezorlib/cli/debug.py b/python/src/trezorlib/cli/debug.py index 00f0c6276b..c4afae6b02 100644 --- a/python/src/trezorlib/cli/debug.py +++ b/python/src/trezorlib/cli/debug.py @@ -52,9 +52,8 @@ def record_screen_from_connection( """Record screen helper to transform TrezorConnection into TrezorClientDebugLink.""" transport = obj.get_transport() debug_client = TrezorClientDebugLink(transport, auto_interact=False) - debug_client.open() record_screen(debug_client, directory, report_func=click.echo) - debug_client.close() + debug_client.close_transport() @cli.command() diff --git a/python/src/trezorlib/cli/trezorctl.py b/python/src/trezorlib/cli/trezorctl.py index 028b01474a..d3f3be940f 100755 --- a/python/src/trezorlib/cli/trezorctl.py +++ b/python/src/trezorlib/cli/trezorctl.py @@ -296,11 +296,14 @@ def list_devices(no_resolve: bool) -> Optional[Iterable["Transport"]]: for transport in enumerate_devices(): try: client = get_client(transport) + transport.open() description = format_device_name(client.features) except DeviceIsBusy: description = "Device is in use by another process" except Exception as e: description = "Failed to read details " + str(type(e)) + finally: + transport.close() click.echo(f"{transport.get_path()} - {description}") return None diff --git a/python/src/trezorlib/client.py b/python/src/trezorlib/client.py index 959be1425c..e4707ff1aa 100644 --- a/python/src/trezorlib/client.py +++ b/python/src/trezorlib/client.py @@ -165,19 +165,13 @@ class TrezorClient: def is_invalidated(self) -> bool: return self._is_invalidated - def refresh_features(self) -> None: + def refresh_features(self) -> messages.Features: self.protocol.update_features() self._features = self.protocol.get_features() + return self._features def _get_protocol(self) -> Channel: - self.transport.open() - protocol = ProtocolV1Channel(self.transport, mapping.DEFAULT_MAPPING) - - protocol.write(messages.Initialize()) - - _ = protocol.read() - self.transport.close() return protocol @@ -189,6 +183,8 @@ def get_default_client( Returns a TrezorClient instance with minimum fuss. + Transport is opened and should be closed after you're done with the client. + If path is specified, does a prefix-search for the specified device. Otherwise, uses the value of TREZOR_PATH env variable, or finds first connected Trezor. If no UI is supplied, instantiates the default CLI UI. @@ -198,5 +194,6 @@ def get_default_client( path = os.getenv("TREZOR_PATH") transport = get_transport(path, prefix_search=True) + transport.open() return TrezorClient(transport, **kwargs) diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index c1ed378d0f..84d199e849 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -476,15 +476,9 @@ class DebugLink: def open(self) -> None: self.transport.open() - # raise NotImplementedError - # TODO is this needed? - # self.transport.deprecated_begin_session() def close(self) -> None: - pass - # raise NotImplementedError - # TODO is this needed? - # self.transport.deprecated_end_session() + self.transport.close() def _write(self, msg: protobuf.MessageType) -> None: if self.waiting_for_layout_change: @@ -1186,26 +1180,35 @@ class TrezorClientDebugLink(TrezorClient): # without special DebugLink interface provided # by the device. - def __init__(self, transport: "Transport", auto_interact: bool = True) -> None: + def __init__( + self, + transport: Transport, + auto_interact: bool = True, + open_transport: bool = True, + debug_transport: Transport | None = None, + ) -> None: try: - debug_transport = transport.find_debug() + debug_transport = debug_transport or transport.find_debug() self.debug = DebugLink(debug_transport, auto_interact) + if open_transport: + self.debug.open() # try to open debuglink, see if it works - self.debug.open() - self.debug.close() + assert self.debug.transport.ping() except Exception: if not auto_interact: self.debug = NullDebugLink() else: raise + if open_transport: + transport.open() + # set transport explicitly so that sync_responses can work super().__init__(transport) self.transport = transport self.ui: DebugUI = DebugUI(self.debug) - self.reset_debug_features(new_seedless_session=True) self.sync_responses() # So that we can choose right screenshotting logic (T1 vs TT) @@ -1219,14 +1222,17 @@ class TrezorClientDebugLink(TrezorClient): def get_new_client(self) -> TrezorClientDebugLink: new_client = TrezorClientDebugLink( - self.transport, self.debug.allow_interactions + self.transport, + self.debug.allow_interactions, + open_transport=False, + debug_transport=self.debug.transport, ) new_client.debug.screenshot_recording_dir = self.debug.screenshot_recording_dir new_client.debug.t1_screenshot_directory = self.debug.t1_screenshot_directory new_client.debug.t1_screenshot_counter = self.debug.t1_screenshot_counter return new_client - def reset_debug_features(self, new_seedless_session: bool = False) -> None: + def reset_debug_features(self) -> None: """ Prepare the debugging client for a new testcase. @@ -1332,21 +1338,9 @@ class TrezorClientDebugLink(TrezorClient): return _callback_passphrase - def ensure_open(self) -> None: - """Only open session if there isn't already an open one.""" - # if self.session_counter == 0: - # self.open() - # TODO check if is this needed - - def open(self) -> None: - pass - # TODO is this needed? - # self.debug.open() - - def close(self) -> None: - pass - # TODO is this needed? - # self.debug.close() + def close_transport(self) -> None: + self.transport.close() + self.debug.close() def lock(self) -> None: s = self.get_seedless_session() @@ -1356,7 +1350,7 @@ class TrezorClientDebugLink(TrezorClient): self, passphrase: str | object | None = "", derive_cardano: bool = False, - session_id: int = 0, + session_id: bytes | None = None, ) -> SessionDebugWrapper: if isinstance(passphrase, str): passphrase = Mnemonic.normalize_string(passphrase) @@ -1445,7 +1439,7 @@ class TrezorClientDebugLink(TrezorClient): else: input_flow = None - self.reset_debug_features(new_seedless_session=False) + self.reset_debug_features() if exc_type is not None and isinstance(input_flow, t.Generator): # Propagate the exception through the input flow, so that we see in @@ -1498,20 +1492,15 @@ class TrezorClientDebugLink(TrezorClient): # prompt, which is in TINY mode and does not respond to `Ping`. if self.protocol_version is ProtocolVersion.V1: assert isinstance(self.protocol, ProtocolV1Channel) - self.transport.open() - try: - self.protocol.write(messages.Cancel()) - resp = self.protocol.read() - message = "SYNC" + secrets.token_hex(8) - self.protocol.write(messages.Ping(message=message)) - while resp != messages.Success(message=message): - try: - resp = self.protocol.read() - except Exception: - pass - finally: - pass - # TODO fix self.transport.end_session() + self.protocol.write(messages.Cancel()) + resp = self.protocol.read() + message = "SYNC" + secrets.token_hex(8) + self.protocol.write(messages.Ping(message=message)) + while resp != messages.Success(message=message): + try: + resp = self.protocol.read() + except Exception: + pass def mnemonic_callback(self, _) -> str: word, pos = self.debug.read_recovery_word() diff --git a/python/src/trezorlib/device.py b/python/src/trezorlib/device.py index a3b24c247d..42f752d4e1 100644 --- a/python/src/trezorlib/device.py +++ b/python/src/trezorlib/device.py @@ -138,8 +138,6 @@ def sd_protect( def wipe(session: "Session") -> str | None: ret = session.call(messages.WipeDevice(), expect=messages.Success) session.invalidate() - # if not session.features.bootloader_mode: - # session.refresh_features() return _return_success(ret) diff --git a/python/src/trezorlib/transport/hid.py b/python/src/trezorlib/transport/hid.py index d37dbcf606..47f1f2558b 100644 --- a/python/src/trezorlib/transport/hid.py +++ b/python/src/trezorlib/transport/hid.py @@ -156,6 +156,9 @@ class HidTransport(Transport): return 1 raise TransportException("Unknown HID version") + def ping(self) -> bool: + return self.handle is not None + def is_wirelink(dev: HidDevice) -> bool: return dev["usage_page"] == 0xFF00 or dev["interface_number"] == 0 diff --git a/python/src/trezorlib/transport/thp/protocol_v1.py b/python/src/trezorlib/transport/thp/protocol_v1.py index 0dd6bdd7e8..6837dc04f0 100644 --- a/python/src/trezorlib/transport/thp/protocol_v1.py +++ b/python/src/trezorlib/transport/thp/protocol_v1.py @@ -60,7 +60,6 @@ class ProtocolV1Channel(Channel): f"received message: {msg.__class__.__name__}", extra={"protobuf": msg}, ) - self.transport.close() return msg def write(self, msg: t.Any) -> None: diff --git a/python/src/trezorlib/transport/udp.py b/python/src/trezorlib/transport/udp.py index c040545d7e..b276f5900b 100644 --- a/python/src/trezorlib/transport/udp.py +++ b/python/src/trezorlib/transport/udp.py @@ -106,8 +106,6 @@ class UdpTransport(Transport): self.socket = None def write_chunk(self, chunk: bytes) -> None: - if self.socket is None: - self.open() assert self.socket is not None if len(chunk) != 64: raise TransportException("Unexpected data length") @@ -115,8 +113,6 @@ class UdpTransport(Transport): self.socket.sendall(chunk) def read_chunk(self, timeout: float | None = None) -> bytes: - if self.socket is None: - self.open() assert self.socket is not None start = time.time() while True: diff --git a/python/src/trezorlib/transport/webusb.py b/python/src/trezorlib/transport/webusb.py index 7919608825..78a4292ca1 100644 --- a/python/src/trezorlib/transport/webusb.py +++ b/python/src/trezorlib/transport/webusb.py @@ -134,8 +134,6 @@ class WebUsbTransport(Transport): self.handle = None def write_chunk(self, chunk: bytes) -> None: - if self.handle is None: - self.open() assert self.handle is not None if len(chunk) != WEBUSB_CHUNK_SIZE: raise TransportException(f"Unexpected chunk size: {len(chunk)}") @@ -180,6 +178,9 @@ class WebUsbTransport(Transport): # For v1 protocol, find debug USB interface for the same serial number return self.__class__(self.device, debug=True) + def ping(self) -> bool: + return self.handle is not None + def is_vendor_class(dev: usb1.USBDevice) -> bool: configurationId = 0 diff --git a/tests/click_tests/test_recovery.py b/tests/click_tests/test_recovery.py index deb2c851e8..8ba72b8eda 100644 --- a/tests/click_tests/test_recovery.py +++ b/tests/click_tests/test_recovery.py @@ -127,6 +127,7 @@ def test_recovery_cancel_issue4613(device_handler: "BackgroundDeviceHandler"): recovery.confirm_recovery(debug, title="recovery__title_dry_run") # select number of words recovery.select_number_of_words(debug, num_of_words=12) + device_handler.client.transport.close() # abort the process running the recovery from host device_handler.kill_task() @@ -134,7 +135,7 @@ def test_recovery_cancel_issue4613(device_handler: "BackgroundDeviceHandler"): # from the host side. # Reopen client and debuglink, closed by kill_task - device_handler.client.open() + device_handler.client.transport.open() debug = device_handler.debuglink() # Ping the Trezor with an Initialize message (listed in DO_NOT_RESTART) @@ -145,7 +146,9 @@ def test_recovery_cancel_issue4613(device_handler: "BackgroundDeviceHandler"): except exceptions.Cancelled: # due to a related problem, the first call in this situation will return # a Cancelled failure. This test does not care, we just retry. - features = device_handler.client.call(messages.Initialize()) + features = device_handler.client.get_seedless_session().call( + messages.Initialize() + ) assert features.recovery_status == messages.RecoveryStatus.Recovery # Trezor is sitting in recovery_homescreen now, waiting for the user to select diff --git a/tests/click_tests/test_repeated_backup.py b/tests/click_tests/test_repeated_backup.py index 7ff1973e3a..c08ade16f4 100644 --- a/tests/click_tests/test_repeated_backup.py +++ b/tests/click_tests/test_repeated_backup.py @@ -202,7 +202,7 @@ def test_repeated_backup( assert features.recovery_status == messages.RecoveryStatus.Nothing # try to unlock backup yet again... - device_handler.run( + device_handler.run_with_session( device.recover, type=messages.RecoveryType.UnlockRepeatedBackup, ) diff --git a/tests/device_handler.py b/tests/device_handler.py index b07ace407b..0f69bff4d2 100644 --- a/tests/device_handler.py +++ b/tests/device_handler.py @@ -11,6 +11,7 @@ from trezorlib.transport import udp if t.TYPE_CHECKING: from trezorlib._internal.emulator import Emulator from trezorlib.debuglink import DebugLink + from trezorlib.debuglink import SessionDebugWrapper as Session from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.messages import Features @@ -53,7 +54,7 @@ class BackgroundDeviceHandler: def run_with_session( self, - function: t.Callable[tx.Concatenate["Client", P], t.Any], + function: t.Callable[tx.Concatenate["Session", P], t.Any], *args: P.args, **kwargs: P.kwargs, ) -> None: @@ -72,7 +73,7 @@ class BackgroundDeviceHandler: def run_with_provided_session( self, session, - function: t.Callable[tx.Concatenate["Client", P], t.Any], + function: t.Callable[tx.Concatenate["Session", P], t.Any], *args: P.args, **kwargs: P.kwargs, ) -> None: @@ -92,8 +93,6 @@ class BackgroundDeviceHandler: # Force close the client, which should raise an exception in a client # waiting on IO. Does not work over Bridge, because bridge doesn't have # a close() method. - # while self.client.session_counter > 0: - # self.client.close() try: self.task.result(timeout=1) except Exception: diff --git a/tests/device_tests/bitcoin/test_authorize_coinjoin.py b/tests/device_tests/bitcoin/test_authorize_coinjoin.py index 549e275358..65157487f4 100644 --- a/tests/device_tests/bitcoin/test_authorize_coinjoin.py +++ b/tests/device_tests/bitcoin/test_authorize_coinjoin.py @@ -793,7 +793,7 @@ def test_get_address(session: Session): def test_multisession_authorization(client: Client): # Authorize CoinJoin with www.example1.com in session 1. - session1 = client.get_session(session_id=1) + session1 = client.get_session() btc.authorize_coinjoin( session1, @@ -805,10 +805,9 @@ def test_multisession_authorization(client: Client): coin_name="Testnet", script_type=messages.InputScriptType.SPENDTAPROOT, ) - session2 = client.get_session(session_id=2) + # Open a second session. - # session_id1 = session.session_id - # TODO client.init_device(new_session=True) + session2 = client.get_session() # Authorize CoinJoin with www.example2.com in session 2. btc.authorize_coinjoin( diff --git a/tests/upgrade_tests/test_firmware_upgrades.py b/tests/upgrade_tests/test_firmware_upgrades.py index f4b2b31ed1..bb35e852fb 100644 --- a/tests/upgrade_tests/test_firmware_upgrades.py +++ b/tests/upgrade_tests/test_firmware_upgrades.py @@ -479,6 +479,7 @@ def test_upgrade_u2f(gen: str, tag: str): storage = emu.get_storage() with EmulatorWrapper(gen, storage=storage) as emu: + session = emu.client.get_seedless_session() counter = fido.get_next_counter(session) assert counter == 12