diff --git a/core/src/trezor/wire/thp/__init__.py b/core/src/trezor/wire/thp/__init__.py index cbdddb67fd..fc5e124f9d 100644 --- a/core/src/trezor/wire/thp/__init__.py +++ b/core/src/trezor/wire/thp/__init__.py @@ -48,6 +48,10 @@ class ThpDeviceLockedError(ThpError): pass +class ThpUnallocatedChannelError(ThpError): + pass + + class ThpUnallocatedSessionError(ThpError): def __init__(self, session_id: int) -> None: diff --git a/core/src/trezor/wire/thp/channel.py b/core/src/trezor/wire/thp/channel.py index 4490ddfc59..0e76bb3f27 100644 --- a/core/src/trezor/wire/thp/channel.py +++ b/core/src/trezor/wire/thp/channel.py @@ -214,7 +214,12 @@ class Channel: self.fallback_decrypt = False self.expected_payload_length = 0 self.bytes_read = 0 - print("FAILED TO FALLBACK") + if __debug__ and utils.ALLOW_DEBUG_MESSAGES: + from ubinascii import hexlify + + log.debug( + __name__, "FAILED TO FALLBACK: %s", hexlify(packet).decode() + ) return to_read_len = min(len(packet) - INIT_HEADER_LENGTH, payload_length) diff --git a/core/src/trezor/wire/thp/received_message_handler.py b/core/src/trezor/wire/thp/received_message_handler.py index 022db83701..275e4c3912 100644 --- a/core/src/trezor/wire/thp/received_message_handler.py +++ b/core/src/trezor/wire/thp/received_message_handler.py @@ -35,6 +35,7 @@ from . import ( ThpError, ThpErrorType, ThpInvalidDataError, + ThpUnallocatedChannelError, ThpUnallocatedSessionError, ) from . import alternating_bit_protocol as ABP @@ -129,6 +130,9 @@ async def handle_received_message( except ThpUnallocatedSessionError as e: error_message = Failure(code=FailureType.ThpUnallocatedSession) await ctx.write(error_message, e.session_id) + except ThpUnallocatedChannelError: + await ctx.write_error(ThpErrorType.UNALLOCATED_CHANNEL) + ctx.clear() except ThpDecryptionError: await ctx.write_error(ThpErrorType.DECRYPTION_FAILED) ctx.clear() @@ -271,8 +275,11 @@ async def _handle_state_TH2(ctx: Channel, message_length: int, ctrl_byte: int) - log.debug(__name__, "handle_state_TH2") if not control_byte.is_handshake_comp_req(ctrl_byte): raise ThpError("Message received is not a handshake completion request!") + if ctx.handshake is None: - raise Exception("Handshake object is not prepared. Retry handshake.") + raise ThpUnallocatedChannelError( + "Handshake object is not prepared. Create new channel." + ) if not config.is_unlocked(): raise ThpDeviceLockedError