diff --git a/core/embed/upymod/modtrezorutils/modtrezorutils.c b/core/embed/upymod/modtrezorutils/modtrezorutils.c index 8d5fb8f8f2..f63ad17810 100644 --- a/core/embed/upymod/modtrezorutils/modtrezorutils.c +++ b/core/embed/upymod/modtrezorutils/modtrezorutils.c @@ -302,6 +302,27 @@ 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_reallocs() -> None: +/// """ +/// Assert that `sys.modules` and `main`'s globals are never +/// reallocated. +/// """ +STATIC mp_obj_t mod_trezorutils_check_reallocs(void) { + mp_obj_dict_t *modules = &MP_STATE_VM(mp_loaded_modules_dict); + if (modules->map.alloc > MICROPY_LOADED_MODULES_DICT_SIZE) { + mp_raise_msg(&mp_type_AssertionError, "sys.modules dict is reallocated"); + } + mp_obj_t main = mp_obj_dict_get(modules, MP_OBJ_NEW_QSTR(MP_QSTR_main)); + size_t main_map_alloc = mp_obj_module_get_globals(main)->map.alloc; + if (main_map_alloc > MICROPY_MAIN_DICT_SIZE) { + mp_raise_msg(&mp_type_AssertionError, "main globals dict is reallocated"); + } + return mp_const_none; +} +STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorutils_check_reallocs_obj, + mod_trezorutils_check_reallocs); + #endif // !PYOPT /// def reboot_to_bootloader( @@ -520,6 +541,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_reallocs), + MP_ROM_PTR(&mod_trezorutils_check_reallocs_obj)}, #endif {MP_ROM_QSTR(MP_QSTR_sd_hotswap_enabled), MP_ROM_PTR(&mod_trezorutils_sd_hotswap_enabled_obj)}, diff --git a/core/mocks/generated/trezorutils.pyi b/core/mocks/generated/trezorutils.pyi index 2ec1af7e11..e2b708baf0 100644 --- a/core/mocks/generated/trezorutils.pyi +++ b/core/mocks/generated/trezorutils.pyi @@ -104,6 +104,12 @@ if __debug__: """ Dump GC info in case of an OOM. """ +if __debug__: + def check_reallocs() -> None: + """ + Assert that `sys.modules` and `main`'s globals are never + reallocated. + """ # upymod/modtrezorutils/modtrezorutils.c diff --git a/core/src/trezor/utils.py b/core/src/trezor/utils.py index ebb0057560..457b661593 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 + from trezorutils import LOG_STACK_USAGE, check_reallocs if LOG_STACK_USAGE: from trezorutils import estimate_unused_stack, zero_unused_stack # noqa: F401 @@ -67,10 +67,9 @@ def unimport_begin() -> set[str]: def unimport_end(mods: set[str], collect: bool = True) -> None: - # static check that the size of sys.modules never grows above value of - # MICROPY_LOADED_MODULES_DICT_SIZE, so that the sys.modules dict is never - # reallocated at run-time - assert len(sys.modules) <= 160, "Please bump preallocated size in mpconfigport.h" + if __debug__: + # Assert that `sys.modules` and `__main__` dict are never reallocated at run-time + check_reallocs() for mod in sys.modules: # pylint: disable=consider-using-dict-items if mod not in mods: