fix(root_tree): non-deterministic root hash (#63)

dashpay · Feb 8, 2022 · 30bed4c · 30bed4c
1 parent cedcfdc
commit 30bed4c
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 13 deletions.
diff --git a/grovedb/src/lib.rs b/grovedb/src/lib.rs
@@ -6,7 +6,7 @@ mod tests;
 
 use std::{
     cell::RefCell,
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     path::Path,
     rc::Rc,
 };
@@ -103,22 +103,22 @@ pub struct Proof {
 
 pub struct GroveDb {
     root_tree: MerkleTree<Sha256>,
-    root_leaf_keys: HashMap<Vec<u8>, usize>,
+    root_leaf_keys: BTreeMap<Vec<u8>, usize>,
     meta_storage: PrefixedRocksDbStorage,
     db: Rc<storage::rocksdb_storage::OptimisticTransactionDB>,
     // Locks the database for writes during the transaction
     is_readonly: bool,
     // Temp trees used for writes during transaction
     temp_root_tree: MerkleTree<Sha256>,
-    temp_root_leaf_keys: HashMap<Vec<u8>, usize>,
+    temp_root_leaf_keys: BTreeMap<Vec<u8>, usize>,
     temp_subtrees: RefCell<HashMap<Vec<u8>, Merk<PrefixedRocksDbStorage>>>,
     temp_deleted_subtrees: RefCell<HashSet<Vec<u8>>>,
 }
 
 impl GroveDb {
     pub fn new(
         root_tree: MerkleTree<Sha256>,
-        root_leaf_keys: HashMap<Vec<u8>, usize>,
+        root_leaf_keys: BTreeMap<Vec<u8>, usize>,
         meta_storage: PrefixedRocksDbStorage,
         db: Rc<storage::rocksdb_storage::OptimisticTransactionDB>,
     ) -> Self {
@@ -128,7 +128,7 @@ impl GroveDb {
             meta_storage,
             db,
             temp_root_tree: MerkleTree::new(),
-            temp_root_leaf_keys: HashMap::new(),
+            temp_root_leaf_keys: BTreeMap::new(),
             temp_subtrees: RefCell::new(HashMap::new()),
             temp_deleted_subtrees: RefCell::new(HashSet::new()),
             is_readonly: false,
@@ -147,14 +147,14 @@ impl GroveDb {
         let meta_storage = PrefixedRocksDbStorage::new(db.clone(), Vec::new())?;
 
         // TODO: owned `get` is not required for deserialization
-        let root_leaf_keys: HashMap<Vec<u8>, usize> = if let Some(root_leaf_keys_serialized) =
+        let root_leaf_keys: BTreeMap<Vec<u8>, usize> = if let Some(root_leaf_keys_serialized) =
             meta_storage.get_meta(ROOT_LEAFS_SERIALIZED_KEY)?
         {
             bincode::deserialize(&root_leaf_keys_serialized).map_err(|_| {
                 Error::CorruptedData(String::from("unable to deserialize root leafs"))
             })?
         } else {
-            HashMap::new()
+            BTreeMap::new()
         };
 
         let temp_subtrees: RefCell<HashMap<Vec<u8>, Merk<PrefixedRocksDbStorage>>> =
@@ -199,7 +199,7 @@ impl GroveDb {
 
     fn build_root_tree(
         subtrees: &Subtrees,
-        root_leaf_keys: &HashMap<Vec<u8>, usize>,
+        root_leaf_keys: &BTreeMap<Vec<u8>, usize>,
         transaction: Option<&OptimisticTransactionDBTransaction>,
     ) -> MerkleTree<Sha256> {
         let mut leaf_hashes: Vec<[u8; 32]> = vec![[0; 32]; root_leaf_keys.len()];
@@ -388,7 +388,9 @@ impl GroveDb {
         // TODO: root tree actually does support transactions, so this
         //  code can be reworked to account for that
         self.root_tree = self.temp_root_tree.clone();
-        self.root_leaf_keys = self.temp_root_leaf_keys.drain().collect();
+
+        self.root_leaf_keys = self.temp_root_leaf_keys.clone();
+        self.temp_root_leaf_keys.clear();
 
         self.cleanup_transactional_data();
 
@@ -434,7 +436,7 @@ impl GroveDb {
 
         // Free transactional data
         self.temp_root_tree = MerkleTree::new();
-        self.temp_root_leaf_keys = HashMap::new();
+        self.temp_root_leaf_keys = BTreeMap::new();
         self.temp_subtrees = RefCell::new(HashMap::new());
     }
 }
diff --git a/grovedb/src/subtrees.rs b/grovedb/src/subtrees.rs
@@ -1,7 +1,7 @@
 //! Module for retrieving subtrees
 use std::{
     cell::RefCell,
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     rc::Rc,
 };
 
@@ -12,7 +12,7 @@ use crate::{Element, Error, GroveDb};
 
 // TODO: should take temp_root_leaf_keys also
 pub struct Subtrees<'a> {
-    pub root_leaf_keys: &'a HashMap<Vec<u8>, usize>,
+    pub root_leaf_keys: &'a BTreeMap<Vec<u8>, usize>,
     pub temp_subtrees: &'a RefCell<HashMap<Vec<u8>, Merk<PrefixedRocksDbStorage>>>,
     pub deleted_subtrees: &'a RefCell<HashSet<Vec<u8>>>,
     pub storage: Rc<storage::rocksdb_storage::OptimisticTransactionDB>,

diff --git a/grovedb/src/tests.rs b/grovedb/src/tests.rs
@@ -235,7 +235,7 @@ fn test_tree_structure_is_persistent() {
 #[test]
 fn test_root_tree_leafs_are_noted() {
     let db = make_grovedb();
-    let mut hm = HashMap::new();
+    let mut hm = BTreeMap::new();
     hm.insert(TEST_LEAF.to_vec(), 0);
     hm.insert(ANOTHER_TEST_LEAF.to_vec(), 1);
     assert_eq!(db.root_leaf_keys, hm);