refactor(core/ethereum): improve API of the rlp module

core/.changelog.d/1704.changed

@@ -0,0 +1 @@
+Refactor RLP codec for better clarity and some small memory savings.

core/src/apps/ethereum/sign_tx.py

@@ -68,19 +68,19 @@ async def sign_tx(ctx, msg, keychain):
     total_length = get_total_length(msg, data_total)
 
     sha = HashWriter(sha3_256(keccak=True))
-    sha.extend(rlp.encode_length(total_length, True))  # total length
+    rlp.write_header(sha, total_length, rlp.LIST_HEADER_BYTE)
 
     if msg.tx_type is not None:
-        sha.extend(rlp.encode(msg.tx_type))
+        rlp.write(sha, msg.tx_type)
 
     for field in (msg.nonce, msg.gas_price, msg.gas_limit, address_bytes, msg.value):
-        sha.extend(rlp.encode(field))
+        rlp.write(sha, field)
 
     if data_left == 0:
-        sha.extend(rlp.encode(data))
+        rlp.write(sha, data)
     else:
-        sha.extend(rlp.encode_length(data_total, False))
-        sha.extend(rlp.encode(data, False))
+        rlp.write_header(sha, data_total, rlp.STRING_HEADER_BYTE, data)
+        sha.extend(data)
 
     while data_left > 0:
         resp = await send_request_chunk(ctx, data_left)
@@ -89,9 +89,9 @@ async def sign_tx(ctx, msg, keychain):
 
     # eip 155 replay protection
     if msg.chain_id:
-        sha.extend(rlp.encode(msg.chain_id))
-        sha.extend(rlp.encode(0))
-        sha.extend(rlp.encode(0))
+        rlp.write(sha, msg.chain_id)
+        rlp.write(sha, 0)
+        rlp.write(sha, 0)
 
     digest = sha.get_digest()
     result = sign_digest(msg, keychain, digest)
@@ -102,29 +102,24 @@ async def sign_tx(ctx, msg, keychain):
 def get_total_length(msg: EthereumSignTx, data_total: int) -> int:
     length = 0
     if msg.tx_type is not None:
-        length += rlp.field_length(1, [msg.tx_type])
+        length += rlp.length(msg.tx_type)
 
-    length += rlp.field_length(len(msg.nonce), msg.nonce[:1])
-    length += rlp.field_length(len(msg.gas_price), msg.gas_price)
-    length += rlp.field_length(len(msg.gas_limit), msg.gas_limit)
-    to = address.bytes_from_address(msg.to)
-    length += rlp.field_length(len(to), to)
-    length += rlp.field_length(len(msg.value), msg.value)
+    for item in (
+        msg.nonce,
+        msg.gas_price,
+        msg.gas_limit,
+        address.bytes_from_address(msg.to),
+        msg.value,
+    ):
+        length += rlp.length(item)
 
     if msg.chain_id:  # forks replay protection
-        if msg.chain_id < 0x100:
-            l = 1
-        elif msg.chain_id < 0x1_0000:
-            l = 2
-        elif msg.chain_id < 0x100_0000:
-            l = 3
-        else:
-            l = 4
-        length += rlp.field_length(l, [msg.chain_id])
-        length += rlp.field_length(0, 0)
-        length += rlp.field_length(0, 0)
+        length += rlp.length(msg.chain_id)
+        length += rlp.length(0)
+        length += rlp.length(0)
 
-    length += rlp.field_length(data_total, msg.data_initial_chunk)
+    length += rlp.header_length(data_total, msg.data_initial_chunk)
+    length += data_total
     return length
 
 

core/src/trezor/crypto/rlp.py

@@ -1,53 +1,97 @@
+from micropython import const
+
+if False:
+    from typing import Union
+    from trezor.utils import Writer
+
+    # what we want:
+    # RLPItem = Union[list["RLPItem"], bytes, int]
+    # what mypy can process:
+    RLPItem = Union[list, bytes, int]
+
+
+STRING_HEADER_BYTE = const(0x80)
+LIST_HEADER_BYTE = const(0xC0)
+
+
+def _byte_size(x: int) -> int:
+    if x < 0:
+        raise ValueError  # only unsigned ints are supported
+    for exp in range(64):
+        if x < 0x100 ** exp:
+            return exp
+    else:
+        raise ValueError  # int is too large
+
+
 def int_to_bytes(x: int) -> bytes:
