From 54c67997dacea9b1e4b9c9e9075e2fa15af07d97 Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Tue, 15 Apr 2025 16:39:19 +0200
Subject: [PATCH] fixup! chore: handle cancel in pairing and credential flows

---
 core/src/apps/thp/pairing.py           | 13 ++++++++++++-
 tests/device_tests/thp/test_pairing.py | 25 ++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/core/src/apps/thp/pairing.py b/core/src/apps/thp/pairing.py
index 790d3d1419..72a872064c 100644
--- a/core/src/apps/thp/pairing.py
+++ b/core/src/apps/thp/pairing.py
@@ -39,6 +39,8 @@ if __debug__:
 if TYPE_CHECKING:
     from typing import Any, Callable, Concatenate, ParamSpec, Tuple
 
+    from trezorui_api import UiResult
+
     P = ParamSpec("P")
     FuncWithContext = Callable[Concatenate[PairingContext, P], Any]
 
@@ -91,6 +93,7 @@ async def handle_pairing_request(
 ) -> ThpEndResponse:
 
     if Cancel.is_type_of(message):
+        ctx.channel_ctx.set_channel_state(ChannelState.TP0)
         raise ActionCancelled()
 
     if not ThpPairingRequest.is_type_of(message):
@@ -123,14 +126,22 @@ async def handle_pairing_request(
         await _prepare_pairing(ctx)
 
         ctx.channel_ctx.set_channel_state(ChannelState.TP3)
+        result: UiResult | None = None
         try:
             # Should raise UnexpectedMessageException
-            await ctx.show_pairing_method_screen()
+            result = await ctx.show_pairing_method_screen()
         except UnexpectedMessageException as e:
             raw_response = e.msg
             req_type = protobuf.type_for_wire(raw_response.type)
             response = message_handler.wrap_protobuf_load(raw_response.data, req_type)
 
+        if result is not None:
+            from trezorui_api import CONFIRMED
+
+            if result is CONFIRMED:
+                # TODO In UI Screen has one button with text "Cancel", but is gets confirmed
+                ctx.channel_ctx.set_channel_state(ChannelState.TP0)
+                raise ActionCancelled
         if Cancel.is_type_of(response):
             ctx.channel_ctx.clear()
             raise SilentError("Action was cancelled by the Host")
diff --git a/tests/device_tests/thp/test_pairing.py b/tests/device_tests/thp/test_pairing.py
index f00f59f95a..7d367000a1 100644
--- a/tests/device_tests/thp/test_pairing.py
+++ b/tests/device_tests/thp/test_pairing.py
@@ -150,6 +150,29 @@ def test_pairing_code_entry(client: Client) -> None:
     protocol._has_valid_channel = True
 
 
+def test_pairing_code_entry_cancel(client: Client) -> None:
+    protocol = prepare_protocol_for_pairing(client)
+
+    handle_pairing_request(client, protocol, "TestTrezor CodeEntry")
+
+    protocol._send_message(
+        ThpSelectMethod(selected_pairing_method=ThpPairingMethod.CodeEntry)
+    )
+
+    _ = protocol._read_message(ThpCodeEntryCommitment)
+
+    challenge = os.urandom(16)
+    protocol._send_message(ThpCodeEntryChallenge(challenge=challenge))
+    protocol._read_message(ThpCodeEntryCpaceTrezor)
+
+    # Code Entry code shown
+
+    # Press Cancel button
+    client.debug.press_yes()
+    failure = protocol._read_message(Failure)
+    assert failure.code is FailureType.ActionCancelled
+
+
 def test_pairing_cancel_1(client: Client) -> None:
     protocol = prepare_protocol_for_pairing(client)
 
@@ -360,7 +383,7 @@ def test_credential_phase(client: Client) -> None:
     client.debug.press_yes()
     # Autoconnect issuance confirmation dialog is shown.
     button_req = protocol._read_message(ButtonRequest)
-    assert button_req.name == "autoconnect_credential_request"
+    assert button_req.name == "thp_autoconnect_credential_request"
     protocol._send_message(ButtonAck())
     client.debug.press_yes()
     # Autoconnect credential is received