From 41cdbf7bc4e4e16219851edae4a22c1d5f133b61 Mon Sep 17 00:00:00 2001 From: M1nd3r Date: Wed, 4 Dec 2024 17:44:35 +0100 Subject: [PATCH] refactor(core): move handling of ThpCreateNewSession message into base.py [no changelog] --- core/src/all_modules.py | 2 - core/src/apps/base.py | 74 ++++++++++++++++++++++++- core/src/apps/thp/create_new_session.py | 59 -------------------- core/src/apps/workflow_handlers.py | 7 --- 4 files changed, 71 insertions(+), 71 deletions(-) delete mode 100644 core/src/apps/thp/create_new_session.py diff --git a/core/src/all_modules.py b/core/src/all_modules.py index 183755390c..9e90c4060b 100644 --- a/core/src/all_modules.py +++ b/core/src/all_modules.py @@ -461,8 +461,6 @@ if utils.USE_THP: import trezor.wire.thp.writer apps.thp import apps.thp - apps.thp.create_new_session - import apps.thp.create_new_session apps.thp.credential_manager import apps.thp.credential_manager apps.thp.pairing diff --git a/core/src/apps/base.py b/core/src/apps/base.py index ca923b7337..1cc546a959 100644 --- a/core/src/apps/base.py +++ b/core/src/apps/base.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING import storage.device as storage_device from storage.cache_common import APP_COMMON_BUSY_DEADLINE_MS, APP_COMMON_SEED from trezor import TR, config, utils, wire, workflow -from trezor.enums import HomescreenFormat, MessageType +from trezor.enums import HomescreenFormat, MessageType, ThpMessageType from trezor.messages import Success, UnlockPath from trezor.ui.layouts import confirm_action from trezor.wire import context @@ -27,6 +27,9 @@ if TYPE_CHECKING: ) from trezor.wire import Handler, Msg + if utils.USE_THP: + from trezor.messages import Failure, ThpCreateNewSession, ThpNewSession + _SCREENSAVER_IS_ON = False @@ -204,7 +207,69 @@ def get_features() -> Features: return f -if not utils.USE_THP: +if utils.USE_THP: + + async def handle_ThpCreateNewSession( + message: ThpCreateNewSession, + ) -> ThpNewSession | Failure: + """ + Creates a new `ThpSession` based on the provided parameters and returns a + `ThpNewSession` message containing the new session ID. + + Returns an appropriate `Failure` message if session creation fails. + """ + from trezor import log, loop + from trezor.enums import FailureType + from trezor.messages import Failure, ThpNewSession + from trezor.wire import NotInitialized + from trezor.wire.context import get_context + from trezor.wire.errors import ActionCancelled, DataError + from trezor.wire.thp import SessionState + from trezor.wire.thp.session_context import GenericSessionContext + from trezor.wire.thp.session_manager import create_new_session + + from apps.common.seed import derive_and_store_roots + + ctx = get_context() + + # Assert that context `ctx` is `GenericSessionContext` + assert isinstance(ctx, GenericSessionContext) + + channel = ctx.channel + + # Do not use `ctx` beyond this point, as it is techically + # allowed to change in between await statements + + new_session = create_new_session(channel) + try: + await derive_and_store_roots(new_session, message) + except DataError as e: + return Failure(code=FailureType.DataError, message=e.message) + except ActionCancelled as e: + return Failure(code=FailureType.ActionCancelled, message=e.message) + except NotInitialized as e: + return Failure(code=FailureType.NotInitialized, message=e.message) + # TODO handle other errors (`Exception`` when "Cardano icarus secret is already set!" + # and `RuntimeError` when accessing storage for mnemonic.get_secret - it actually + # happens for locked devices) + + new_session.set_session_state(SessionState.ALLOCATED) + channel.sessions[new_session.session_id] = new_session + loop.schedule(new_session.handle()) + new_session_id: int = new_session.session_id + + if __debug__ and utils.ALLOW_DEBUG_MESSAGES: + log.debug( + __name__, + "create_new_session - new session created. Passphrase: %s, Session id: %d\n%s", + message.passphrase if message.passphrase is not None else "", + new_session.session_id, + str(channel.sessions), + ) + + return ThpNewSession(new_session_id=new_session_id) + +else: async def handle_Initialize(msg: Initialize) -> Features: import storage.cache_codec as cache_codec @@ -468,7 +533,10 @@ def boot() -> None: MT = MessageType # local_cache_global # Register workflow handlers - if not utils.USE_THP: + if utils.USE_THP: + TMT = ThpMessageType + workflow_handlers.register(TMT.ThpCreateNewSession, handle_ThpCreateNewSession) + else: workflow_handlers.register(MT.Initialize, handle_Initialize) for msg_type, handler in [ (MT.GetFeatures, handle_GetFeatures), diff --git a/core/src/apps/thp/create_new_session.py b/core/src/apps/thp/create_new_session.py deleted file mode 100644 index 156b852d46..0000000000 --- a/core/src/apps/thp/create_new_session.py +++ /dev/null @@ -1,59 +0,0 @@ -from trezor import log, loop -from trezor.enums import FailureType -from trezor.messages import Failure, ThpCreateNewSession, ThpNewSession -from trezor.wire.context import get_context -from trezor.wire.errors import ActionCancelled, DataError -from trezor.wire.thp import SessionState - - -async def create_new_session(message: ThpCreateNewSession) -> ThpNewSession | Failure: - """ - Creates a new `ThpSession` based on the provided parameters and returns a - `ThpNewSession` message containing the new session ID. - - Returns an appropriate `Failure` message if session creation fails. - """ - from trezor.wire import NotInitialized - from trezor.wire.thp.session_context import GenericSessionContext - from trezor.wire.thp.session_manager import create_new_session - - from apps.common.seed import derive_and_store_roots - - ctx = get_context() - - # Assert that context `ctx` is `GenericSessionContext` - assert isinstance(ctx, GenericSessionContext) - - channel = ctx.channel - - # Do not use `ctx` beyond this point, as it is techically - # allowed to change in between await statements - - new_session = create_new_session(channel) - try: - await derive_and_store_roots(new_session, message) - except DataError as e: - return Failure(code=FailureType.DataError, message=e.message) - except ActionCancelled as e: - return Failure(code=FailureType.ActionCancelled, message=e.message) - except NotInitialized as e: - return Failure(code=FailureType.NotInitialized, message=e.message) - # TODO handle other errors (`Exception`` when "Cardano icarus secret is already set!" - # and `RuntimeError` when accessing storage for mnemonic.get_secret - it actually - # happens for locked devices) - - new_session.set_session_state(SessionState.ALLOCATED) - channel.sessions[new_session.session_id] = new_session - loop.schedule(new_session.handle()) - new_session_id: int = new_session.session_id - - if __debug__: - log.debug( - __name__, - "create_new_session - new session created. Passphrase: %s, Session id: %d\n%s", - message.passphrase if message.passphrase is not None else "", - new_session.session_id, - str(channel.sessions), - ) - - return ThpNewSession(new_session_id=new_session_id) diff --git a/core/src/apps/workflow_handlers.py b/core/src/apps/workflow_handlers.py index 3013516382..b65c853c93 100644 --- a/core/src/apps/workflow_handlers.py +++ b/core/src/apps/workflow_handlers.py @@ -35,13 +35,6 @@ def _find_message_handler_module(msg_type: int) -> str: if __debug__ and msg_type == MessageType.BenchmarkRun: return "apps.benchmark.run" - if utils.USE_THP: - from trezor.enums import ThpMessageType - - # thp management - if msg_type == ThpMessageType.ThpCreateNewSession: - return "apps.thp.create_new_session" - # management if msg_type == MessageType.ResetDevice: return "apps.management.reset_device"