using HashWriter object to collect the final hash continously

core/src/apps/ethereum/typed_data.py

@@ -14,8 +14,12 @@
 from .address import address_from_bytes
 
 
+def get_hash_writer() -> HashWriter:
+    return HashWriter(sha3_256(keccak=True))
+
+
 def keccak256(message: bytes) -> bytes:
-    h = HashWriter(sha3_256(keccak=True))
+    h = get_hash_writer()
     h.extend(message)
     return h.get_digest()
 
@@ -29,21 +33,19 @@ def hash_struct(
     """
     Encodes and hashes an object using Keccak256
     """
-    # TODO: create hashwriter object and pass it through other functions
-    # w: Writer
-    # w.append, w.extend methods
-    # return w.digest()
-    type_hash = hash_type(primary_type, types)
-    encoded_data = encode_data(primary_type, data, types, metamask_v4_compat)
-    return keccak256(type_hash + encoded_data)
+    w = get_hash_writer()
+    hash_type(w, primary_type, types)
+    encode_data(w, primary_type, data, types, metamask_v4_compat)
+    return w.get_digest()
 
 
 def encode_data(
+    w: HashWriter,
     primary_type: str,
     data: dict,
     types: Dict[str, EthereumTypedDataStructAck],
     metamask_v4_compat: bool = True,
-) -> bytes:
+) -> None:
     """
     Encodes an object by encoding and concatenating each of its members
 
@@ -56,29 +58,26 @@ def encode_data(
     data - Object to encode
     types - Type definitions
     """
-    result = b""
-
     type_members = types[primary_type].members
     for member in type_members:
