Fix imports in thp_v1.py

M1nd3r 1 month ago
parent 5761512ce7
commit 975f020d80

@ -1,7 +1,7 @@
from typing import TYPE_CHECKING # pyright: ignore[reportShadowedImports]
from trezor import utils
from trezor.wire import codec_v1, thp_v1
from trezor.wire import codec_v1
from trezor.wire.protocol_common import MessageWithId
if TYPE_CHECKING:
@ -10,13 +10,12 @@ if TYPE_CHECKING:
async def read_message(iface: WireInterface, buffer: utils.BufferType) -> MessageWithId:
if utils.USE_THP:
return await thp_v1.read_message(iface, buffer)
raise Exception("THP protocol should be used instead")
return await codec_v1.read_message(iface, buffer)
async def write_message(iface: WireInterface, message: MessageWithId) -> None:
if utils.USE_THP:
await thp_v1.write_message_with_sync_control(iface, message)
return
raise Exception("THP protocol should be used instead")
await codec_v1.write_message(iface, message.type, message.data)
return

@ -13,11 +13,9 @@ from . import thp_session as THP
from .checksum import CHECKSUM_LENGTH
from .thp_messages import (
ACK_MESSAGE,
CONT_DATA_OFFSET,
CONTINUATION_PACKET,
ENCRYPTED_TRANSPORT,
HANDSHAKE_INIT,
INIT_DATA_OFFSET,
)
from .thp_session import ThpError
@ -30,8 +28,12 @@ _MOCK_INTERFACE_HID = b"\x00"
_PUBKEY_LENGTH = const(32)
_REPORT_LENGTH = const(64)
_MAX_PAYLOAD_LEN = const(60000)
INIT_DATA_OFFSET = const(5)
CONT_DATA_OFFSET = const(3)
REPORT_LENGTH = const(64)
MAX_PAYLOAD_LEN = const(60000)
class ChannelContext(Context):
@ -267,7 +269,7 @@ def _encode_iface(iface: WireInterface) -> bytes:
def _get_buffer_for_payload(
payload_length: int, existing_buffer: utils.BufferType, max_length=_MAX_PAYLOAD_LEN
payload_length: int, existing_buffer: utils.BufferType, max_length=MAX_PAYLOAD_LEN
) -> utils.BufferType:
if payload_length > max_length:
raise ThpError("Message too large")
@ -276,7 +278,7 @@ def _get_buffer_for_payload(
try:
payload: utils.BufferType = bytearray(payload_length)
except MemoryError:
payload = bytearray(_REPORT_LENGTH)
payload = bytearray(REPORT_LENGTH)
raise ThpError("Message too large")
return payload

@ -12,8 +12,6 @@ HANDSHAKE_INIT = 0x00
ACK_MESSAGE = 0x20
_ERROR = 0x41
_CHANNEL_ALLOCATION_RES = 0x40
INIT_DATA_OFFSET = const(5)
CONT_DATA_OFFSET = const(3)
class InitHeader:

@ -9,18 +9,18 @@ from .protocol_common import MessageWithId
from .thp import ChannelState, ack_handler, checksum, thp_messages
from .thp import thp_session as THP
from .thp.channel_context import (
_MAX_PAYLOAD_LEN,
_REPORT_LENGTH,
CONT_DATA_OFFSET,
INIT_DATA_OFFSET,
MAX_PAYLOAD_LEN,
REPORT_LENGTH,
ChannelContext,
load_cached_channels,
)
from .thp.checksum import CHECKSUM_LENGTH
from .thp.thp_messages import (
CODEC_V1,
CONT_DATA_OFFSET,
CONTINUATION_PACKET,
ENCRYPTED_TRANSPORT,
INIT_DATA_OFFSET,
InitHeader,
InterruptingInitPacket,
)
@ -192,7 +192,7 @@ def _get_loop_wait_read(iface: WireInterface):
def _get_buffer_for_payload(
payload_length: int, existing_buffer: utils.BufferType, max_length=_MAX_PAYLOAD_LEN
payload_length: int, existing_buffer: utils.BufferType, max_length=MAX_PAYLOAD_LEN
) -> utils.BufferType:
if payload_length > max_length:
raise ThpError("Message too large")
@ -201,7 +201,7 @@ def _get_buffer_for_payload(
try:
payload: utils.BufferType = bytearray(payload_length)
except MemoryError:
payload = bytearray(_REPORT_LENGTH)
payload = bytearray(REPORT_LENGTH)
raise ThpError("Message too large")
return payload
@ -302,7 +302,7 @@ async def write_to_wire(
payload_length = len(payload)
# prepare the report buffer with header data
report = bytearray(_REPORT_LENGTH)
report = bytearray(REPORT_LENGTH)
header.pack_to_buffer(report)
# write initial report

@ -322,7 +322,7 @@ class TestWireTrezorHostProtocolV1(unittest.TestCase):
# ensure that a message this big won't fit into memory
# Note: this control is changed, because THP has only 2 byte length field
self.assertTrue(message_size > thp_v1._MAX_PAYLOAD_LEN)
self.assertTrue(message_size > thp_v1.MAX_PAYLOAD_LEN)
# self.assertRaises(MemoryError, bytearray, message_size)
header = make_header(PLAINTEXT_1, COMMON_CID, message_size)
packet = header + MESSAGE_TYPE_BYTES + (b"\x00" * INIT_MESSAGE_DATA_LENGTH)

Loading…
Cancel
Save