From 7ca67cc4d9d3adebf313923e875ac4609653c2c8 Mon Sep 17 00:00:00 2001 From: matejcik Date: Mon, 22 Mar 2021 15:47:43 +0100 Subject: [PATCH] feat(core): clear out memory space after every workflow A small fixed list of modules is kept pre-loaded in the GC arena. These must not keep references to anything else, as all other modules are unloaded and the memory is cleared. --- core/src/main.py | 82 +++++++------------------------- core/src/session.py | 28 +++++++++++ core/src/trezor/utils.py | 24 ++++++++-- core/src/trezor/wire/__init__.py | 11 ++++- 4 files changed, 74 insertions(+), 71 deletions(-) create mode 100644 core/src/session.py diff --git a/core/src/main.py b/core/src/main.py index e1a721fdf..54606e6db 100644 --- a/core/src/main.py +++ b/core/src/main.py @@ -1,73 +1,25 @@ # isort:skip_file +# fmt: off -# unlock the device -import boot # noqa: F401 +# Import always-active modules +import storage +import storage.device +from trezor import config, pin, utils # noqa: F401 -# prepare the USB interfaces, but do not connect to the host yet +# Prepare the USB interfaces first. Do not connect to the host yet. import usb -from trezor import loop, utils, wire, workflow +unimport_manager = utils.unimport() + +# unlock the device, unload the boot module afterwards +with unimport_manager: + import boot + del boot # start the USB -usb.bus.open() +usb.bus.open(storage.device.get_device_id()) - -def _boot_apps() -> None: - # load applications - import apps.base - import apps.management - import apps.bitcoin - import apps.misc - - if not utils.BITCOIN_ONLY: - import apps.ethereum - import apps.lisk - import apps.monero - import apps.nem - import apps.stellar - import apps.ripple - import apps.cardano - import apps.tezos - import apps.eos - import apps.binance - import apps.webauthn - - if __debug__: - import apps.debug - - # boot applications - apps.base.boot() - apps.management.boot() - apps.bitcoin.boot() - apps.misc.boot() - if not utils.BITCOIN_ONLY: - apps.ethereum.boot() - apps.lisk.boot() - apps.monero.boot() - apps.nem.boot() - apps.stellar.boot() - apps.ripple.boot() - apps.cardano.boot() - apps.tezos.boot() - apps.eos.boot() - apps.binance.boot() - apps.webauthn.boot() - if __debug__: - apps.debug.boot() - - # run main event loop and specify which screen is the default - apps.base.set_homescreen() - workflow.start_default() - - -_boot_apps() - -# initialize the wire codec -wire.setup(usb.iface_wire) -if __debug__: - wire.setup(usb.iface_debug, is_debug_session=True) - -loop.run() - -# loop is empty. That should not happen -utils.halt("All tasks have died.") +while True: + with unimport_manager: + import session # noqa: F401 + del session diff --git a/core/src/session.py b/core/src/session.py new file mode 100644 index 000000000..1fbf120c3 --- /dev/null +++ b/core/src/session.py @@ -0,0 +1,28 @@ +from trezor import loop, utils, wire, workflow + +# load applications +import apps.base +import usb + +apps.base.boot() + +if not utils.BITCOIN_ONLY and usb.ENABLE_IFACE_WEBAUTHN: + import apps.webauthn + + apps.webauthn.boot() + +if __debug__: + import apps.debug + + apps.debug.boot() + +# run main event loop and specify which screen is the default +apps.base.set_homescreen() +workflow.start_default() + +# initialize the wire codec +wire.setup(usb.iface_wire) +if __debug__: + wire.setup(usb.iface_debug, is_debug_session=True) + +loop.run() diff --git a/core/src/trezor/utils.py b/core/src/trezor/utils.py index 491300863..ce8f2865a 100644 --- a/core/src/trezor/utils.py +++ b/core/src/trezor/utils.py @@ -27,20 +27,20 @@ if __debug__: if False: from typing import ( Any, - Iterable, Iterator, Protocol, Union, TypeVar, Sequence, + Set, ) -def unimport_begin() -> Iterable[str]: +def unimport_begin() -> Set[str]: return set(sys.modules) -def unimport_end(mods: Iterable[str]) -> None: +def unimport_end(mods: Set[str], collect: bool = True) -> None: for mod in sys.modules: if mod not in mods: # remove reference from sys.modules @@ -58,7 +58,23 @@ def unimport_end(mods: Iterable[str]) -> None: # referenced from the parent package. both is fine. pass # collect removed modules - gc.collect() + if collect: + gc.collect() + + +class unimport: + def __init__(self) -> None: + self.mods: Set[str] | None = None + + def __enter__(self) -> None: + self.mods = unimport_begin() + + 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() def ensure(cond: bool, msg: str | None = None) -> None: diff --git a/core/src/trezor/wire/__init__.py b/core/src/trezor/wire/__init__.py index 35db4fa0a..4d73d1e2a 100644 --- a/core/src/trezor/wire/__init__.py +++ b/core/src/trezor/wire/__init__.py @@ -284,7 +284,7 @@ class UnexpectedMessageError(Exception): async def handle_session( - iface: WireInterface, session_id: int, is_debug_session: bool = True + iface: WireInterface, session_id: int, is_debug_session: bool = False ) -> None: if __debug__ and is_debug_session: ctx_buffer = WIRE_BUFFER_DEBUG @@ -295,6 +295,8 @@ async def handle_session( res_msg: protobuf.MessageType | None = None req_type = None req_msg = None + + modules = utils.unimport_begin() while True: try: if next_msg is None: @@ -330,7 +332,8 @@ async def handle_session( # Take a mark of modules that are imported at this point, so we can # roll back and un-import any others. Should not raise. - modules = utils.unimport_begin() + if is_debug_session: + modules = utils.unimport_begin() # We need to find a handler for this message type. Should not # raise. @@ -430,6 +433,10 @@ async def handle_session( # Unload modules imported by the workflow. Should not raise. utils.unimport_end(modules) + if not is_debug_session and next_msg is None: # and msg_type != 0: + loop.clear() + return + except Exception as exc: # The session handling should never exit, just log and continue. if __debug__: