From bcab08ad29cafe297795ee5fe23513314b6881c3 Mon Sep 17 00:00:00 2001 From: cchudant Date: Mon, 30 Sep 2024 09:02:46 +0000 Subject: [PATCH] feat: tree max_height --- src/error.rs | 5 +++++ src/lib.rs | 7 +++++-- src/tests/madara_comparison.rs | 2 +- src/tests/merge.rs | 5 +++-- src/tests/merkle_tree.rs | 3 ++- src/tests/proptest.rs | 2 +- src/tests/simple.rs | 28 ++++++++++++++-------------- src/tests/transactional_state.rs | 30 +++++++++++++++++------------- src/tests/trie_log.rs | 12 ++++++------ src/trie/iterator.rs | 1 + src/trie/proof.rs | 24 +++++++++++++++--------- src/trie/tree.rs | 18 +++++++++++++++++- src/trie/trees.rs | 17 ++++++++++------- 13 files changed, 97 insertions(+), 57 deletions(-) diff --git a/src/error.rs b/src/error.rs index 747ecb4..771b833 100644 --- a/src/error.rs +++ b/src/error.rs @@ -23,6 +23,8 @@ where NodeDecodeError(parity_scale_codec::Error), /// Error when creating a storage proof. CreateProof(String), + /// Malformated trie key. + KeyLength { expected: usize, got: usize }, } impl core::convert::From @@ -55,6 +57,9 @@ where BonsaiStorageError::Database(e) => write!(f, "Database error: {}", e), BonsaiStorageError::NodeDecodeError(e) => write!(f, "Node decode error: {}", e), BonsaiStorageError::CreateProof(e) => write!(f, "Proof creation error: {}", e), + BonsaiStorageError::KeyLength { expected, got } => { + write!(f, "Malformated key length: expected {expected}, got {got}") + } } } } diff --git a/src/lib.rs b/src/lib.rs index 144764e..cf0a5a8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -237,21 +237,23 @@ where pub fn new( db: DB, config: BonsaiStorageConfig, + max_height: u8, ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), None); Ok(Self { - tries: MerkleTrees::new(key_value_db), + tries: MerkleTrees::new(key_value_db, max_height), }) } pub fn new_from_transactional_state( db: DB, config: BonsaiStorageConfig, + max_height: u8, created_at: ChangeID, _identifiers: impl IntoIterator>, ) -> Result> { let key_value_db = KeyValueDB::new(db, config.into(), Some(created_at)); - let tries = MerkleTrees::::new(key_value_db); + let tries = MerkleTrees::::new(key_value_db, max_height); Ok(Self { tries }) } @@ -516,6 +518,7 @@ where Ok(Some(BonsaiStorage::new_from_transactional_state( transaction, config, + self.tries.max_height, change_id, self.tries.get_identifiers(), )?)) diff --git a/src/tests/madara_comparison.rs b/src/tests/madara_comparison.rs index 529083b..793faaf 100644 --- a/src/tests/madara_comparison.rs +++ b/src/tests/madara_comparison.rs @@ -15,7 +15,7 @@ fn trie_height_251() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); for i in 0..251 { let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec(); key.set(i, true); diff --git a/src/tests/merge.rs b/src/tests/merge.rs index 0b12c79..ff6c33a 100644 --- a/src/tests/merge.rs +++ b/src/tests/merge.rs @@ -60,8 +60,9 @@ fn init_test( let identifier = vec![]; let config = BonsaiStorageConfig::default(); - let mut bonsai_storage = BonsaiStorage::new(RocksDB::new(db, RocksDBConfig::default()), config) - .expect("Failed to create BonsaiStorage"); + let mut bonsai_storage = + BonsaiStorage::new(RocksDB::new(db, RocksDBConfig::default()), config, 251) + .expect("Failed to create BonsaiStorage"); let mut id_builder = BasicIdBuilder::new(); diff --git a/src/tests/merkle_tree.rs b/src/tests/merkle_tree.rs index d37f269..9fcd768 100644 --- a/src/tests/merkle_tree.rs +++ b/src/tests/merkle_tree.rs @@ -17,7 +17,8 @@ fn test_key_retrieval() { let rocksdb = create_rocks_db(tempdir.path()).unwrap(); let db = RocksDB::new(&rocksdb, RocksDBConfig::default()); let mut bonsai = - BonsaiStorage::::new(db, BonsaiStorageConfig::default()).unwrap(); + BonsaiStorage::::new(db, BonsaiStorageConfig::default(), 251) + .unwrap(); let block_0 = vec![ ( diff --git a/src/tests/proptest.rs b/src/tests/proptest.rs index 6ce724c..cb0e7b9 100644 --- a/src/tests/proptest.rs +++ b/src/tests/proptest.rs @@ -73,7 +73,7 @@ impl MerkleTreeInsertProblem { let mut ckv = HashMap::new(); // apply steps - let mut tree = MerkleTree::::new(smallvec![]); + let mut tree = MerkleTree::::new(smallvec![], 251); for step in &self.0 { match step { Step::Insert(k, v) => { diff --git a/src/tests/simple.rs b/src/tests/simple.rs index 1133127..62e391e 100644 --- a/src/tests/simple.rs +++ b/src/tests/simple.rs @@ -14,7 +14,7 @@ fn basics() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 1], @@ -69,7 +69,7 @@ fn root_hash_similar_rocks_db() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 1], @@ -116,7 +116,7 @@ fn root_hash_similar_rocks_db() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 3], @@ -156,13 +156,13 @@ fn starknet_specific() { let db1 = create_rocks_db(tempdir1.path()).unwrap(); let config1 = BonsaiStorageConfig::default(); let mut bonsai_storage1: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db1, RocksDBConfig::default()), config1).unwrap(); + BonsaiStorage::new(RocksDB::new(&db1, RocksDBConfig::default()), config1, 251).unwrap(); let tempdir2 = tempfile::tempdir().unwrap(); let db2 = create_rocks_db(tempdir2.path()).unwrap(); let config2 = BonsaiStorageConfig::default(); let mut bonsai_storage2: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db2, RocksDBConfig::default()), config2).unwrap(); + BonsaiStorage::new(RocksDB::new(&db2, RocksDBConfig::default()), config2, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let contract_states = vec![ @@ -237,7 +237,7 @@ fn root_hash_similar_hashmap_db() { let db = HashMapDb::::default(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(db, config).unwrap(); + BonsaiStorage::new(db, config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 1], @@ -283,7 +283,7 @@ fn root_hash_similar_hashmap_db() { let db = HashMapDb::::default(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(db, config).unwrap(); + BonsaiStorage::new(db, config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( vec![1, 2, 3], @@ -322,7 +322,7 @@ fn double_insert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let contract_states = vec![ ContractState { @@ -385,7 +385,7 @@ fn double_identifier() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let contract_states = vec![ ContractState { @@ -450,7 +450,7 @@ fn get_changes() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = (vec![1, 2, 1], Felt::from_hex("0x01").unwrap()); let bitvec = BitVec::from_vec(pair1.0.clone()); @@ -501,7 +501,7 @@ fn keyer(felt: Felt) -> BitVec { fn test_insert_zero() { let config = BonsaiStorageConfig::default(); let bonsai_db = HashMapDb::::default(); - let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes(); @@ -645,7 +645,7 @@ fn test_block_7_starknet() { let _ = env_logger::builder().is_test(true).try_init(); let config = BonsaiStorageConfig::default(); let bonsai_db = HashMapDb::::default(); - let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x056e4fed965fccd7fb01fcadd827470338f35ced62275328929d0d725b5707ba".as_bytes(); @@ -857,7 +857,7 @@ fn test_block_7_starknet() { fn test_block_7_starknet_2() { let config = BonsaiStorageConfig::default(); let bonsai_db = HashMapDb::::default(); - let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x421203c58e1b4a6c3675be26cfaa18d2b6b42695ca206be1f08ce29f7f1bc7c".as_bytes(); @@ -988,7 +988,7 @@ fn test_block_7_starknet_2() { fn test_block_9() { let config = BonsaiStorageConfig::default(); let bonsai_db = HashMapDb::::default(); - let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config) + let mut bonsai_storage = BonsaiStorage::<_, _, Pedersen>::new(bonsai_db, config, 251) .expect("Failed to create bonsai storage"); let identifier = "0x06F3C934BA4EC49245CB9A42FC715E4D589AA502AF69BE13916127A538D525CE".as_bytes(); diff --git a/src/tests/transactional_state.rs b/src/tests/transactional_state.rs index 91cb6b2..30c119c 100644 --- a/src/tests/transactional_state.rs +++ b/src/tests/transactional_state.rs @@ -14,7 +14,7 @@ fn basics() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -58,8 +58,12 @@ fn test_thread() { let tempdir = tempfile::tempdir().unwrap(); let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); - let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config.clone()).unwrap(); + let mut bonsai_storage = BonsaiStorage::new( + RocksDB::new(&db, RocksDBConfig::default()), + config.clone(), + 251, + ) + .unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -119,7 +123,7 @@ fn remove() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -164,7 +168,7 @@ fn merge() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -207,7 +211,7 @@ fn merge_with_uncommitted_insert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -269,7 +273,7 @@ fn merge_with_uncommitted_remove() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -327,7 +331,7 @@ fn transactional_state_after_uncommitted() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -369,7 +373,7 @@ fn merge_override() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -412,7 +416,7 @@ fn merge_remove() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -451,7 +455,7 @@ fn merge_txn_revert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -511,7 +515,7 @@ fn merge_invalid() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -558,7 +562,7 @@ fn many_snapshots() { ..Default::default() }; let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( diff --git a/src/tests/trie_log.rs b/src/tests/trie_log.rs index ffd9022..f83ce8b 100644 --- a/src/tests/trie_log.rs +++ b/src/tests/trie_log.rs @@ -13,7 +13,7 @@ fn basics() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -62,7 +62,7 @@ fn unrecorded_revert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -87,7 +87,7 @@ fn in_place_revert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = (vec![1, 2, 3], &BonsaiTrieHash::default()); @@ -110,7 +110,7 @@ fn truncated_revert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -150,7 +150,7 @@ fn double_revert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( @@ -192,7 +192,7 @@ fn remove_and_reinsert() { let db = create_rocks_db(tempdir.path()).unwrap(); let config = BonsaiStorageConfig::default(); let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> = - BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config, 251).unwrap(); let mut id_builder = BasicIdBuilder::new(); let pair1 = ( diff --git a/src/trie/iterator.rs b/src/trie/iterator.rs index 8621d36..6306e65 100644 --- a/src/trie/iterator.rs +++ b/src/trie/iterator.rs @@ -283,6 +283,7 @@ mod tests { let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( RocksDB::::new(&db, RocksDBConfig::default()), BonsaiStorageConfig::default(), + 8, ) .unwrap(); diff --git a/src/trie/proof.rs b/src/trie/proof.rs index 4f6ff77..b108937 100644 --- a/src/trie/proof.rs +++ b/src/trie/proof.rs @@ -64,16 +64,16 @@ impl MultiProof { &'b self, root: Felt, key_values: impl IntoIterator, Felt)> + 'a, + tree_height: u8, ) -> impl Iterator + 'a { let mut checked_cache: HashSet = Default::default(); let mut current_path = BitVec::with_capacity(251); key_values.into_iter().map(move |(k, v)| { let k = k.as_ref(); - // todo: find a way to disable this check in non-251bit key tests. - // if key.len() != 251 { - // return Err(BonsaiStorageError::CreateProof(format!("Key {key:b} is not the correct length."))); - // } + if k.len() != tree_height as _ { + return Membership::NonMember; + } // Go down the tree, starting from the root. current_path.clear(); // hoisted alloc @@ -145,6 +145,8 @@ impl MerkleTree { db: &KeyValueDB, keys: impl IntoIterator>, ) -> Result> { + let max_height = self.max_height; + struct ProofVisitor(MultiProof, PhantomData); impl NodeVisitor for ProofVisitor { fn visit_node( @@ -179,10 +181,12 @@ impl MerkleTree { let mut iter = self.iter(db); for key in keys { let key = key.as_ref(); - // todo: find a way to disable this check in non-251bit key tests. - // if key.len() != 251 { - // return Err(BonsaiStorageError::CreateProof(format!("Key {key:b} is not the correct length."))); - // } + if key.len() != max_height as _ { + return Err(BonsaiStorageError::KeyLength { + expected: self.max_height as _, + got: key.len(), + }); + } iter.traverse_to(&mut visitor, key)?; // We should have found a leaf here. iter.leaf_hash.ok_or_else(|| { @@ -218,6 +222,7 @@ mod tests { let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( RocksDB::::new(&db, RocksDBConfig::default()), BonsaiStorageConfig::default(), + 8, ) .unwrap(); @@ -256,7 +261,8 @@ mod tests { assert!(proof .verify_proof::( tree.root_hash(&bonsai_storage.tries.db).unwrap(), - [(bits![u8, Msb0; 0,0,0,1,0,0,0,0], ONE)] + [(bits![u8, Msb0; 0,0,0,1,0,0,0,0], ONE)], + 8 ) .all(|v| v.into())); } diff --git a/src/trie/tree.rs b/src/trie/tree.rs index 43da376..e824658 100644 --- a/src/trie/tree.rs +++ b/src/trie/tree.rs @@ -65,6 +65,8 @@ pub struct MerkleTree { pub(crate) death_row: HashSet, /// The list of leaves that have been modified during the current commit. pub(crate) cache_leaf_modified: HashMap>, + /// The maximum height of the tree. This is an u8 because we may rely on the fact that it's less than 256 in the future for optimizations. + pub(crate) max_height: u8, /// The hasher used to hash the nodes. _hasher: PhantomData, } @@ -86,6 +88,7 @@ impl fmt::Debug for MerkleTree { impl Clone for MerkleTree { fn clone(&self) -> Self { Self { + max_height: self.max_height, root_node: self.root_node, nodes: self.nodes.clone(), identifier: self.identifier.clone(), @@ -107,13 +110,14 @@ enum NodeOrFelt<'a> { } impl MerkleTree { - pub fn new(identifier: ByteVec) -> Self { + pub fn new(identifier: ByteVec, max_height: u8) -> Self { Self { root_node: None, nodes: Default::default(), identifier, death_row: HashSet::new(), cache_leaf_modified: HashMap::new(), + max_height, _hasher: PhantomData, } } @@ -568,6 +572,12 @@ impl MerkleTree { if value == Felt::ZERO { return self.delete_leaf(db, key); } + if key.len() != self.max_height as _ { + return Err(BonsaiStorageError::KeyLength { + expected: self.max_height as _, + got: key.len(), + }); + } let key_bytes = bitslice_to_bytes(key); log::trace!("key_bytes: {:?}", key_bytes); @@ -759,6 +769,12 @@ impl MerkleTree { db: &KeyValueDB, key: &BitSlice, ) -> Result<(), BonsaiStorageError> { + if key.len() != self.max_height as _ { + return Err(BonsaiStorageError::KeyLength { + expected: self.max_height as _, + got: key.len(), + }); + } log::trace!("delete leaf"); // Algorithm explanation: // diff --git a/src/trie/trees.rs b/src/trie/trees.rs index 25cbf45..b6b5a81 100644 --- a/src/trie/trees.rs +++ b/src/trie/trees.rs @@ -9,6 +9,7 @@ use starknet_types_core::{felt::Felt, hash::StarkHash}; pub(crate) struct MerkleTrees { pub db: KeyValueDB, pub trees: HashMap>, + pub max_height: u8, } impl fmt::Debug @@ -30,15 +31,17 @@ impl Clone Self { db: self.db.clone(), trees: self.trees.clone(), + max_height: self.max_height, } } } impl MerkleTrees { - pub(crate) fn new(db: KeyValueDB) -> Self { + pub(crate) fn new(db: KeyValueDB, tree_height: u8) -> Self { Self { db, trees: HashMap::new(), + max_height: tree_height, } } @@ -51,7 +54,7 @@ impl MerkleTrees MerkleTrees::new(identifier.into()).get(&self.db, key) + MerkleTree::::new(identifier.into(), self.max_height).get(&self.db, key) } } @@ -77,7 +80,7 @@ impl MerkleTrees::new(identifier.into()).get_at(&self.db, key, id) + MerkleTree::::new(identifier.into(), self.max_height).get_at(&self.db, key, id) } } @@ -89,7 +92,7 @@ impl MerkleTrees::new(identifier.into()).contains(&self.db, key) + MerkleTree::::new(identifier.into(), self.max_height).contains(&self.db, key) } } @@ -124,7 +127,7 @@ impl MerkleTrees::new(identifier.into()).root_hash(&self.db) + MerkleTree::::new(identifier.into(), self.max_height).root_hash(&self.db) } } @@ -237,7 +240,7 @@ impl MerkleTrees