diff --git a/eth2/utils/ssz/src/cached_tree_hash.rs b/eth2/utils/ssz/src/cached_tree_hash.rs index 7a722766c..8a7b07f50 100644 --- a/eth2/utils/ssz/src/cached_tree_hash.rs +++ b/eth2/utils/ssz/src/cached_tree_hash.rs @@ -1,4 +1,5 @@ use hashing::hash; +use int_to_bytes::int_to_bytes32; use std::fmt::Debug; use std::iter::Iterator; use std::ops::Range; @@ -231,6 +232,15 @@ impl TreeHashCache { Ok(hash(&child_bytes)) } + pub fn mix_in_length(&self, chunk: usize, length: usize) -> Result, Error> { + let mut bytes = Vec::with_capacity(2 * BYTES_PER_CHUNK); + + bytes.append(&mut self.get_chunk(chunk)?.to_vec()); + bytes.append(&mut int_to_bytes32(length as u64)); + + Ok(hash(&bytes)) + } + pub fn into_merkle_tree(self) -> Vec { self.cache } @@ -300,6 +310,10 @@ impl OffsetHandler { }) } + pub fn root(&self) -> usize { + self.first_node + } + pub fn height(&self) -> usize { self.num_leaf_nodes.trailing_zeros() as usize } diff --git a/eth2/utils/ssz/src/cached_tree_hash/impls.rs b/eth2/utils/ssz/src/cached_tree_hash/impls.rs index 0377649cb..558b4dde5 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/impls.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/impls.rs @@ -212,6 +212,12 @@ where } } + // If the root node or the length has changed, mix in the length of the list. + let root_node = offset_handler.root(); + if cache.changed(root_node)? | (self.len() != other.len()) { + cache.modify_chunk(root_node, &cache.mix_in_length(root_node, self.len())?)?; + } + Ok(offset_handler.next_node()) } } diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs index f09fac419..b85c16587 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/tests.rs @@ -343,6 +343,13 @@ fn outer_builds() { assert_eq!(merkle, cache); } +fn mix_in_length(root: &mut [u8], len: usize) { + let mut bytes = root.to_vec(); + bytes.append(&mut int_to_bytes32(len as u64)); + + root.copy_from_slice(&hash(&bytes)); +} + /// Generic test that covers: /// /// 1. Produce a new cache from `original`. @@ -367,7 +374,9 @@ fn test_u64_vec_modifications(original: Vec, modified: Vec) { data.append(&mut int_to_bytes8(*i)); } let data = sanitise_bytes(data); - let expected = merkleize(data); + let mut expected = merkleize(data); + + mix_in_length(&mut expected[0..HASHSIZE], modified.len()); assert_eq!(expected, modified_cache); } @@ -490,6 +499,8 @@ fn test_inner_vec_modifications(original: Vec, modified: Vec, refe expected.append(&mut vec![0; HASHSIZE]); } + mix_in_length(&mut expected[0..HASHSIZE], modified.len()); + // Compare the cached tree to the reference tree. assert_trees_eq(&expected, &modified_cache); }