Enable GC dump on OOM

core/SConscript.firmware

@@ -395,9 +395,17 @@ ui.init_ui(TREZOR_MODEL, "firmware", RUST_UI_FEATURES)
 
 SOURCE_QSTR = SOURCE_MOD + SOURCE_MICROPYTHON + SOURCE_MICROPYTHON_SPEED
 
+if PYOPT == '0':
+    DEBUG_FLAGS = "-DMICROPY_OOM_CALLBACK=1 -DSTATIC="
+else:
+    DEBUG_FLAGS = "-DMICROPY_OOM_CALLBACK=0"
+
 env = Environment(
     ENV=os.environ,
-    CFLAGS=f"{ARGUMENTS.get('CFLAGS', '')} -DPRODUCTION={int(PRODUCTION)} -DPYOPT={PYOPT} -DBOOTLOADER_QA={int(BOOTLOADER_QA)} -DBITCOIN_ONLY={BITCOIN_ONLY}",
+    CFLAGS=(
+        f"{ARGUMENTS.get('CFLAGS', '')} -DPRODUCTION={int(PRODUCTION)} -DPYOPT={PYOPT} "
+        f"-DBOOTLOADER_QA={int(BOOTLOADER_QA)} -DBITCOIN_ONLY={BITCOIN_ONLY} {DEBUG_FLAGS}"
+    ),
     CPPDEFINES_IMPLICIT=[],
     CPPDEFPREFIX="-D'",
     CPPDEFSUFFIX="'",

core/SConscript.unix

@@ -422,13 +422,13 @@ ui.init_ui(TREZOR_MODEL, "firmware", RUST_UI_FEATURES)
 SOURCE_QSTR = SOURCE_MOD + SOURCE_MICROPYTHON + SOURCE_UNIX
 
 if PYOPT == '0' or not FROZEN:
-    STATIC="-DSTATIC="
+    DEBUG_FLAGS = "-DMICROPY_OOM_CALLBACK=1 -DSTATIC="
 else:
-    STATIC=""
+    DEBUG_FLAGS = "-DMICROPY_OOM_CALLBACK=0"
 
 env = Environment(
     ENV=os.environ,
-    CFLAGS=ARGUMENTS.get("CFLAGS", "") + f" -DCONFIDENTIAL= -DPYOPT={PYOPT} -DBITCOIN_ONLY={BITCOIN_ONLY} {STATIC}",
+    CFLAGS=ARGUMENTS.get("CFLAGS", "") + f" -DCONFIDENTIAL= -DPYOPT={PYOPT} -DBITCOIN_ONLY={BITCOIN_ONLY} {DEBUG_FLAGS}",
     CPPDEFPREFIX="-D'",
     CPPDEFSUFFIX="'",
 )

core/embed/rust/librust.h

@@ -2,7 +2,7 @@
 
 #include "librust_qstr.h"
 
-#ifdef TREZOR_EMULATOR
+#if !PYOPT
 mp_obj_t protobuf_debug_msg_type();
 mp_obj_t protobuf_debug_msg_def_type();
 #endif
@@ -11,6 +11,6 @@ extern mp_obj_module_t mp_module_trezorproto;
 extern mp_obj_module_t mp_module_trezorui_api;
 extern mp_obj_module_t mp_module_trezortranslate;
 
-#ifdef TREZOR_EMULATOR
+#if !PYOPT
 mp_obj_t ui_debug_layout_type();
 #endif

core/embed/upymod/modtrezorutils/modtrezorutils-meminfo.h

@@ -17,7 +17,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-#if !TREZOR_EMULATOR || PYOPT
+#if PYOPT
 #define MEMINFO_DICT_ENTRIES /* empty */
 
 #else
@@ -38,6 +38,14 @@
 #include "embed/rust/librust.h"
 #include "embed/upymod/trezorobj.h"
 
+#if !TREZOR_EMULATOR
+#define fopen(path, mode) &mp_plat_print
+#define fprintf mp_printf
+#define fflush(f)
+#define fclose(f)
+#define FILE const mp_print_t
+#endif
+
 #define WORDS_PER_BLOCK ((MICROPY_BYTES_PER_GC_BLOCK) / MP_BYTES_PER_OBJ_WORD)
 #define BYTES_PER_BLOCK (MICROPY_BYTES_PER_GC_BLOCK)
 
@@ -149,24 +157,38 @@ bool is_short(mp_const_obj_t value) {
          mp_obj_is_small_int(value) || !VERIFY_PTR(value);
 }
 
+static void escape_and_dump_string(FILE *out, const char *unescaped) {
+  fprintf(out, "\"");
+  for (; *unescaped; ++unescaped) {
+    char c = *unescaped;
+    if (c == '\n') {
+      fprintf(out, "\\n");
+    } else if (c == '\r') {
+      fprintf(out, "\\r");
+    } else if (c == '\"') {
+      fprintf(out, "\\\"");
+    } else if (c == '\\') {
+      fprintf(out, "\\\\");
+    } else if (c >= 0x20 && c < 0x7F) {
+      fprintf(out, "%c", c);
+    } else {
+      fprintf(out, "\\u%04x", c);
+    }
+  }
+  fprintf(out, "\"");
+}
+
 static void print_type(FILE *out, const char *typename, const char *shortval,
                        const void *ptr, bool end) {
-  static char unescaped[1000];
   size_t size = 0;
   if (!is_short(ptr)) {
     size = find_allocated_size(ptr);
   }
   fprintf(out, "{\"type\": \"%s\", \"alloc\": %ld, \"ptr\": \"%p\"", typename,
           size, ptr);
   if (shortval) {
-    assert(strlen(shortval) < 1000);
-    char *c = unescaped;
-    while (*shortval) {
-      if (*shortval == '\\' || *shortval == '"') *c++ = '\\';
-      *c++ = *shortval++;
-    }
-    *c = 0;
-    fprintf(out, ", \"shortval\": \"%s\"", unescaped);
+    fprintf(out, ", \"shortval\": ");
+    escape_and_dump_string(out, shortval);
   } else {
     fprintf(out, ", \"shortval\": null");
   }
@@ -199,7 +221,7 @@ void dump_short(FILE *out, mp_const_obj_t value) {
 
   } else if (mp_obj_is_small_int(value)) {
     static char num_buf[100];
-    snprintf(num_buf, 100, "%ld", MP_OBJ_SMALL_INT_VALUE(value));
+    snprintf(num_buf, 100, INT_FMT, MP_OBJ_SMALL_INT_VALUE(value));
     print_type(out, "smallint", num_buf, NULL, true);
 
   } else if (!VERIFY_PTR(value)) {
@@ -680,10 +702,11 @@ void dump_qstr_pool(FILE *out, const qstr_pool_t *pool) {
   for (const char *const *q = pool->qstrs, *const *q_top =
                                                pool->qstrs + pool->len;
        q < q_top; q++) {
+    escape_and_dump_string(out, Q_GET_DATA(*q));
     if (q < (q_top - 1))
-      fprintf(out, "\"%s\",\n", Q_GET_DATA(*q));
+      fprintf(out, ",\n");
     else
-      fprintf(out, "\"%s\"]\n", Q_GET_DATA(*q));
+      fprintf(out, "]\n");
   }
   fprintf(out, "},\n");
   for (const char *const *q = pool->qstrs, *const *q_top =
@@ -709,15 +732,19 @@ void dump_qstrdata(FILE *out) {
   }
 }
 
-/// def meminfo(filename: str) -> None:
-///     """Dumps map of micropython GC arena to a file.
-///     The JSON file can be decoded by analyze-memory-dump.py
-///     Only available in the emulator.
-///      """
-STATIC mp_obj_t mod_trezorutils_meminfo(mp_obj_t filename) {
-  size_t fn_len;
-  FILE *out = fopen(mp_obj_str_get_data(filename, &fn_len), "w");
-  fprintf(out, "[");
+static void dump_meminfo_json(FILE *out) {
+  bool should_close = true;
+  if (out == NULL) {
+    should_close = false;
+#if TREZOR_EMULATOR
+    out = stdout;
+#else
+    out = &mp_plat_print;
+#endif
+  }
+  fprintf(out, "\n[\n[" UINT_FMT ", " UINT_FMT ", " UINT_FMT "],\n",
+          (mp_uint_t)MP_STATE_MEM(gc_pool_start),
+          (mp_uint_t)MP_STATE_MEM(gc_pool_end), BYTES_PER_BLOCK);
 
   // void **ptrs = (void **)(void *)&mp_state_ctx;
   // size_t root_start = offsetof(mp_state_ctx_t, thread.dict_locals);
@@ -768,8 +795,12 @@ STATIC mp_obj_t mod_trezorutils_meminfo(mp_obj_t filename) {
     pool = pool->prev;
   }
 
-  fprintf(out, "null]\n");
-  fclose(out);
+  fprintf(out, "null\n]\n");
+  if (should_close) {
+    fclose(out);
+  } else {
+    fflush(out);
+  }
   for (size_t block = 0;
        block < MP_STATE_MEM(gc_alloc_table_byte_len) * BLOCKS_PER_ATB;
        block++) {
@@ -779,6 +810,19 @@ STATIC mp_obj_t mod_trezorutils_meminfo(mp_obj_t filename) {
   }
 
   gc_dump_alloc_table();
+}
+
+/// def meminfo(filename: str | None) -> None:
+///     """Dumps map of micropython GC arena to a file.
+///     The JSON file can be decoded by analyze-memory-dump.py
+///      """
+STATIC mp_obj_t mod_trezorutils_meminfo(mp_obj_t filename) {
+  size_t fn_len;
+  FILE *out = (filename == mp_const_none)
+                  ? NULL
+                  : fopen(mp_obj_str_get_data(filename, &fn_len), "w");
+  (void)fn_len;
+  dump_meminfo_json(out);
   return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorutils_meminfo_obj,

core/embed/upymod/modtrezorutils/modtrezorutils.c

@@ -20,6 +20,9 @@
 #include <trezor_model.h>
 #include <trezor_rtl.h>
 
+#if MICROPY_OOM_CALLBACK
+#include <py/gc.h>
+#endif
 #include "py/objstr.h"
 #include "py/runtime.h"
 
@@ -243,7 +246,8 @@ 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
+#if !PYOPT
+#if LOG_STACK_USAGE
 /// def zero_unused_stack() -> None:
 ///     """
 ///     Zero unused stack memory.
@@ -275,7 +279,30 @@ STATIC mp_obj_t mod_trezorutils_estimate_unused_stack(void) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_estimate_unused_stack_obj,
                                  mod_trezorutils_estimate_unused_stack);
 
-#endif  // !PYOPT && LOG_STACK_USAGE
+#endif  // LOG_STACK_USAGE
+
+#if MICROPY_OOM_CALLBACK
+static void gc_oom_callback(void) {
+  gc_dump_info();
+#if BLOCK_ON_VCP || TREZOR_EMULATOR
+  dump_meminfo_json(NULL);  // dump to stdout
+#endif
+}
+
+/// if __debug__:
+///     def enable_oom_dump() -> None:
+///         """
+///         Dump GC info in case of an OOM.
+///         """
+STATIC mp_obj_t mod_trezorutils_enable_oom_dump(void) {
+  gc_set_oom_callback(gc_oom_callback);
+  return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_enable_oom_dump_obj,
+                                 mod_trezorutils_enable_oom_dump);
+#endif  // MICROPY_OOM_CALLBACK
+
+#endif  // !PYOPT
 
 /// def reboot_to_bootloader(
 ///     boot_command : int = 0,
@@ -482,11 +509,17 @@ 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
+#if !PYOPT
+#if 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
+#if MICROPY_OOM_CALLBACK
+    {MP_ROM_QSTR(MP_QSTR_enable_oom_dump),
+     MP_ROM_PTR(&mod_trezorutils_enable_oom_dump_obj)},
+#endif
 #endif
     {MP_ROM_QSTR(MP_QSTR_sd_hotswap_enabled),
      MP_ROM_PTR(&mod_trezorutils_sd_hotswap_enabled_obj)},

core/mocks/generated/trezorutils.pyi

@@ -2,10 +2,9 @@ from typing import *
 
 
 # upymod/modtrezorutils/modtrezorutils-meminfo.h
-def meminfo(filename: str) -> None:
+def meminfo(filename: str | None) -> None:
     """Dumps map of micropython GC arena to a file.
     The JSON file can be decoded by analyze-memory-dump.py
-    Only available in the emulator.
      """
 
 
@@ -100,6 +99,11 @@ def estimate_unused_stack() -> int:
     """
     Estimate unused stack size.
     """
+if __debug__:
+    def enable_oom_dump() -> None:
+        """
+        Dump GC info in case of an OOM.
+        """
 
 
 # upymod/modtrezorutils/modtrezorutils.c

core/src/trezor/utils.py

@@ -125,6 +125,13 @@ def presize_module(modname: str, size: int) -> None:
 if __debug__:
     from ubinascii import hexlify
 
+    try:
+        from trezorutils import enable_oom_dump
+
+        enable_oom_dump()
+    except ImportError:
+        pass
+
     def mem_dump(filename: str) -> None:
         from micropython import mem_info
 

core/tools/analyze-memory-dump.py

@@ -30,7 +30,8 @@
 
 
 with open(sys.argv[1]) as f:
-    MEMMAP = json.load(f)
+    MEMMAP = iter(json.load(f))
+    (min_ptr, max_ptr, bytes_per_block) = next(MEMMAP)
 
 
 # filter out notices and comments
@@ -57,7 +58,13 @@ def ptr_or_shortval(maybe_ptr):
 
 
 def is_ignored_ptr(ptr):
-    return (ptr == "(nil)" or ptr.startswith("0x5") or ptr.startswith("0x6"))
+    if ptr == "(nil)":
+        return True
+
+    if isinstance(ptr, str):
+        ptr = int(ptr, 16)
+
+    return not (min_ptr <= ptr < max_ptr)
 
 
 def deref_or_shortval(maybe_ptr):
@@ -155,13 +162,6 @@ def ptrval(self):
 
 allobjs = list(MEMORY.values())
 allobjs.sort(key=lambda x: x.ptr)
-min_ptr = min(
-    item.ptrval()
-    for item in allobjs
-    if not is_ignored_ptr(item.ptr)
-)
-max_ptr = max(item.ptrval() for item in allobjs if item.ptr != "(nil)")
-
 
 types = {
     "anystr": "S",
@@ -201,9 +201,10 @@ def ptrval(self):
 pixels_per_line = len(
     "................................................................"
 )
-pixelsize = 0x800 // pixels_per_line
-maxline = ((max_ptr - min_ptr) & ~0x7FF) + (0x800 * 2)
-pixelmap = [None] * (maxline // pixelsize)
+pixelsize = bytes_per_block
+bytes_per_line = bytes_per_block * pixels_per_line
+maxline = ((max_ptr - min_ptr) & ~(bytes_per_line - 1)) + (bytes_per_line * 2)
+pixelmap = [None] * 2*(maxline // pixelsize)
 
 
 def pixel_index(ptrval):