From 34ae443d82b71b2be9518ab4e5b5b2a1e6d3e200 Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Wed, 23 Oct 2024 21:46:56 +0200
Subject: [PATCH] wip
---
core/src/trezor/utils.py | 21 +++++++++++++++++++
core/src/trezor/wire/__init__.py | 2 +-
.../wire/thp/received_message_handler.py | 14 ++++++++++++-
3 files changed, 35 insertions(+), 2 deletions(-)
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