diff --git a/core/Makefile b/core/Makefile index 6a5fe4d97e..7df43189e7 100644 --- a/core/Makefile +++ b/core/Makefile @@ -40,6 +40,7 @@ PYTEST_TIMEOUT ?= 500 TEST_LANG ?= "en" THP ?= 0 BENCHMARK ?= 0 +LOG_STACK_USAGE ?= 0 TREZOR_EMULATOR_DEBUGGABLE ?= 0 QUIET_MODE ?= 0 TREZOR_DISABLE_ANIMATION ?= $(if $(filter 0,$(PYOPT)),1,0) @@ -121,6 +122,7 @@ SCONS_VARS = \ CFLAGS="$(CFLAGS)" \ CMAKELISTS="$(CMAKELISTS)" \ DISABLE_OPTIGA="$(DISABLE_OPTIGA)" \ + LOG_STACK_USAGE="$(LOG_STACK_USAGE)" \ PRODUCTION="$(PRODUCTION)" \ PYOPT="$(PYOPT)" \ QUIET_MODE="$(QUIET_MODE)" \ diff --git a/core/SConscript.firmware b/core/SConscript.firmware index 43d23d0972..6d2ff61a34 100644 --- a/core/SConscript.firmware +++ b/core/SConscript.firmware @@ -19,6 +19,7 @@ HW_REVISION = ARGUMENTS.get('HW_REVISION', None) SCM_REVISION = ARGUMENTS.get('SCM_REVISION', None) THP = ARGUMENTS.get('THP', '0') == '1' # Trezor-Host Protocol BENCHMARK = ARGUMENTS.get('BENCHMARK', '0') == '1' +LOG_STACK_USAGE=ARGUMENTS.get('LOG_STACK_USAGE', '0') == '1' DISABLE_ANIMATION = ARGUMENTS.get('TREZOR_DISABLE_ANIMATION', '0') == '1' UI_DEBUG_OVERLAY = ARGUMENTS.get('UI_DEBUG_OVERLAY', '0') == '1' @@ -33,6 +34,10 @@ if BENCHMARK and PYOPT != '0': print("BENCHMARK=1 works only with PYOPT=0.") exit(1) +if LOG_STACK_USAGE and PYOPT != '0': + print("LOG_STACK_USAGE=1 works only with PYOPT=0.") + exit(1) + FEATURE_FLAGS = { "RDI": True, "SECP256K1_ZKP": True, # required for trezor.crypto.curve.bip340 (BIP340/Taproot) @@ -82,6 +87,7 @@ CPPDEFINES_MOD += [ ('USE_NEM', '1' if (EVERYTHING and TREZOR_MODEL == "T2T1") else '0'), ('USE_EOS', '1' if (EVERYTHING and TREZOR_MODEL == "T2T1") else '0'), ('DISABLE_ANIMATION', '1' if DISABLE_ANIMATION else '0'), + ('LOG_STACK_USAGE', '1' if LOG_STACK_USAGE else '0'), ] SOURCE_MOD += [ 'embed/upymod/trezorobj.c', diff --git a/core/embed/upymod/modtrezorutils/modtrezorutils.c b/core/embed/upymod/modtrezorutils/modtrezorutils.c index 8869206a41..e09b752abd 100644 --- a/core/embed/upymod/modtrezorutils/modtrezorutils.c +++ b/core/embed/upymod/modtrezorutils/modtrezorutils.c @@ -42,6 +42,10 @@ #include #endif +#if !PYOPT && LOG_STACK_USAGE +#include +#endif + static void ui_progress(void *context, uint32_t current, uint32_t total) { mp_obj_t ui_wait_callback = (mp_obj_t)context; @@ -239,6 +243,40 @@ STATIC mp_obj_t mod_trezorutils_sd_hotswap_enabled(void) { STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_sd_hotswap_enabled_obj, mod_trezorutils_sd_hotswap_enabled); +#if !PYOPT && LOG_STACK_USAGE +/// def zero_unused_stack() -> None: +/// """ +/// Zero unused stack memory. +/// """ +STATIC mp_obj_t mod_trezorutils_zero_unused_stack(void) { + clear_unused_stack(); + return mp_const_none; +} +STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_zero_unused_stack_obj, + mod_trezorutils_zero_unused_stack); + +/// def estimate_unused_stack() -> int: +/// """ +/// Estimate unused stack size. +/// """ +STATIC mp_obj_t mod_trezorutils_estimate_unused_stack(void) { + const uint8_t *stack_top = (const uint8_t *)MP_STATE_THREAD(stack_top); + size_t stack_limit = MP_STATE_THREAD(stack_limit); + + const uint8_t *stack = stack_top - stack_limit; + size_t offset = 0; + for (; offset < stack_limit; ++offset) { + if (stack[offset] != 0) { + break; + } + } + return mp_obj_new_int_from_uint(offset); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_estimate_unused_stack_obj, + mod_trezorutils_estimate_unused_stack); + +#endif // !PYOPT && LOG_STACK_USAGE + /// def reboot_to_bootloader( /// boot_command : int = 0, /// boot_args : bytes | None = None, @@ -418,6 +456,8 @@ STATIC mp_obj_tuple_t mod_trezorutils_version_obj = { /// if __debug__: /// DISABLE_ANIMATION: bool /// """Whether the firmware should disable animations.""" +/// LOG_STACK_USAGE: bool +/// """Whether the firmware should log estimated stack usage.""" STATIC const mp_rom_map_elem_t mp_module_trezorutils_globals_table[] = { {MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_trezorutils)}, @@ -440,6 +480,12 @@ STATIC const mp_rom_map_elem_t mp_module_trezorutils_globals_table[] = { MP_ROM_PTR(&mod_trezorutils_unit_packaging_obj)}, {MP_ROM_QSTR(MP_QSTR_unit_btconly), MP_ROM_PTR(&mod_trezorutils_unit_btconly_obj)}, +#if !PYOPT && LOG_STACK_USAGE + {MP_ROM_QSTR(MP_QSTR_zero_unused_stack), + MP_ROM_PTR(&mod_trezorutils_zero_unused_stack_obj)}, + {MP_ROM_QSTR(MP_QSTR_estimate_unused_stack), + MP_ROM_PTR(&mod_trezorutils_estimate_unused_stack_obj)}, +#endif {MP_ROM_QSTR(MP_QSTR_sd_hotswap_enabled), MP_ROM_PTR(&mod_trezorutils_sd_hotswap_enabled_obj)}, // various built-in constants @@ -523,6 +569,11 @@ STATIC const mp_rom_map_elem_t mp_module_trezorutils_globals_table[] = { #else {MP_ROM_QSTR(MP_QSTR_DISABLE_ANIMATION), mp_const_false}, #endif // TREZOR_DISABLE_ANIMATION +#if LOG_STACK_USAGE + {MP_ROM_QSTR(MP_QSTR_LOG_STACK_USAGE), mp_const_true}, +#else + {MP_ROM_QSTR(MP_QSTR_LOG_STACK_USAGE), mp_const_false}, +#endif // LOG_STACK_USAGE #endif // PYOPT }; diff --git a/core/mocks/generated/trezorutils.pyi b/core/mocks/generated/trezorutils.pyi index 8cf7cc44ec..3bab3d7cf1 100644 --- a/core/mocks/generated/trezorutils.pyi +++ b/core/mocks/generated/trezorutils.pyi @@ -88,6 +88,20 @@ def sd_hotswap_enabled() -> bool: """ +# upymod/modtrezorutils/modtrezorutils.c +def zero_unused_stack() -> None: + """ + Zero unused stack memory. + """ + + +# upymod/modtrezorutils/modtrezorutils.c +def estimate_unused_stack() -> int: + """ + Estimate unused stack size. + """ + + # upymod/modtrezorutils/modtrezorutils.c def reboot_to_bootloader( boot_command : int = 0, @@ -159,3 +173,5 @@ USE_THP: bool if __debug__: DISABLE_ANIMATION: bool """Whether the firmware should disable animations.""" + LOG_STACK_USAGE: bool + """Whether the firmware should log estimated stack usage.""" diff --git a/core/src/trezor/utils.py b/core/src/trezor/utils.py index 6fcfb085f1..76603efa15 100644 --- a/core/src/trezor/utils.py +++ b/core/src/trezor/utils.py @@ -36,18 +36,24 @@ from trezorutils import ( # noqa: F401 from typing import TYPE_CHECKING if __debug__: + from trezorutils import LOG_STACK_USAGE + + if LOG_STACK_USAGE: + from trezorutils import estimate_unused_stack, zero_unused_stack # noqa: F401 + if EMULATOR: import uos DISABLE_ANIMATION = uos.getenv("TREZOR_DISABLE_ANIMATION") == "1" LOG_MEMORY = uos.getenv("TREZOR_LOG_MEMORY") == "1" else: - from trezorutils import DISABLE_ANIMATION # noqa: F401 + from trezorutils import DISABLE_ANIMATION LOG_MEMORY = 0 else: DISABLE_ANIMATION = False + LOG_STACK_USAGE = False if TYPE_CHECKING: from typing import Any, Iterator, Protocol, Sequence, TypeVar diff --git a/core/src/trezor/wire/message_handler.py b/core/src/trezor/wire/message_handler.py index 21c901dc90..13be61e8cc 100644 --- a/core/src/trezor/wire/message_handler.py +++ b/core/src/trezor/wire/message_handler.py @@ -97,6 +97,10 @@ async def handle_single_message(ctx: Context, msg: Message) -> bool: # `req_type`. Raises if the message is malformed. req_msg = wrap_protobuf_load(msg.data, req_type) + if __debug__ and utils.LOG_STACK_USAGE: + utils.zero_unused_stack() + unused_stack_before = utils.estimate_unused_stack() + # Create the handler task. task = handler(req_msg) @@ -109,6 +113,17 @@ async def handle_single_message(ctx: Context, msg: Message) -> bool: # workflows are shut down. res_msg = await workflow.spawn(with_context(ctx, task)) + if __debug__ and utils.LOG_STACK_USAGE: + unused_stack_after = utils.estimate_unused_stack() + log.debug( + __name__, + "<%s> estimated stack usage=%d (unused before=%d, after=%d)", + msg_type, + unused_stack_before - unused_stack_after, + unused_stack_before, + unused_stack_after, + ) + except UnexpectedMessageException: # Workflow was trying to read a message from the wire, and # something unexpected came in. See Context.read() for