-        encoded_value = encode_field(
+        encode_field(
+            w=w,
             field=member.type,
             value=data[member.name],
             types=types,
             in_array=False,
             metamask_v4_compat=metamask_v4_compat,
         )
-        result += encoded_value
-
-    return result
 
 
 def encode_field(
+    w: HashWriter,
     field: EthereumFieldType,
     value: bytes,
     types: Dict[str, EthereumTypedDataStructAck],
     in_array: bool,
     metamask_v4_compat: bool,
-) -> bytes:
+) -> None:
     """
     SPEC:
     Atomic types:
@@ -99,53 +98,57 @@ def encode_field(
 
     # Arrays and structs need special recursive handling
     if data_type == EthereumDataType.ARRAY:
-        buf = b""
+        arr_w = get_hash_writer()
         for element in value:
-            buf += encode_field(
+            encode_field(
+                w=arr_w,
                 field=field.entry_type,
                 value=element,
                 types=types,
                 in_array=True,
                 metamask_v4_compat=metamask_v4_compat,
             )
-        return keccak256(buf)
+        w.extend(arr_w.get_digest())
     elif data_type == EthereumDataType.STRUCT:
         # Metamask V4 implementation has a bug, that causes the
         # behavior of structs in array be different from SPEC
         # Explanation at https://github.com/MetaMask/eth-sig-util/pull/107
         # encode_data() is the way to process structs in arrays, but
         # Metamask V4 is using hash_struct() even in this case
         if in_array and not metamask_v4_compat:
-            return encode_data(
+            encode_data(
+                w=w,
                 primary_type=field.struct_name,
                 data=value,
                 types=types,
                 metamask_v4_compat=metamask_v4_compat,
             )
         else:
-            return hash_struct(
-                primary_type=field.struct_name,
-                data=value,
-                types=types,
-                metamask_v4_compat=metamask_v4_compat,
+            w.extend(
+                hash_struct(
+                    primary_type=field.struct_name,
+                    data=value,
+                    types=types,
+                    metamask_v4_compat=metamask_v4_compat,
+                )
             )
     elif data_type == EthereumDataType.BYTES:
         # TODO: is not tested
         if field.size is None:
-            return keccak256(value)
+            w.extend(keccak256(value))
         else:
-            return rightpad32(value)
+            w.extend(rightpad32(value))
     elif data_type == EthereumDataType.STRING:
-        return keccak256(value)
+        w.extend((keccak256(value)))
     elif data_type in [
         EthereumDataType.UINT,
         EthereumDataType.INT,
         EthereumDataType.BOOL,
         EthereumDataType.ADDRESS,
     ]:
-        return leftpad32(value)
-
-    raise ValueError  # Unsupported data type for field encoding
+        w.extend(leftpad32(value))
+    else:
+        raise ValueError  # Unsupported data type for field encoding
 
 
 def leftpad32(value: bytes) -> bytes:
@@ -199,11 +202,12 @@ def validate_field(field: EthereumFieldType, field_name: str, value: bytes) -> N
             raise wire.DataError("{}: invalid UTF-8".format(field_name))
 
 
-def hash_type(primary_type: str, types: Dict[str, EthereumTypedDataStructAck]) -> bytes:
+def hash_type(w: HashWriter, primary_type: str, types: Dict[str, EthereumTypedDataStructAck]) -> None:
     """
     Encodes and hashes a type using Keccak256
     """
-    return keccak256(encode_type(primary_type, types))
+    result = keccak256(encode_type(primary_type, types))
+    w.extend(result)
 
 
 def encode_type(

core/tests/test_apps.ethereum.typed_data.py

@@ -21,6 +21,7 @@
         keccak256,
         get_type_name,
         decode_data,
+        get_hash_writer,
     )
 
 DOMAIN_TYPES = {
@@ -409,13 +410,15 @@ def test_encode_data(self):
             ),
         )
         for primary_type, data, types, expected in VECTORS:
-            res = encode_data(
+            w = get_hash_writer()
+            encode_data(
+                w=w,
                 primary_type=primary_type,
                 data=data,
                 types=types,
                 metamask_v4_compat=True,
             )
-            self.assertEqual(res, expected)
+            self.assertEqual(w.get_digest(), keccak256(expected))
 
     def test_encode_type(self):
         VECTORS = (  # primary_type, types, expected
@@ -464,8 +467,9 @@ def test_hash_type(self):
         )
 
         for primary_type, expected in VECTORS:
-            res = hash_type(primary_type=primary_type, types=ALL_TYPES_BASIC)
-            self.assertEqual(res, expected)
+            w = get_hash_writer()
+            hash_type(w=w, primary_type=primary_type, types=ALL_TYPES_BASIC)
+            self.assertEqual(w.get_digest(), keccak256(expected))
 
     def test_find_typed_dependencies(self):
         VECTORS = (  # primary_type, expected
@@ -498,7 +502,6 @@ def test_find_typed_dependencies(self):
             self.assertEqual(res, expected)
 
     def test_encode_field(self):
-        # TODO: need to add a fake writer and check it really got written
         VECTORS = (  # field, value, expected
             (
                 EthereumFieldType(data_type=EthereumDataType.STRING, size=None),
@@ -565,14 +568,16 @@ def test_encode_field(self):
             # metamask_v4_compat should not have any effect on the
             # result for items outside of arrays
             for metamask_v4_compat in [True, False]:
-                res = encode_field(
+                w = get_hash_writer()
+                encode_field(
+                    w=w,
                     field=field,
                     value=value,
                     types=MESSAGE_TYPES_BASIC,
                     in_array=False,
                     metamask_v4_compat=metamask_v4_compat,
                 )
-                self.assertEqual(res, expected)
+                self.assertEqual(w.get_digest(), keccak256(expected))
 
         # metamask_v4_compat makes a difference in arrays of structs when False
         field = EthereumFieldType(
@@ -584,14 +589,26 @@ def test_encode_field(self):
         }
         expected_not_in_array = b"5'H\x1b_\xa5a5\x06\x04\xa6\rsOI\xee\x90\x7f\x17O[\xa6\xbby\x1a\xabAun\xce~\xd1"
         expected_in_array = b"(\xca\xc3\x18\xa8l\x8a\nj\x91V\xc2\xdb\xa2\xc8\xc266w\xba\x05\x14\xefae\x92\xd8\x15W\xe6y\xb6\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00T\xb0\xfaf\xa0et\x8c@\xdc\xa2\xc7\xfe\x12Z (\xcf\x99\x82"
-        res_not_in_array = encode_field(
-            field, value, MESSAGE_TYPES_BASIC, in_array=False, metamask_v4_compat=False
+        w_not_in_array = get_hash_writer()
+        w_in_array = get_hash_writer()
+        encode_field(
+            w=w_not_in_array,
+            field=field,
+            value=value,
+            types=MESSAGE_TYPES_BASIC,
+            in_array=False,
+            metamask_v4_compat=False,
         )
-        res_in_array = encode_field(
-            field, value, MESSAGE_TYPES_BASIC, in_array=True, metamask_v4_compat=False
+        encode_field(
+            w=w_in_array,
+            field=field,
+            value=value,
+            types=MESSAGE_TYPES_BASIC,
+            in_array=True,
+            metamask_v4_compat=False,
         )
-        self.assertEqual(res_not_in_array, expected_not_in_array)
-        self.assertEqual(res_in_array, expected_in_array)
+        self.assertEqual(w_not_in_array.get_digest(), keccak256(expected_not_in_array))
+        self.assertEqual(w_in_array.get_digest(), keccak256(expected_in_array))
 
     def test_validate_field(self):
         VECTORS_VALID_INVALID = (  # field, valid_values, invalid_values