diff --git a/core/src/trezor/utils.py b/core/src/trezor/utils.py index ff977d27d7..e53d68c8f2 100644 --- a/core/src/trezor/utils.py +++ b/core/src/trezor/utils.py @@ -39,6 +39,27 @@ DISABLE_ENCRYPTION: bool = False ALLOW_DEBUG_MESSAGES: bool = False +last_allocation_count: int = 0 + + +def print_and_update_alloc(name: str = "unknown") -> None: + try: + import micropython + + global last_allocation_count + current_count = micropython.alloc_count() + print("------------------------") + micropython.mem_info() + print("### Called from: ", name) + print("### Allocation count:", current_count) + print("### Allocation diff: ", current_count - last_allocation_count) + print("------------------------") + + last_allocation_count = current_count + except Exception: + print("Memory profiling failed") + + if __debug__: if EMULATOR: import uos diff --git a/core/src/trezor/wire/__init__.py b/core/src/trezor/wire/__init__.py index 8cc1e8bf09..23207ffa3a 100644 --- a/core/src/trezor/wire/__init__.py +++ b/core/src/trezor/wire/__init__.py @@ -60,7 +60,7 @@ def setup(iface: WireInterface) -> None: if utils.USE_THP: async def handle_thp_session(iface: WireInterface): - + utils.print_and_update_alloc(__name__) thp_main.set_read_buffer(WIRE_BUFFER) thp_main.set_write_buffer(WIRE_BUFFER_2) diff --git a/core/src/trezor/wire/thp/received_message_handler.py b/core/src/trezor/wire/thp/received_message_handler.py index 77f80fbb55..a0cdc62818 100644 --- a/core/src/trezor/wire/thp/received_message_handler.py +++ b/core/src/trezor/wire/thp/received_message_handler.py @@ -65,6 +65,7 @@ async def handle_received_message( ctx: Channel, message_buffer: utils.BufferType ) -> None: """Handle a message received from the channel.""" + utils.print_and_update_alloc(__name__) if __debug__ and utils.ALLOW_DEBUG_MESSAGES: log.debug(__name__, "handle_received_message") @@ -99,7 +100,10 @@ async def handle_received_message( # 1: Handle ACKs if control_byte.is_ack(ctrl_byte): + utils.print_and_update_alloc("ACK before handling") await _handle_ack(ctx, ack_bit) + utils.print_and_update_alloc("ACK after handling") + return if _should_have_ctrl_byte_encrypted_transport( @@ -173,11 +177,19 @@ async def _handle_ack(ctx: Channel, ack_bit: int): log.debug(__name__, "Stopped transmission loop") ABP.set_sending_allowed(ctx.channel_cache, True) - + utils.print_and_update_alloc("Handle ACK, after cache edit") if ctx.write_task_spawn is not None: if __debug__ and utils.ALLOW_DEBUG_MESSAGES: log.debug(__name__, 'Control to "write_encrypted_payload_loop" task') await ctx.write_task_spawn + try: + ctx.write_task_spawn.close() + ctx.write_task_spawn = None + except Exception as e: + print("\nTHE KILLING FAILED, SO IT IS ONLY ATTEMPTED MURDER") + print(type(e)) + utils.print_and_update_alloc("Handle ACK, after write task spawn") + # Note that no the write_task_spawn could result in loop.clear(), # which will result in termination of this function - any code after # this await might not be executed