From 41db573ef9882933abcf6c5d5a35e6d9a74c213f Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@satoshilabs.com>
Date: Mon, 21 Apr 2025 13:10:27 +0300
Subject: [PATCH] feat(core): check that GC restores free heap memory

Enabled only for frozen debug builds.

[no changelog]
---
 .../upymod/modtrezorutils/modtrezorutils.c    | 27 +++++++++++++++++++
 core/mocks/generated/trezorutils.pyi          |  7 +++++
 core/src/main.py                              |  6 ++---
 core/src/trezor/utils.py                      |  9 +++++--
 4 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/core/embed/upymod/modtrezorutils/modtrezorutils.c b/core/embed/upymod/modtrezorutils/modtrezorutils.c
index 5799f9ac07..eef23b41e1 100644
--- a/core/embed/upymod/modtrezorutils/modtrezorutils.c
+++ b/core/embed/upymod/modtrezorutils/modtrezorutils.c
@@ -302,6 +302,31 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_enable_oom_dump_obj,
                                  mod_trezorutils_enable_oom_dump);
 #endif  // MICROPY_OOM_CALLBACK
 
+/// if __debug__:
+///     def check_free_heap(previous: int) -> int:
+///         """
+///         Assert that free heap memory doesn't decrease.
+///         Returns current free heap memory (in bytes).
+///         Enabled only for frozen debug builds.
+///         """
+STATIC mp_obj_t mod_trezorutils_check_free_heap(mp_obj_t arg) {
+  mp_uint_t free_heap = trezor_obj_get_uint(arg);
+#if MICROPY_MODULE_FROZEN_MPY
+  gc_info_t info;
+  gc_info(&info);
+  if (free_heap > info.free) {
+    gc_dump_info();
+    mp_raise_msg_varg(&mp_type_AssertionError,
+                      "Free heap decreased by " UINT_FMT " bytes",
+                      free_heap - info.free);
+  }
+  free_heap = info.free;  // current free heap
+#endif
+  return mp_obj_new_int_from_uint(free_heap);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorutils_check_free_heap_obj,
+                                 mod_trezorutils_check_free_heap);
+
 /// if __debug__:
 ///     def check_heap_fragmentation() -> None:
 ///         """
@@ -559,6 +584,8 @@ STATIC const mp_rom_map_elem_t mp_module_trezorutils_globals_table[] = {
     {MP_ROM_QSTR(MP_QSTR_enable_oom_dump),
      MP_ROM_PTR(&mod_trezorutils_enable_oom_dump_obj)},
 #endif
+    {MP_ROM_QSTR(MP_QSTR_check_free_heap),
+     MP_ROM_PTR(&mod_trezorutils_check_free_heap_obj)},
     {MP_ROM_QSTR(MP_QSTR_check_heap_fragmentation),
      MP_ROM_PTR(&mod_trezorutils_check_heap_fragmentation_obj)},
 #endif
diff --git a/core/mocks/generated/trezorutils.pyi b/core/mocks/generated/trezorutils.pyi
index 2fae11f5e1..7da33c4078 100644
--- a/core/mocks/generated/trezorutils.pyi
+++ b/core/mocks/generated/trezorutils.pyi
@@ -104,6 +104,13 @@ if __debug__:
         """
         Dump GC info in case of an OOM.
         """
+if __debug__:
+    def check_free_heap(previous: int) -> int:
+        """
+        Assert that free heap memory doesn't decrease.
+        Returns current free heap memory (in bytes).
+        Enabled only for frozen debug builds.
+        """
 if __debug__:
     def check_heap_fragmentation() -> None:
         """
diff --git a/core/src/main.py b/core/src/main.py
index ed6a0311af..e70aab913f 100644
--- a/core/src/main.py
+++ b/core/src/main.py
@@ -36,11 +36,8 @@ import trezor.pin  # noqa: F401
 # usb imports trezor.utils and trezor.io which is a C module
 import usb
 
-# create an unimport manager that will be reused in the main loop
-unimport_manager = utils.unimport()
-
 # unlock the device, unload the boot module afterwards
-with unimport_manager:
+with utils.unimport():
     import boot
     del boot
 
@@ -63,6 +60,7 @@ if utils.USE_BLE:
 
 
 # run the endless loop
+unimport_manager = utils.unimport()
 while True:
     with unimport_manager:
         import session  # noqa: F401
diff --git a/core/src/trezor/utils.py b/core/src/trezor/utils.py
index fd7ddede04..55126e262a 100644
--- a/core/src/trezor/utils.py
+++ b/core/src/trezor/utils.py
@@ -37,7 +37,7 @@ from trezorutils import (  # noqa: F401
 from typing import TYPE_CHECKING
 
 if __debug__:
-    from trezorutils import LOG_STACK_USAGE, check_heap_fragmentation
+    from trezorutils import LOG_STACK_USAGE, check_free_heap, check_heap_fragmentation
 
     if LOG_STACK_USAGE:
         from trezorutils import estimate_unused_stack, zero_unused_stack  # noqa: F401
@@ -94,17 +94,22 @@ def unimport_end(mods: set[str], collect: bool = True) -> None:
 class unimport:
     def __init__(self) -> None:
         self.mods: set[str] | None = None
+        if __debug__:
+            self.free_heap = 0
 
     def __enter__(self) -> None:
         self.mods = unimport_begin()
 
-    def __exit__(self, _exc_type: Any, _exc_value: Any, _tb: Any) -> None:
+    def __exit__(self, exc_type: Any, _exc_value: Any, _tb: Any) -> None:
         assert self.mods is not None
         unimport_end(self.mods, collect=False)
         self.mods.clear()
         self.mods = None
         gc.collect()
 
+        if __debug__ and exc_type is not SystemExit:
+            self.free_heap = check_free_heap(self.free_heap)
+
 
 def presize_module(modname: str, size: int) -> None:
     """Ensure the module's dict is preallocated to an expected size.