THP buffers optimization, unit test improvements#6536
Conversation
WalkthroughReassembler.handle_packet now checks for a None read buffer and returns False when buffer retrieval fails. ThpBuffer.get was changed to return memoryview | None and logs a debug warning instead of asserting when requested length exceeds available data. Channel.write now pre-checks the encoded message length against MAX_PAYLOAD_LEN and routes oversized messages to be sent as a Failure(FirmwareError); it also catches MemoryError during write-buffer allocation, logging and writing a FirmwareError Failure unless the outgoing message is already a Failure. CHECKSUM_LENGTH constant was removed. Test infra: global DISABLE_LOG/TestRunner added, TrackedChannel and _encrypt_patch added for write-tracking, and many tests removed ad-hoc debug-log suppression; a new test module exercises oversized-message behavior. Sequence Diagram(s)sequenceDiagram
participant App as App
participant Channel as Channel
participant Writer as Writer
participant MemoryManager as MemoryManager
App->>Channel: write(msg)
Channel->>Writer: encode(msg) -> encoded_len, encoded_bytes
alt encoded_len > MAX_PAYLOAD_LEN
Channel->>Channel: log warning (oversized)
Channel->>Channel: create Failure(FirmwareError)
Channel->>Channel: write(Failure) -- recursive (guarded)
else encoded_len <= MAX_PAYLOAD_LEN
Channel->>MemoryManager: get(encoded_len)
MemoryManager-->>Channel: buffer or None
alt buffer is None
Channel->>Channel: log warning (allocation/retrieval failed)
Channel-->>App: return False
else buffer available
Channel->>Writer: fill buffer and send packet(s)
Channel-->>App: return True
end
end
sequenceDiagram
participant Test as Test
participant Tracked as TrackedChannel
participant Channel as Channel
participant Writer as Writer
Test->>Tracked: set_expected_messages([...])
Test->>Tracked: __enter__()
Test->>Tracked: write(msg)
Tracked->>Tracked: record(msg type)
Tracked->>Channel: write(msg) (encryption patched via _encrypt_patch)
Channel->>Writer: encode and attempt send
Test->>Tracked: __exit__()
Tracked->>Tracked: validate recorded messages (assert expected)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
| model | device_test | click_test | persistence_test |
|---|---|---|---|
| T2T1 | test(all) main(all) ![]() |
test(all) main(all) ![]() |
test(all) main(all) ![]() |
| T3B1 | test(all) main(all) ![]() |
test(all) main(all) ![]() |
test(all) main(all) ![]() |
| T3T1 | test(all) main(all) ![]() |
test(all) main(all) ![]() |
test(all) main(all) ![]() |
| T3W1 | test(all) main(all) ![]() |
test(all) main(all) ![]() |
test(all) main(all) ![]() |
Latest CI run: 22572485050
[no changelog]
c518b19 to
a1756b6
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
core/tests/test_trezor.wire.thp.channel.py (1)
42-80:xtest_*methods won’t be executed by discovery.If these are intentionally disabled, prefer
@unittest.skip(...)ontest_*methods so they remain visible in test reports and easier to re-enable.💡 Suggested pattern
- def xtest_too_big_message_mocked_size(self): + `@unittest.skip`("TODO: stabilize and re-enable") + def test_too_big_message_mocked_size(self): @@ - def xtest_big_message_mocked_size_pass(self): + `@unittest.skip`("TODO: stabilize and re-enable") + def test_big_message_mocked_size_pass(self): @@ - def xtest_big_message(self): + `@unittest.skip`("TODO: stabilize and re-enable") + def test_big_message(self):Also applies to: 107-124
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/tests/test_trezor.wire.thp.channel.py` around lines 42 - 80, The two disabled tests xtest_too_big_message_mocked_size and xtest_big_message_mocked_size_pass are not discovered by unittest; change them to regular test names (test_too_big_message_mocked_size, test_big_message_mocked_size_pass) and mark them with `@unittest.skip`("reason") if you intend to keep them disabled but visible in test reports, or remove the skip to re-enable; update the method names and add the unittest.skip decorator above the corresponding functions (or use unittest.skipIf/unittest.skipUnless as appropriate) so they appear in discovery and test output.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 423-432: The write() fallback path can recurse if sending
Failure(code=FailureType.FirmwareError) itself fails; add a one-shot guard
(e.g., a boolean flag like _sending_fallback or an extra parameter
allow_fallback=False) inside the Channel.write method to detect when a fallback
Failure is being sent and prevent re-entering the fallback path: set the guard
before invoking write(Failure(...), session_id), clear it after, and if the
guard is already set then short-circuit (return/raise) instead of attempting
another fallback. Apply the same guard to the other fallback site that
constructs Failure(FirmwareError) (the block around the 446-449 occurrence) so
recursive loops are impossible.
In `@core/tests/README.md`:
- Around line 34-35: The README section beginning with the "## Logging" heading
has incorrect file-ending style; normalize the file to use LF line endings and
ensure the file ends with a single final newline character so
editorconfig-checker passes (replace CRLF with LF and add a trailing '\n' at
EOF).
In `@core/tests/thp_common.py`:
- Around line 37-39: The shared _encrypt_patch defined for Channel._encrypt is
non-reentrant and can leak a monkey-patched state when TrackedChannel.write is
entered recursively; replace the module-level/shared _encrypt_patch with code
that constructs and applies a fresh patch instance each time it's needed (e.g.,
inside the context where TrackedChannel.write is invoked) so orig_value isn't
overwritten and Channel._encrypt is always restored; update the similar shared
patch usage referenced around the other occurrence (lines corresponding to the
second instance) to follow the same per-use patch creation approach.
---
Nitpick comments:
In `@core/tests/test_trezor.wire.thp.channel.py`:
- Around line 42-80: The two disabled tests xtest_too_big_message_mocked_size
and xtest_big_message_mocked_size_pass are not discovered by unittest; change
them to regular test names (test_too_big_message_mocked_size,
test_big_message_mocked_size_pass) and mark them with `@unittest.skip`("reason")
if you intend to keep them disabled but visible in test reports, or remove the
skip to re-enable; update the method names and add the unittest.skip decorator
above the corresponding functions (or use unittest.skipIf/unittest.skipUnless as
appropriate) so they appear in discovery and test output.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
core/src/trezor/wire/thp/channel.pycore/src/trezor/wire/thp/memory_manager.pycore/src/trezor/wire/thp/writer.pycore/tests/README.mdcore/tests/fixtures.pycore/tests/test_apps.bitcoin.approver.pycore/tests/test_apps.bitcoin.authorization.pycore/tests/test_apps.bitcoin.keychain.pycore/tests/test_apps.common.keychain.pycore/tests/test_apps.ethereum.keychain.pycore/tests/test_storage.cache.pycore/tests/test_trezor.wire.thp.channel.pycore/tests/test_trezor.wire.thp.crypto.pycore/tests/test_trezor.wire.thp.writer.pycore/tests/thp_common.pycore/tests/unittest.py
💤 Files with no reviewable changes (9)
- core/tests/test_apps.ethereum.keychain.py
- core/tests/test_trezor.wire.thp.writer.py
- core/tests/test_apps.common.keychain.py
- core/tests/test_apps.bitcoin.approver.py
- core/tests/test_storage.cache.py
- core/tests/test_trezor.wire.thp.crypto.py
- core/tests/test_apps.bitcoin.keychain.py
- core/tests/test_apps.bitcoin.authorization.py
- core/src/trezor/wire/thp/writer.py
core/tests/README.md
Outdated
| ## Logging | ||
| By default, logging in unittests (using `trezor.log`) is disabled _even_ on debug builds. To enable logging, set `DISABLE_LOG` in `unittest.py` to `False`. No newline at end of file |
There was a problem hiding this comment.
Fix file-ending style to unblock CI.
Prebuild is failing with editorconfig-checker (wrong line endings or missing final newline). Please normalize to LF and ensure the file ends with a newline.
💡 Minimal fix
## Logging
By default, logging in unittests (using `trezor.log`) is disabled _even_ on debug builds. To enable logging, set `DISABLE_LOG` in `unittest.py` to `False`.
+📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| ## Logging | |
| By default, logging in unittests (using `trezor.log`) is disabled _even_ on debug builds. To enable logging, set `DISABLE_LOG` in `unittest.py` to `False`. | |
| ## Logging | |
| By default, logging in unittests (using `trezor.log`) is disabled _even_ on debug builds. To enable logging, set `DISABLE_LOG` in `unittest.py` to `False`. | |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@core/tests/README.md` around lines 34 - 35, The README section beginning with
the "## Logging" heading has incorrect file-ending style; normalize the file to
use LF line endings and ensure the file ends with a single final newline
character so editorconfig-checker passes (replace CRLF with LF and add a
trailing '\n' at EOF).
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
core/tests/thp_common.py (1)
37-39:⚠️ Potential issue | 🟠 MajorUse a fresh patch context per write to avoid non-reentrant monkey-patch state.
The shared
_encrypt_patchobject is still reused across recursive writes, which can leakChannel._encryptpatch state.✅ Reentrant-safe fix
- _encrypt_patch = patch( - Channel, "_encrypt", lambda self, buffer, noise_payload_len: None - ) + def _no_op_encrypt(_self, _buffer, _noise_payload_len) -> None: + return None @@ async def write( self, msg: protobuf.MessageType, session_id: int = 0, ) -> None: self.messages_to_write.append(msg) - with _encrypt_patch: + with patch(Channel, "_encrypt", _no_op_encrypt): # encryption is disabled return await self.original_write(msg, session_id)Also applies to: 81-83
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/tests/thp_common.py` around lines 37 - 39, The shared _encrypt_patch patch object is reused across recursive writes causing non-reentrant monkey-patch leaks; instead create and use a fresh patch context each time you need to stub Channel._encrypt (e.g., replace the module-level _encrypt_patch with a helper that returns a new patch: call patch.object(Channel, "_encrypt", lambda self, buffer, noise_payload_len: None) inside a with/context manager or factory function and use it around each write), and apply the same change to the other occurrence (the block at the second occurrence previously at lines 81-83) so each write uses its own patch context.core/src/trezor/wire/thp/channel.py (1)
423-432:⚠️ Potential issue | 🟠 MajorAdd a one-shot guard to prevent fallback recursion loops.
The fallback
write(Failure(FirmwareError))still re-enterswrite()without a recursion guard. If fallback send also fails, this can recurse repeatedly.🔧 Minimal guard pattern
async def write( self, msg: protobuf.MessageType, session_id: int = 0, ) -> None: assert ABP.is_sending_allowed(self.channel_cache) + from trezor.enums import FailureType + from trezor.messages import Failure + + is_firmware_failure = ( + Failure.is_type_of(msg) and msg.code == FailureType.FirmwareError + ) @@ if length > MAX_PAYLOAD_LEN: if __debug__: log.warning(__name__, "Failed to write, message is too big.") - from trezor.enums import FailureType - from trezor.messages import Failure + if is_firmware_failure: + return return await self.write( Failure(code=FailureType.FirmwareError), session_id, ) @@ except MemoryError: if __debug__: log.warning( __name__, "Failed to allocate a sufficienty large write buffer." ) - from trezor.enums import FailureType - from trezor.messages import Failure + if is_firmware_failure: + return return await self.write( Failure(code=FailureType.FirmwareError), session_id, )Also applies to: 446-449
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/src/trezor/wire/thp/channel.py` around lines 423 - 432, The fallback call in Channel.write that does "return await self.write(Failure(...), session_id)" can re-enter write and recurse if the fallback send fails; add a one‑shot recursion guard (e.g., an instance flag like self._sending_failure or an extra parameter fallback=False) and check it at the start of write to short‑circuit further fallback attempts; when invoking the fallback set the guard to True before awaiting the secondary write and clear it after, and if the guard is already set simply return/raise without re-entering; apply the same guard pattern to the other identical fallback block (the block around the other Failure(FirmwareError) call) so fallback recursion cannot loop.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 435-438: The dynamic-allocation path currently creates a
bytearray(buffer = bytearray(length)) and later slices it, causing an extra
large copy right before calling write_encrypted_payload(); change the logic to
keep and pass a memoryview (e.g., mv = memoryview(buffer) or allocate as
bytearray and immediately wrap with memoryview) so slicing uses the memoryview
instead of creating a new bytearray; update both dynamic-allocation branches
that set buffer to bytearray(length) and any subsequent slices used in
write_encrypted_payload(...) to operate on the memoryview to avoid the extra
copy.
---
Duplicate comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 423-432: The fallback call in Channel.write that does "return
await self.write(Failure(...), session_id)" can re-enter write and recurse if
the fallback send fails; add a one‑shot recursion guard (e.g., an instance flag
like self._sending_failure or an extra parameter fallback=False) and check it at
the start of write to short‑circuit further fallback attempts; when invoking the
fallback set the guard to True before awaiting the secondary write and clear it
after, and if the guard is already set simply return/raise without re-entering;
apply the same guard pattern to the other identical fallback block (the block
around the other Failure(FirmwareError) call) so fallback recursion cannot loop.
In `@core/tests/thp_common.py`:
- Around line 37-39: The shared _encrypt_patch patch object is reused across
recursive writes causing non-reentrant monkey-patch leaks; instead create and
use a fresh patch context each time you need to stub Channel._encrypt (e.g.,
replace the module-level _encrypt_patch with a helper that returns a new patch:
call patch.object(Channel, "_encrypt", lambda self, buffer, noise_payload_len:
None) inside a with/context manager or factory function and use it around each
write), and apply the same change to the other occurrence (the block at the
second occurrence previously at lines 81-83) so each write uses its own patch
context.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
core/src/trezor/wire/thp/channel.pycore/src/trezor/wire/thp/memory_manager.pycore/src/trezor/wire/thp/writer.pycore/tests/README.mdcore/tests/fixtures.pycore/tests/test_apps.bitcoin.approver.pycore/tests/test_apps.bitcoin.authorization.pycore/tests/test_apps.bitcoin.keychain.pycore/tests/test_apps.common.keychain.pycore/tests/test_apps.ethereum.keychain.pycore/tests/test_storage.cache.pycore/tests/test_trezor.wire.thp.channel.pycore/tests/test_trezor.wire.thp.crypto.pycore/tests/test_trezor.wire.thp.writer.pycore/tests/thp_common.pycore/tests/unittest.py
💤 Files with no reviewable changes (9)
- core/tests/test_apps.common.keychain.py
- core/src/trezor/wire/thp/writer.py
- core/tests/test_apps.bitcoin.authorization.py
- core/tests/test_trezor.wire.thp.writer.py
- core/tests/test_apps.bitcoin.keychain.py
- core/tests/test_storage.cache.py
- core/tests/test_apps.bitcoin.approver.py
- core/tests/test_trezor.wire.thp.crypto.py
- core/tests/test_apps.ethereum.keychain.py
🚧 Files skipped from review as they are similar to previous changes (2)
- core/tests/fixtures.py
- core/tests/test_trezor.wire.thp.channel.py
a1756b6 to
5ef58f6
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
core/src/trezor/wire/thp/channel.py (1)
423-432:⚠️ Potential issue | 🟡 MinorMissing recursion guard on oversized message path.
The MemoryError fallback path (line 446) correctly guards against recursion by checking
Failure.is_type_of(msg), but the oversized-message path here lacks the same guard. While aFailuremessage is unlikely to exceedMAX_PAYLOAD_LEN, adding the same guard ensures consistent defensive coding.🛡️ Proposed fix for consistency
if length > MAX_PAYLOAD_LEN: if __debug__: log.warning(__name__, "Failed to write, message is too big.") from trezor.enums import FailureType from trezor.messages import Failure + if Failure.is_type_of(msg): + # prevent infinite recursion + return + return await self.write( Failure(code=FailureType.FirmwareError), session_id, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/src/trezor/wire/thp/channel.py` around lines 423 - 432, The oversized-message branch in Channel.write (the if length > MAX_PAYLOAD_LEN block) needs the same recursion guard used in the MemoryError fallback: before returning await self.write(Failure(...), session_id) check whether the original msg is already a Failure by calling Failure.is_type_of(msg) (or equivalent) and if so avoid recursing; implement the guard inside the oversized path so you only synthesize and resend a Failure when the original message is not already a Failure to prevent infinite recursion.
🧹 Nitpick comments (3)
core/tests/test_trezor.wire.thp.channel.py (1)
3-3: Consider movingCallableimport underTYPE_CHECKING.Per MicroPython guidelines, types used only for annotations should be conditionally imported under
TYPE_CHECKINGto avoid runtime overhead.♻️ Proposed fix
# flake8: noqa: F403,F405 from common import * # isort:skip -from typing import Callable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Callable from mock import patchBased on learnings: "In MicroPython (runtime ignores annotations), avoid importing types solely for type checking at runtime."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/tests/test_trezor.wire.thp.channel.py` at line 3, Move the runtime import of Callable into a conditional TYPE_CHECKING block to avoid importing typing at runtime: add a from typing import TYPE_CHECKING at top and wrap from typing import Callable inside if TYPE_CHECKING: so that Callable is only imported for type-checking (references to Callable in functions or annotations such as in the test_trezor.wire.thp.channel.py file remain unchanged).core/tests/thp_common.py (1)
40-82: Consider restoringinner_channel.writeon cleanup.
TrackedChannel.__init__patchesinner_channel.writebut there's no cleanup to restoreoriginal_write. If the channel is reused across tests (unlikely givensetUprecreates context), the patched write would persist.♻️ Proposed fix: add a cleanup method or restore in __exit__
def __exit__(self, exc_type, exc_value, tb) -> None: + # Restore original write method + self.inner_channel.write = self.original_write + if not self.expected_messages_to_write: return🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/tests/thp_common.py` around lines 40 - 82, TrackedChannel replaces inner_channel.write in __init__ but never restores it; modify TrackedChannel to restore inner_channel.write to original_write during cleanup (e.g., in __exit__ or a dedicated close/cleanup method) so the patched write does not persist across contexts or tests; ensure restoration happens even on exceptions (use finally or always-run __exit__ logic) and reference inner_channel, original_write, __init__, and __exit__ when implementing the fix.core/src/trezor/wire/thp/channel.py (1)
446-448: Chain the exception for clearer tracebacks.Per static analysis (Ruff B904), when re-raising within an
exceptclause, useraise ... fromto preserve the exception chain.♻️ Proposed fix
if Failure.is_type_of(msg): # prevent infinite recursion - raise MemoryError + raise MemoryError from None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/src/trezor/wire/thp/channel.py` around lines 446 - 448, Inside the except handler where Failure.is_type_of(msg) is checked, change the bare re-raise "raise MemoryError" to chain the original caught exception using "raise MemoryError from <exc>" (use the same exception variable used in the surrounding except clause, e.g., "e"). Update the raise to "raise MemoryError from e" so Failure.is_type_of(msg) and MemoryError preserve the exception chain for better tracebacks.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 439-442: Fix the typo in the warning message emitted in
channel.py: update the string passed to log.warning in the block guarded by
__debug__ (the warning logging call in the code around the write buffer
allocation logic) to replace "sufficienty" with "sufficiently" so the message
reads "...Failed to allocate a sufficiently large write buffer." and keep the
rest of the call (log.warning(__name__, ...)) unchanged.
In `@core/tests/test_trezor.wire.thp.channel.py`:
- Line 91: The test contains a typo in the ping_message string; update the
ping_message variable used in the test (ping_message) to a grammatical phrase
such as "This message will fail to be sent" (or "This message should fail to be
sent") so the test message reads correctly.
---
Duplicate comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 423-432: The oversized-message branch in Channel.write (the if
length > MAX_PAYLOAD_LEN block) needs the same recursion guard used in the
MemoryError fallback: before returning await self.write(Failure(...),
session_id) check whether the original msg is already a Failure by calling
Failure.is_type_of(msg) (or equivalent) and if so avoid recursing; implement the
guard inside the oversized path so you only synthesize and resend a Failure when
the original message is not already a Failure to prevent infinite recursion.
---
Nitpick comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 446-448: Inside the except handler where Failure.is_type_of(msg)
is checked, change the bare re-raise "raise MemoryError" to chain the original
caught exception using "raise MemoryError from <exc>" (use the same exception
variable used in the surrounding except clause, e.g., "e"). Update the raise to
"raise MemoryError from e" so Failure.is_type_of(msg) and MemoryError preserve
the exception chain for better tracebacks.
In `@core/tests/test_trezor.wire.thp.channel.py`:
- Line 3: Move the runtime import of Callable into a conditional TYPE_CHECKING
block to avoid importing typing at runtime: add a from typing import
TYPE_CHECKING at top and wrap from typing import Callable inside if
TYPE_CHECKING: so that Callable is only imported for type-checking (references
to Callable in functions or annotations such as in the
test_trezor.wire.thp.channel.py file remain unchanged).
In `@core/tests/thp_common.py`:
- Around line 40-82: TrackedChannel replaces inner_channel.write in __init__ but
never restores it; modify TrackedChannel to restore inner_channel.write to
original_write during cleanup (e.g., in __exit__ or a dedicated close/cleanup
method) so the patched write does not persist across contexts or tests; ensure
restoration happens even on exceptions (use finally or always-run __exit__
logic) and reference inner_channel, original_write, __init__, and __exit__ when
implementing the fix.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
core/src/trezor/wire/thp/channel.pycore/src/trezor/wire/thp/memory_manager.pycore/src/trezor/wire/thp/writer.pycore/tests/fixtures.pycore/tests/test_trezor.wire.thp.channel.pycore/tests/thp_common.py
💤 Files with no reviewable changes (1)
- core/src/trezor/wire/thp/writer.py
🚧 Files skipped from review as they are similar to previous changes (1)
- core/tests/fixtures.py
5ef58f6 to
9bcffd8
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
core/tests/thp_common.py (1)
42-49: Scopeinner_channel.writemonkey-patch to the context manager lifecycle.
TrackedChannelpatcheswritein__init__and never restores it. Prefer patching in__enter__and restoring in__exit__so tracking is contained towith channel:.♻️ Proposed refactor
class TrackedChannel: @@ def __init__(self) -> None: self.messages_to_write = [] self.expected_messages_to_write = [] self.inner_channel = context.get_channel_context() self.original_write = self.inner_channel.write - # To track recursive calls in channel.write - self.inner_channel.write = self.write @@ - def __enter__(self) -> None: + def __enter__(self) -> None: self.messages_to_write: list[protobuf.MessageType] = [] + # To track recursive calls in channel.write + self.inner_channel.write = self.write def __exit__(self, exc_type, exc_value, tb) -> None: + self.inner_channel.write = self.original_write if not self.expected_messages_to_write: returnAlso applies to: 61-63
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/tests/thp_common.py` around lines 42 - 49, The TrackedChannel class currently monkey-patches inner_channel.write in __init__ and never restores it; change the patch to occur in the context manager lifecycle by moving the assignment self.original_write = self.inner_channel.write and self.inner_channel.write = self.write into __enter__, and restore the original with self.inner_channel.write = self.original_write inside __exit__ (also clear or null out originals if needed); update any analogous patching at the other location (the block around lines 61-63) to follow the same __enter__/__exit__ pattern so the write override is only active for the with-context.core/src/trezor/wire/thp/channel.py (1)
438-449: Use explicit exception chaining for clarity.At line 449, consider using
raise MemoryError from errto make the relationship between the original and new exception explicit. While Python's implicit context chaining preserves the original exception, explicit chaining documents the causal intent more clearly in code that intentionally re-raises to prevent infinite recursion.🔧 Proposed fix
- except MemoryError: + except MemoryError as err: if __debug__: log.warning( __name__, "Failed to allocate a sufficiently large write buffer.", ) from trezor.enums import FailureType from trezor.messages import Failure if Failure.is_type_of(msg): # prevent infinite recursion - raise MemoryError + raise MemoryError from err🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/src/trezor/wire/thp/channel.py` around lines 438 - 449, The except block currently swallows the original MemoryError; change the handler to capture the original exception (use "except MemoryError as err") and re-raise with explicit chaining ("raise MemoryError from err") instead of plain "raise MemoryError" in the branch that prevents recursion (the block around Failure.is_type_of(msg)); update only the except clause and the re-raise so the imports and Failure.is_type_of(msg) logic remain unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/tests/test_trezor.wire.thp.channel.py`:
- Line 10: Add postponed evaluation of annotations to avoid importing protobuf
at module import time when utils.USE_THP is false: insert "from __future__
import annotations" at the top of the module (before any other code) so
annotations in _encoded_len_patch (and any references to protobuf.MessageType)
are not evaluated immediately; alternatively, change the type annotations in
_encoded_len_patch to use quoted forward references (e.g.,
"protobuf.MessageType") or guard imports with typing.TYPE_CHECKING and only
import protobuf inside that block, ensuring references to protobuf.MessageType
no longer trigger NameError when THP is disabled.
---
Nitpick comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 438-449: The except block currently swallows the original
MemoryError; change the handler to capture the original exception (use "except
MemoryError as err") and re-raise with explicit chaining ("raise MemoryError
from err") instead of plain "raise MemoryError" in the branch that prevents
recursion (the block around Failure.is_type_of(msg)); update only the except
clause and the re-raise so the imports and Failure.is_type_of(msg) logic remain
unchanged.
In `@core/tests/thp_common.py`:
- Around line 42-49: The TrackedChannel class currently monkey-patches
inner_channel.write in __init__ and never restores it; change the patch to occur
in the context manager lifecycle by moving the assignment self.original_write =
self.inner_channel.write and self.inner_channel.write = self.write into
__enter__, and restore the original with self.inner_channel.write =
self.original_write inside __exit__ (also clear or null out originals if
needed); update any analogous patching at the other location (the block around
lines 61-63) to follow the same __enter__/__exit__ pattern so the write override
is only active for the with-context.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
core/src/trezor/wire/thp/channel.pycore/src/trezor/wire/thp/memory_manager.pycore/src/trezor/wire/thp/writer.pycore/tests/fixtures.pycore/tests/test_trezor.wire.thp.channel.pycore/tests/thp_common.py
💤 Files with no reviewable changes (1)
- core/src/trezor/wire/thp/writer.py
🚧 Files skipped from review as they are similar to previous changes (2)
- core/src/trezor/wire/thp/memory_manager.py
- core/tests/fixtures.py
- Allows creation of a new (bigger) buffer when message exceeds standard buffer size. - Fail to read big messages in more controlled way. - Remove duplicated constant from `writer.py`. - Mostly solves the issue with listing of FIDO credentials. Now it is possible to list almost 60kB of credentials compared to previous 8kB. (#6349) [no changelog]
[no changelog]
9bcffd8 to
a1810ad
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/tests/thp_common.py (1)
56-60: Returnselffrom__enter__to comply with context manager protocol.The
TrackedChannelcontext manager currently returnsNonefrom__enter__, violating Python's context manager protocol (PEP 343). While the existing tests usewith channel:without anasbinding and thus avoid the issue, the code should follow standard practice. If anyone attemptswith channel as tracked:, the variable would bind toNoneinstead of the channel instance.♻️ Proposed fix
- def __enter__(self) -> None: + def __enter__(self) -> "TrackedChannel": self.messages_to_write: list[protobuf.MessageType] = [] # To track recursive calls in channel.write self.inner_channel.write = self.write + return self🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/tests/thp_common.py` around lines 56 - 60, The TrackedChannel context manager's __enter__ currently returns None; update the TrackedChannel.__enter__ method to return self so that "with channel as tracked:" binds to the TrackedChannel instance; locate the __enter__ implementation (it sets self.messages_to_write and replaces self.inner_channel.write) and add a final "return self" so the context manager conforms to PEP 343.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@core/tests/thp_common.py`:
- Around line 56-60: The TrackedChannel context manager's __enter__ currently
returns None; update the TrackedChannel.__enter__ method to return self so that
"with channel as tracked:" binds to the TrackedChannel instance; locate the
__enter__ implementation (it sets self.messages_to_write and replaces
self.inner_channel.write) and add a final "return self" so the context manager
conforms to PEP 343.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
core/src/trezor/wire/thp/channel.pycore/src/trezor/wire/thp/memory_manager.pycore/src/trezor/wire/thp/writer.pycore/tests/fixtures.pycore/tests/test_trezor.wire.thp.channel.pycore/tests/thp_common.py
💤 Files with no reviewable changes (1)
- core/src/trezor/wire/thp/writer.py
🚧 Files skipped from review as they are similar to previous changes (2)
- core/tests/test_trezor.wire.thp.channel.py
- core/tests/fixtures.py
core/tests/unittest.py
Outdated
| if __debug__ and DISABLE_LOG: | ||
| from trezor import log, utils | ||
|
|
||
| log.debug = log._no_op |
core/src/trezor/wire/thp/channel.py
Outdated
|
|
||
| if Failure.is_type_of(msg): | ||
| # prevent infinite recursion | ||
| raise MemoryError from e |
There was a problem hiding this comment.
| raise MemoryError from e | |
| raise |
| return await self.write( | ||
| Failure(code=FailureType.FirmwareError), | ||
| session_id, | ||
| ) |
There was a problem hiding this comment.
IIUC, we should still raise an exception - otherwise the current workflow will continue.
(here and below)
There was a problem hiding this comment.
What exception type would you like to raise here?
I suggest the following change:
From:
return await self.write(
Failure(code=FailureType.FirmwareError),
session_id,
)
to
await self.write(
Failure(code=FailureType.FirmwareError),
session_id,
)
raise #SomeTypeOfException
WDYT?
[no changelog]
There was a problem hiding this comment.
♻️ Duplicate comments (1)
core/src/trezor/wire/thp/channel.py (1)
423-432:⚠️ Potential issue | 🟡 MinorAdd a one-shot guard in the oversized fallback path too.
Line 423 recursively calls
write(Failure(...))without guarding the case where the current message is already aFailure. This is the same recursion class previously discussed; please short-circuit before re-entering fallback.🔧 Suggested guard
if length > MAX_PAYLOAD_LEN: if __debug__: log.warning(__name__, "Failed to write, message is too big.") from trezor.enums import FailureType from trezor.messages import Failure + if Failure.is_type_of(msg): + raise ValueError("Cannot write oversized Failure message.") return await self.write( Failure(code=FailureType.FirmwareError), session_id, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/src/trezor/wire/thp/channel.py` around lines 423 - 432, The oversized-payload fallback in write(...) can recurse when the current message is already a Failure; add a one-shot guard: inside the length > MAX_PAYLOAD_LEN branch, detect if the message is an instance of trezor.messages.Failure (or has code == FailureType.FirmwareError) and short-circuit (e.g., return or raise) instead of calling self.write(Failure(...)) again; modify the code in the write method (the oversized fallback) to check Failure before constructing/recursing into Failure and handle that case directly to avoid infinite recursion.
🧹 Nitpick comments (1)
core/tests/unittest.py (1)
10-10: Consider making test log suppression configurable.Hardcoding
DISABLE_LOG = Truemakes local debugging noisier to re-enable. A small runtime switch (e.g., env/config) would keep defaults while allowing quick overrides.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/tests/unittest.py` at line 10, Replace the hardcoded constant DISABLE_LOG = True with a runtime-configurable setting by reading an environment variable (e.g., os.getenv('DISABLE_LOG')) or a test config, parsing common truthy/falsey values, and falling back to the current default (True) if unset; update the symbol DISABLE_LOG in core/tests/unittest.py to derive its value from that env/config so developers can temporarily override logging without editing code.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@core/src/trezor/wire/thp/channel.py`:
- Around line 423-432: The oversized-payload fallback in write(...) can recurse
when the current message is already a Failure; add a one-shot guard: inside the
length > MAX_PAYLOAD_LEN branch, detect if the message is an instance of
trezor.messages.Failure (or has code == FailureType.FirmwareError) and
short-circuit (e.g., return or raise) instead of calling
self.write(Failure(...)) again; modify the code in the write method (the
oversized fallback) to check Failure before constructing/recursing into Failure
and handle that case directly to avoid infinite recursion.
---
Nitpick comments:
In `@core/tests/unittest.py`:
- Line 10: Replace the hardcoded constant DISABLE_LOG = True with a
runtime-configurable setting by reading an environment variable (e.g.,
os.getenv('DISABLE_LOG')) or a test config, parsing common truthy/falsey values,
and falling back to the current default (True) if unset; update the symbol
DISABLE_LOG in core/tests/unittest.py to derive its value from that env/config
so developers can temporarily override logging without editing code.




































This PR:
- Fixes broken suppression of logging in core python unit tests
- Allows creation of a new (bigger) buffer when message exceeds standard buffer size.
- Fail to read big messages in more controlled way.
- Remove duplicated constant from
writer.py.- Mostly solves the issue with listing of FIDO credentials. Now it is possible to list almost 60kB of credentials compared to previous 8kB. (#6349)
- Introduces new unittests for the channel write function.