feat(python): enable per-workflow THP ACK piggybacking

Following #6563 (comment).

[no changelog]

Author: romanz

python/src/trezorlib/client.py

@@ -24,6 +24,7 @@
 import unicodedata
 import warnings
 from abc import ABCMeta, abstractmethod
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass
 
 import typing_extensions as tx
@@ -245,6 +246,7 @@ def __init__(
         self._mapping = mapping
         self._features = None
         self.pairing = pairing
+        self._interact_ctx = self._interact()
 
     # ===== Internal methods for overriding in subclasses =====
 
@@ -284,11 +286,15 @@ def _get_session(
         """
         raise NotImplementedError
 
+    def _interact(self, *, force_flush: bool = False) -> AbstractContextManager:
+        return nullcontext()
+
     # ===== Common implementations =====
 
     def __enter__(self) -> tx.Self:
         """(Re)Open a connection to the device."""
         self.transport.__enter__()
+        self._interact_ctx.__enter__()
         return self
 
     def __exit__(
@@ -297,7 +303,10 @@ def __exit__(
         exc_value: BaseException | None,
         traceback: t.Any,
     ) -> None:
-        self.transport.__exit__(exc_type, exc_value, traceback)
+        try:
+            self._interact_ctx.__exit__(exc_type, exc_value, traceback)
+        finally:
+            self.transport.__exit__(exc_type, exc_value, traceback)
 
     def connect(self) -> None:
         """Establish a connection to the device.

python/src/trezorlib/thp/channel.py

@@ -22,7 +22,6 @@
 import secrets
 import time
 import typing as t
-from contextlib import contextmanager
 from enum import Enum, IntEnum, auto
 
 import typing_extensions as tx
@@ -111,6 +110,30 @@ def __init__(self) -> None:
         super().__init__(self.__doc__)
 
 
+class _ThpInteractiveContext:
+
+    def __init__(self, channel: Channel, force_flush: bool = False) -> None:
+        self.channel = channel
+        self.force_flush = force_flush
+
+    def __enter__(self) -> tx.Self:
+        self.channel._active_contexts.append(self)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: t.Any,
+    ) -> None:
+        assert self.channel._active_contexts.pop() is self
+        if not self.channel.is_ack_piggybacking_allowed:
+            return
+
+        if not self.channel._active_contexts or self.force_flush:
+            self.channel._flush_ack()
+
+
 class Channel:
     CHUNK_SIZE: t.ClassVar[int | None] = None
 
@@ -141,7 +164,7 @@ def __init__(
         self._noise: NoiseConnection | None = None
         self.state = channel_state
         self.trezor_public_keys: TrezorPublicKeys | None = None
-        self._active_workflow: object | None = None
+        self._active_contexts = []
 
     @functools.cached_property
     def is_ack_piggybacking_allowed(self) -> bool:
@@ -300,8 +323,7 @@ def _read_handshake_init_response(self) -> None:
             if e.code == exceptions.ThpErrorCode.DEVICE_LOCKED:
                 raise DeviceLockedError from e
             raise
-        if not self.is_ack_piggybacking_allowed:
-            self._send_ack(message)
+        self._send_ack(message)
         if not message.is_handshake_init_response():
             raise ProtocolError(f"Not a valid handshake init response: {message}")
 
@@ -409,16 +431,19 @@ def should_back_off() -> bool:
                     continue
                 raise
 
-    def _send_ack(self, acked_message: Message | None) -> None:
-        if self.is_ack_piggybacking_allowed and self._active_workflow is not None:
+    def _send_ack(self, acked_message: Message) -> None:
+        if self.is_ack_piggybacking_allowed and self._active_contexts:
+            # Skip explicit THP ACK during workflow - the next request will piggyback the correct ACK bit
             return
 
-        if acked_message is not None:
-            ack = control_byte.make_ack_for(acked_message.ctrl_byte)
-            ack_message = Message(ack, acked_message.cid, b"")
-        else:
-            ack = control_byte.make_ack(not self.sync_bit_receive)
-            ack_message = Message(ack, self.channel_id, b"")
+        ack = control_byte.make_ack_for(acked_message.ctrl_byte)
+        ack_message = Message(ack, acked_message.cid, b"")
+
+        thp_io.write_payload_to_wire(self.transport, ack_message)
+
+    def _flush_ack(self) -> None:
+        ack = control_byte.make_ack(not self.sync_bit_receive)
+        ack_message = Message(ack, self.channel_id, b"")
 
         thp_io.write_payload_to_wire(self.transport, ack_message)
 
@@ -441,23 +466,6 @@ def _read_ack(self, message: Message) -> None:
             f"Failed to read ACK in {retries} retries for message: {message}"
         )
 
-    @contextmanager
-    def piggyback_acks(self, marker: object) -> t.Generator[None, None, None]:
-        # Make sure the previous workflow is over.
-        assert self._active_workflow is None
-        self._active_workflow = marker
-        # Skip explicit ACKs during this workflow
-        try:
-            yield
-        finally:
-            active = self._active_workflow
-            self._active_workflow = None
-            assert active is marker
-            if self.is_ack_piggybacking_allowed:
-                # Explicitly ACK the latest received message.  The device may restart
-                # the event loop, so the next request will be sent in a separate message.
-                self._send_ack(None)
-
     def write_chunk(self, data: bytes, /) -> None:
         self._assert_handshake_done()
         encrypted_data = self.noise.encrypt(data)

python/src/trezorlib/thp/client.py

@@ -23,7 +23,7 @@
 
 from .. import client, exceptions, messages, models, protobuf
 from ..log import DUMP_BYTES
-from .channel import Channel
+from .channel import Channel, _ThpInteractiveContext
 from .pairing import PairingController
 
 if t.TYPE_CHECKING:
@@ -67,6 +67,7 @@ def __init__(
         mapping: ProtobufMapping | None,
         model: models.TrezorModel | None,
     ) -> None:
+        # used to override channel creation logic in tests
         channel = Channel.allocate(transport)
         try:
             # try to open the channel
@@ -135,6 +136,14 @@ def _get_session(
         session.derive(passphrase, derive_cardano)
         return session
 
+    def _interact(self, *, force_flush: bool = False) -> _ThpInteractiveContext:
+        """
+        Used internally by `TrezorClient` to create an THP ACK piggybacking context.
+
+        Can be also used by tests for unconditionally sending THP ACKs.
+        """
+        return _ThpInteractiveContext(channel=self.channel, force_flush=force_flush)
+
     def _invalidate(self) -> None:
         super()._invalidate()
         # Close the channel. The client cannot be used until a channel is

python/src/trezorlib/thp/pairing.py

@@ -262,6 +262,7 @@ def setup(self) -> None:
             messages.ThpCodeEntryChallenge(challenge=challenge),
             expect=messages.ThpCodeEntryCpaceTrezor,
         )
+        self.controller.channel._flush_ack()
         self.code_entry_state = CodeEntryState(
             challenge=challenge,
             commitment=commitment_msg.commitment,

tests/device_tests/evolu/common.py

@@ -47,13 +47,12 @@ def __init__(self, session, credential):
         self.credential: ThpCredentialResponse = credential
 
 
-def pair_and_get_credential(client: Client) -> ThpPairingResult:
-    from ..thp.connect import nfc_pairing, prepare_channel_for_pairing
+def pair_and_get_credential(test_ctx: Client) -> ThpPairingResult:
+    from ..thp.connect import new_thp_client
 
-    prepare_channel_for_pairing(
-        client, host_static_privkey=TEST_host_static_private_key
+    client = new_thp_client(
+        test_ctx, host_static_privkey=TEST_host_static_private_key, nfc_pairing=True
     )
-    nfc_pairing(client)
     credential = client.pairing.request_credential(autoconnect=False)
     client.pairing.finish()
 
@@ -78,7 +77,7 @@ def get_delegated_identity_key(client: Client) -> bytes:
     if client.is_thp():
         pairing_data = pair_and_get_credential(client)
         return evolu.get_delegated_identity_key(
-            client.get_session(),
+            session=pairing_data.session,
             thp_credential=pairing_data.credential.credential,
         )
     elif client.is_protocol_v1():

tests/device_tests/test_msg_backup_device.py

@@ -60,10 +60,13 @@ def _try_to_cancel(
     next(gen)
     while True:
         br = yield
-        # Try to cancel the backup flow on Core
-        resp = session.call_raw(messages.Cancel())
-        # Following #6483, backup is not cancellable
-        assert resp == BACKUP_IN_PROGRESS
+        # Entering session's context will send an explicit THP ACK after `BACKUP_IN_PROGRESS` is received.
+        with session.client._interact(force_flush=True):
+            # Try to cancel the backup flow on Core
+            with pytest.raises(TrezorFailure) as exc_info:
+                session.call(messages.Cancel(), expect=messages.Failure)
+            # Following #6483, backup is not cancellable
+            assert exc_info.value.failure == BACKUP_IN_PROGRESS
         try:
             gen.send(br)
         except StopIteration:

tests/device_tests/thp/connect.py

@@ -39,63 +39,62 @@ def prepare_channel_for_handshake(test_ctx: TrezorTestContext) -> None:
     test_ctx.channel._init_noise()
 
 
-def prepare_channel_for_pairing(
+def new_thp_client(
     test_ctx: TrezorTestContext,
+    *,
     credential: Credential | None = None,
     host_static_privkey: bytes | None = None,
     fixed_entropy: bool = False,
-) -> None:
-    """Create a fresh channel, perform the handshake using the provided fixed entropy
-    and credentials, and leave it in the pairing phase.
-    """
+    nfc_pairing: bool = False,
+) -> TrezorClientThp:
     # set up a fresh channel
-    prepare_channel_for_handshake(test_ctx)
+    assert isinstance(test_ctx.client, TrezorClientThp)
+    test_ctx.channel = Channel.allocate(test_ctx.transport)
+
     assert host_static_privkey is None or not fixed_entropy, "don't use both"
     if host_static_privkey is not None:
         test_ctx.channel._init_noise(static_privkey=host_static_privkey)
     elif fixed_entropy:
         test_ctx.channel._init_noise(
             static_privkey=b"\x12" * 32, ephemeral_privkey=b"\x24" * 32
         )
+    else:
+        test_ctx.channel._init_noise()
+
     credentials = []
     if credential is not None:
         credentials.append(credential)
 
     # run the handshake
     test_ctx.channel.open(credentials)
-    assert isinstance(test_ctx.client, TrezorClientThp)
+
+    test_ctx.client._interact_ctx = test_ctx.client._interact()
     test_ctx.client.pairing = test_ctx.pairing = PairingController(test_ctx.client)
 
+    if nfc_pairing:
+        method = Nfc(test_ctx.pairing)
+        # NFC screen shown
+
+        # Read `nfc_secret` and `handshake_hash` from Trezor using debuglink
+        pairing_info = test_ctx.debug.pairing_info(
+            thp_channel_id=test_ctx.channel.channel_id.to_bytes(2, "big"),
+            handshake_hash=test_ctx.channel.handshake_hash,
+            nfc_secret_host=method.nfc_host_secret,
+        )
+        assert pairing_info.handshake_hash is not None
+        assert pairing_info.nfc_secret_trezor is not None
+        assert pairing_info.handshake_hash[:16] == test_ctx.channel.handshake_hash[:16]
 
-def get_encrypted_transport_protocol(test_ctx: TrezorTestContext) -> None:
-    prepare_channel_for_pairing(test_ctx)
-    test_ctx.pairing.skip()
+        method.send_nfc_tag(pairing_info.nfc_secret_trezor)
 
+    return test_ctx.client
 
-def break_channel(test_ctx: TrezorTestContext) -> None:
-    cse = test_ctx.channel._noise.noise_protocol.cipher_state_encrypt
+
+def break_channel(client: TrezorClientThp) -> None:
+    cse = client.channel._noise.noise_protocol.cipher_state_encrypt
     cse.n = cse.n + 1
 
-    session = test_ctx.client._get_any_session()
+    session = client._get_any_session()
     session.write(messages.ButtonAck())
     with pytest.raises(ThpError):
-        session.read(1)
-
-
-def nfc_pairing(test_ctx: TrezorTestContext) -> None:
-    assert isinstance(test_ctx.client, TrezorClientThp)
-    method = Nfc(test_ctx.client.pairing)
-
-    # NFC screen shown
-
-    # Read `nfc_secret` and `handshake_hash` from Trezor using debuglink
-    pairing_info = test_ctx.debug.pairing_info(
-        thp_channel_id=test_ctx.channel.channel_id.to_bytes(2, "big"),
-        handshake_hash=test_ctx.channel.handshake_hash,
-        nfc_secret_host=method.nfc_host_secret,
-    )
-    assert pairing_info.handshake_hash is not None
-    assert pairing_info.nfc_secret_trezor is not None
-    assert pairing_info.handshake_hash[:16] == test_ctx.channel.handshake_hash[:16]
-
-    method.send_nfc_tag(pairing_info.nfc_secret_trezor)
+        session.read(timeout=1)