-    if x == 0:
-        return b""
-    r = bytearray()
-    while x:
-        r.append(x % 256)
-        x //= 256
-    return bytes(reversed(r))
+    return x.to_bytes(_byte_size(x), "big")
 
 
-def encode_length(l: int, is_list: bool) -> bytes:
-    offset = 0xC0 if is_list else 0x80
-    if l < 56:
-        return bytes([l + offset])
-    elif l < 256 ** 8:
-        bl = int_to_bytes(l)
-        return bytes([len(bl) + offset + 55]) + bl
-    else:
-        raise ValueError("Input too long")
+def write_header(
+    w: Writer,
+    length: int,
+    header_byte: int,
+    data_start: bytes | None = None,
+) -> None:
+    if length == 1 and data_start is not None and data_start[0] <= 0x7F:
+        # no header when encoding one byte below 0x80
+        pass
 
-
-def encode(data, include_length=True) -> bytes:
-    if isinstance(data, int):
-        data = int_to_bytes(data)
-    if isinstance(data, bytearray):
-        data = bytes(data)
-    if isinstance(data, bytes):
-        if (len(data) == 1 and ord(data) < 128) or not include_length:
-            return data
-        else:
-            return encode_length(len(data), is_list=False) + data
-    elif isinstance(data, list):
-        output = b""
-        for item in data:
-            output += encode(item)
-        if include_length:
-            return encode_length(len(output), is_list=True) + output
-        else:
-            return output
-    else:
-        raise TypeError("Invalid input of type " + str(type(data)))
-
-
-def field_length(length: int, first_byte: bytearray) -> int:
-    if length == 1 and first_byte[0] <= 0x7F:
-        return 1
     elif length <= 55:
-        return 1 + length
-    elif length <= 0xFF:
-        return 2 + length
-    elif length <= 0xFFFF:
-        return 3 + length
-    return 4 + length
+        w.append(header_byte + length)
+
+    else:
+        encoded_length = int_to_bytes(length)
+        w.append(header_byte + 55 + len(encoded_length))
+        w.extend(encoded_length)
+
+
+def header_length(length: int, data_start: bytes | None = None) -> int:
+    if length == 1 and data_start is not None and data_start[0] <= 0x7F:
+        # no header when encoding one byte below 0x80
+        return 0
+
+    if length <= 55:
+        return 1
+
+    return 1 + _byte_size(length)
+
+
+def length(item: RLPItem) -> int:
+    data: bytes | None = None
+    if isinstance(item, int):
+        data = int_to_bytes(item)
+        item_length = len(data)
+    elif isinstance(item, (bytes, bytearray)):
+        data = item
+        item_length = len(item)
+    elif isinstance(item, list):
+        item_length = sum(length(i) for i in item)
+    else:
+        raise TypeError
+
+    return header_length(item_length, data) + item_length
+
+
+def write_string(w: Writer, string: bytes) -> None:
+    write_header(w, len(string), STRING_HEADER_BYTE, string)
+    w.extend(string)
+
+
+def write_list(w: Writer, lst: list[RLPItem]) -> None:
+    payload_length = sum(length(item) for item in lst)
+    write_header(w, payload_length, LIST_HEADER_BYTE)
+    for item in lst:
+        write(w, item)
+
+
+def write(w: Writer, item: RLPItem) -> None:
+    if isinstance(item, int):
+        write_string(w, int_to_bytes(item))
+    elif isinstance(item, (bytes, bytearray)):
+        write_string(w, item)
+    elif isinstance(item, list):
+        write_list(w, item)
+    else:
+        raise TypeError

core/tests/test_trezor.crypto.rlp.py

