feat(core): estimate workflow stack usage

By zeroing the stack memory before the workflow runs,
we can estimate how much of it has been used (by reading
the stack memory and looking for the first non-zero value).

[no changelog]

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)" \

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',

core/embed/upymod/modtrezorutils/modtrezorutils.c

@@ -42,6 +42,10 @@
 #include <sec/secret.h>
 #endif
 
+#if !PYOPT && LOG_STACK_USAGE
+#include <sys/stack_utils.h>
+#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
 };
 

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."""

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

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