From 3a69524c5245c383da1fa3d0fafc2b27e2429643 Mon Sep 17 00:00:00 2001 From: Jan Pochyla Date: Wed, 20 Jul 2016 15:33:42 +0200 Subject: [PATCH] simplify read_wire_msg a bit --- src/trezor/wire.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/src/trezor/wire.py b/src/trezor/wire.py index 36397e544..6e6fe08a9 100644 --- a/src/trezor/wire.py +++ b/src/trezor/wire.py @@ -39,23 +39,19 @@ def _write_report(rep): def read_wire_msg(): rep = yield from _read_report() - assert rep[0] == _HEADER_MAGIC, 'Incorrect report magic' - - # Parse message header - sid, mtype, mlen = ustruct.unpack_from('>LLL', rep, 1) # Skip magic + magic, sid, mtype, mlen = ustruct.unpack('>BLLL', rep) + assert magic == _HEADER_MAGIC, 'Incorrect report magic' assert mlen < _MAX_DATA_LEN, 'Message too large to read' mlen += 4 # Account for the checksum data = rep[13:][:mlen] # Skip magic and header, trim to data len - buffered = bytearray(data) # Resulting message data - remaining = mlen - len(buffered) + remaining = mlen - len(data) + buffered = bytearray(data) if remaining > 0 else data # Avoid the copy if we don't append while remaining > 0: rep = yield from _read_report() - assert rep[0] == _DATA_MAGIC, 'Incorrect report magic' - - # Compare the session IDs - rsid = ustruct.unpack_from('>L', rep, 1) + magic, rsid = ustruct.unpack('>BL', rep) + assert magic == _DATA_MAGIC, 'Incorrect report magic' assert rsid == sid, 'Session ID mismatch' data = rep[5:][:remaining] # Skip magic and session ID, trim