@@ -17,6 +17,8 @@ class TestCryptoRlp(unittest.TestCase):
          'b8384c6f72656d20697073756d20646f6c6f722073697420616d65742c20636f6e7365637465747572206164697069736963696e6720656c6974'),
         (b'Lorem ipsum dolor sit amet, consectetur adipiscing elit. Curabitur mauris magna, suscipit sed vehicula non, iaculis faucibus tortor. Proin suscipit ultricies malesuada. Duis tortor elit, dictum quis tristique eu, ultrices at risus. Morbi a est imperdiet mi ullamcorper aliquet suscipit nec lorem. Aenean quis leo mollis, vulputate elit varius, consequat enim. Nulla ultrices turpis justo, et posuere urna consectetur nec. Proin non convallis metus. Donec tempor ipsum in mauris congue sollicitudin. Vestibulum ante ipsum primis in faucibus orci luctus et ultrices posuere cubilia Curae; Suspendisse convallis sem vel massa faucibus, eget lacinia lacus tempor. Nulla quis ultricies purus. Proin auctor rhoncus nibh condimentum mollis. Aliquam consequat enim at metus luctus, a eleifend purus egestas. Curabitur at nibh metus. Nam bibendum, neque at auctor tristique, lorem libero aliquet arcu, non interdum tellus lectus sit amet eros. Cras rhoncus, metus ac ornare cursus, dolor justo ultrices metus, at ullamcorper volutpat',
          'b904004c6f72656d20697073756d20646f6c6f722073697420616d65742c20636f6e73656374657475722061646970697363696e6720656c69742e20437572616269747572206d6175726973206d61676e612c20737573636970697420736564207665686963756c61206e6f6e2c20696163756c697320666175636962757320746f72746f722e2050726f696e20737573636970697420756c74726963696573206d616c6573756164612e204475697320746f72746f7220656c69742c2064696374756d2071756973207472697374697175652065752c20756c7472696365732061742072697375732e204d6f72626920612065737420696d70657264696574206d6920756c6c616d636f7270657220616c6971756574207375736369706974206e6563206c6f72656d2e2041656e65616e2071756973206c656f206d6f6c6c69732c2076756c70757461746520656c6974207661726975732c20636f6e73657175617420656e696d2e204e756c6c6120756c74726963657320747572706973206a7573746f2c20657420706f73756572652075726e6120636f6e7365637465747572206e65632e2050726f696e206e6f6e20636f6e76616c6c6973206d657475732e20446f6e65632074656d706f7220697073756d20696e206d617572697320636f6e67756520736f6c6c696369747564696e2e20566573746962756c756d20616e746520697073756d207072696d697320696e206661756369627573206f726369206c756374757320657420756c74726963657320706f737565726520637562696c69612043757261653b2053757370656e646973736520636f6e76616c6c69732073656d2076656c206d617373612066617563696275732c2065676574206c6163696e6961206c616375732074656d706f722e204e756c6c61207175697320756c747269636965732070757275732e2050726f696e20617563746f722072686f6e637573206e69626820636f6e64696d656e74756d206d6f6c6c69732e20416c697175616d20636f6e73657175617420656e696d206174206d65747573206c75637475732c206120656c656966656e6420707572757320656765737461732e20437572616269747572206174206e696268206d657475732e204e616d20626962656e64756d2c206e6571756520617420617563746f72207472697374697175652c206c6f72656d206c696265726f20616c697175657420617263752c206e6f6e20696e74657264756d2074656c6c7573206c65637475732073697420616d65742065726f732e20437261732072686f6e6375732c206d65747573206163206f726e617265206375727375732c20646f6c6f72206a7573746f20756c747269636573206d657475732c20617420756c6c616d636f7270657220766f6c7574706174'),
+        (0,
+         '80'),
         (1,
          '01'),
         (16,
@@ -27,10 +29,28 @@ class TestCryptoRlp(unittest.TestCase):
          '7f'),
         (128,
          '8180'),
+        (254,
+         '81fe'),
+        (255,
+         '81ff'),
+        (256,
+         '820100'),
         (1000,
          '8203e8'),
         (100000,
          '830186a0'),
+        (0xffff,
+         '82ffff'),
+        (0x1_0000,
+         '83010000'),
+        (0xff_ffff,
+         '83ffffff'),
+        (0x100_0000,
+         '8401000000'),
+        (0xffff_ffff,
+         '84ffffffff'),
+        (0x1_0000_0000,
+         '850100000000'),
         (83729609699884896815286331701780722,
          '8f102030405060708090a0b0c0d0e0f2'),
         (105315505618206987246253880190783558935785933862974822347068935681,
@@ -57,12 +77,19 @@ class TestCryptoRlp(unittest.TestCase):
          'a1010000000000000000000000000000000000000000000000000000000000000000'),
     ]
 
-    def test_rlp_encode(self):
-
+    def test_rlp_write(self):
         for i, o in self.vectors:
             o = unhexlify(o)
-            o2 = rlp.encode(i)
-            self.assertEqual(o, o2)
+            w = bytearray()
+            rlp.write(w, i)
+            self.assertEqual(w, o)
+
+
+    def test_rlp_length(self):
+        for i, o in self.vectors:
+            length = rlp.length(i)
+            self.assertEqual(length, len(o) // 2)
+
 
 
 if __name__ == '__main__':