refactor: Add Prefix enum and buffer error checking (#131)

* Prefix enum

* Simplify

* Revert accidental change

* Clippy

* impl AsRef

* Update prefix.rs

* Error handling

* Error checks

* Collapse use statements

* Manual impl From StorageError

* Manual from prefix error

* Update node.rs

* Revert renaming Internal to Node

* Update binary_proofs.rs

* Revert "Update binary_proofs.rs"

This reverts commit d9996de6072072d73e6b5c0fd7d6fcfe939cf625.

* Update node.rs

fuel-merkle/src/binary/hash.rs

@@ -1,4 +1,4 @@
-use crate::common::{self, Bytes32, LEAF, NODE};
+use crate::common::{empty_sum_sha256, Bytes32, Prefix};
 
 use digest::Digest;
 use sha2::Sha256;
@@ -8,15 +8,15 @@ type Hash = Sha256;
 // Merkle Tree hash of an empty list
 // MTH({}) = Hash()
 pub const fn empty_sum() -> &'static Bytes32 {
-    common::empty_sum_sha256()
+    empty_sum_sha256()
 }
 
 // Merkle tree hash of an n-element list D[n]
 // MTH(D[n]) = Hash(0x01 || MTH(D[0:k]) || MTH(D[k:n])
 pub fn node_sum(lhs_data: &[u8], rhs_data: &[u8]) -> Bytes32 {
     let mut hash = Hash::new();
 
-    hash.update([NODE]);
+    hash.update(Prefix::Node);
     hash.update(lhs_data);
     hash.update(rhs_data);
 
@@ -28,7 +28,7 @@ pub fn node_sum(lhs_data: &[u8], rhs_data: &[u8]) -> Bytes32 {
 pub fn leaf_sum(data: &[u8]) -> Bytes32 {
     let mut hash = Hash::new();
 
-    hash.update([LEAF]);
+    hash.update(Prefix::Leaf);
     hash.update(data);
 
     hash.finalize().into()

fuel-merkle/src/binary/merkle_tree.rs

@@ -1,10 +1,11 @@
-use crate::binary::{self, Node};
-use crate::common::{Bytes32, Position, ProofSet, Subtree};
+use crate::{
+    binary::{empty_sum, Node},
+    common::{Bytes32, Position, ProofSet, Subtree},
+};
 
 use fuel_storage::{Mappable, StorageMutate};
 
-use alloc::boxed::Box;
-use alloc::vec::Vec;
+use alloc::{boxed::Box, vec::Vec};
 use core::fmt;
 
 #[derive(Debug, Clone)]
@@ -19,7 +20,7 @@ pub enum MerkleTreeError<StorageError> {
     )]
     LoadError(u64),
 
-    #[cfg_attr(feature = "std", error("a storage error was thrown: {0}"))]
+    #[cfg_attr(feature = "std", error(transparent))]
     StorageError(StorageError),
 }
 
@@ -29,6 +30,7 @@ impl<StorageError> From<StorageError> for MerkleTreeError<StorageError> {
     }
 }
 
+#[derive(Debug)]
 pub struct MerkleTree<StorageType> {
     storage: StorageType,
     head: Option<Box<Subtree<Node>>>,
@@ -38,6 +40,7 @@ pub struct MerkleTree<StorageType> {
 /// The table of the Binary Merkle Tree's nodes. [`MerkleTree`] works with it as
 /// a binary array, where the storage key of the node is the `u64` index and
 /// value is the [`Node`](crate::binary::Node).
+#[derive(Debug)]
 pub struct NodesTable;
 
 impl Mappable for NodesTable {
@@ -78,7 +81,7 @@ where
     pub fn root(&mut self) -> Result<Bytes32, MerkleTreeError<StorageError>> {
         let root_node = self.root_node()?;
         let root = match root_node {
-            None => *binary::empty_sum(),
+            None => *empty_sum(),
             Some(ref node) => *node.hash(),
         };
 
@@ -312,13 +315,16 @@ where
 
 #[cfg(test)]
 mod test {
-    use super::{MerkleTree, NodesTable};
-    use crate::binary::{empty_sum, leaf_sum, node_sum};
-    use crate::common::StorageMap;
-    use alloc::vec::Vec;
+    use super::{MerkleTree, MerkleTreeError, NodesTable};
+    use crate::{
+        binary::{empty_sum, leaf_sum, node_sum},
+        common::StorageMap,
+    };
     use fuel_merkle_test_helpers::TEST_DATA;
     use fuel_storage::StorageInspect;
 
+    use alloc::vec::Vec;
+
     #[test]
     fn test_push_builds_internal_tree_structure() {
         let mut storage_map = StorageMap::<NodesTable>::new();
@@ -435,8 +441,9 @@ mod test {
             let _ = tree.push(datum);
         }
 
-        let tree = MerkleTree::load(&mut storage_map, LEAVES_COUNT * 2);
-        assert!(tree.is_err());
+        let err = MerkleTree::load(&mut storage_map, LEAVES_COUNT * 2)
+            .expect_err("Expected load() to return Error; got Ok");
+        assert!(matches!(err, MerkleTreeError::LoadError(_)));
     }
 
     #[test]
@@ -513,8 +520,10 @@ mod test {
         let mut storage_map = StorageMap::new();
         let mut tree = MerkleTree::new(&mut storage_map);
 
-        let proof = tree.prove(0);
-        assert!(proof.is_err());
+        let err = tree
+            .prove(0)
+            .expect_err("Expected prove() to return Error; got Ok");
+        assert!(matches!(err, MerkleTreeError::InvalidProofIndex(0)));
     }
 
     #[test]
@@ -527,8 +536,10 @@ mod test {
             let _ = tree.push(datum);
         }
 
-        let proof = tree.prove(10);
-        assert!(proof.is_err());
+        let err = tree
+            .prove(10)
+            .expect_err("Expected prove() to return Error; got Ok");
+        assert!(matches!(err, MerkleTreeError::InvalidProofIndex(10)));
     }
 
     #[test]

fuel-merkle/src/common.rs

@@ -3,9 +3,12 @@ mod node;
 mod path_iterator;
 mod position;
 mod position_path;
+mod prefix;
 mod storage_map;
 mod subtree;
 
+pub(crate) mod error;
+
 pub use msb::Msb;
 pub use node::{Node, ParentNode};
 pub use path_iterator::AsPathIterator;
@@ -14,9 +17,7 @@ pub use storage_map::StorageMap;
 pub use subtree::Subtree;
 
 pub(crate) use position_path::PositionPath;
-
-pub const NODE: u8 = 0x01;
-pub const LEAF: u8 = 0x00;
+pub(crate) use prefix::{Prefix, PrefixError};
 
 pub type Bytes1 = [u8; 1];
 pub type Bytes2 = [u8; 2];

fuel-merkle/src/common/error.rs

@@ -0,0 +1,14 @@
+use crate::common::PrefixError;
+
+#[derive(Debug, Clone)]
+#[cfg_attr(feature = "std", derive(thiserror::Error))]
+pub enum DeserializeError {
+    #[cfg_attr(feature = "std", error(transparent))]
+    PrefixError(PrefixError),
+}
+
+impl From<PrefixError> for DeserializeError {
+    fn from(err: PrefixError) -> Self {
+        DeserializeError::PrefixError(err)
+    }
+}

fuel-merkle/src/common/prefix.rs

@@ -0,0 +1,56 @@
+use crate::common::prefix::PrefixError::InvalidPrefix;
+
+const NODE: u8 = 0x01;
+const LEAF: u8 = 0x00;
+
+#[derive(Debug, Clone)]
+#[cfg_attr(feature = "std", derive(thiserror::Error))]
+pub enum PrefixError {
+    #[cfg_attr(feature = "std", error("prefix {0} is not valid"))]
+    InvalidPrefix(u8),
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum Prefix {
+    Node,
+    Leaf,
+}
+
+impl From<Prefix> for u8 {
+    fn from(prefix: Prefix) -> Self {
+        match prefix {
+            Prefix::Node => NODE,
+            Prefix::Leaf => LEAF,
+        }
+    }
+}
+
+impl TryFrom<u8> for Prefix {
+    type Error = PrefixError;
+
+    fn try_from(byte: u8) -> Result<Self, Self::Error> {
+        match byte {
+            NODE => Ok(Prefix::Node),
+            LEAF => Ok(Prefix::Leaf),
+            _ => Err(InvalidPrefix(byte)),
+        }
+    }
+}
+
+impl AsRef<[u8]> for Prefix {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            Prefix::Node => &[NODE],
+            Prefix::Leaf => &[LEAF],
+        }
+    }
+}
+
+impl AsRef<[u8; 1]> for Prefix {
+    fn as_ref(&self) -> &[u8; 1] {
+        match self {
+            Prefix::Node => &[NODE],
+            Prefix::Leaf => &[LEAF],
+        }
+    }
+}

fuel-merkle/src/sparse/merkle_tree.rs

@@ -1,9 +1,10 @@
-use crate::common::{AsPathIterator, Bytes32};
-use crate::sparse::{zero_sum, Buffer, Node, StorageNode};
+use crate::{
+    common::{error::DeserializeError, AsPathIterator, Bytes32},
+    sparse::{zero_sum, Buffer, Node, StorageNode},
+};
 use fuel_storage::{Mappable, StorageMutate};
 
-use alloc::string::String;
-use alloc::vec::Vec;
+use alloc::{string::String, vec::Vec};
 use core::{cmp, fmt, iter};
 
 #[derive(Debug, Clone)]
@@ -15,8 +16,11 @@ pub enum MerkleTreeError<StorageError> {
     )]
     LoadError(String),
 
-    #[cfg_attr(feature = "std", error("a storage error was thrown: {0}"))]
+    #[cfg_attr(feature = "std", error(transparent))]
     StorageError(StorageError),
+
+    #[cfg_attr(feature = "std", error(transparent))]
+    DeserializeError(DeserializeError),
 }
 
 impl<StorageError> From<StorageError> for MerkleTreeError<StorageError> {
@@ -25,6 +29,7 @@ impl<StorageError> From<StorageError> for MerkleTreeError<StorageError> {
     }
 }
 
+#[derive(Debug)]
 pub struct MerkleTree<StorageType> {
     root_node: Node,
     storage: StorageType,
@@ -33,6 +38,7 @@ pub struct MerkleTree<StorageType> {
 /// The table of the Sparse Merkle tree's nodes. [`MerkleTree`] works with it as a sparse merkle
 /// tree, where the storage key is `Bytes32` and the value is the [`Buffer`](crate::sparse::Buffer)
 /// (raw presentation of the [`Node`](crate::sparse::Node)).
+#[derive(Debug)]
 pub struct NodesTable;
 
 impl Mappable for NodesTable {
@@ -44,7 +50,6 @@ impl Mappable for NodesTable {
     type GetValue = Self::SetValue;
 }
 
-// TODO: Process each `unwrap` as error
 impl<StorageType, StorageError> MerkleTree<StorageType>
 where
     StorageType: StorageMutate<NodesTable, Error = StorageError>,
@@ -66,7 +71,9 @@ where
             .ok_or_else(|| MerkleTreeError::LoadError(hex::encode(root)))?
             .into_owned();
         let tree = Self {
-            root_node: Node::from_buffer(buffer),
+            root_node: buffer
+                .try_into()
+                .map_err(MerkleTreeError::DeserializeError)?,
             storage,
         };
         Ok(tree)
@@ -108,7 +115,10 @@ where
         }
 
         if let Some(buffer) = self.storage.get(key)? {
-            let leaf_node = Node::from_buffer(*buffer);
+            let buffer = buffer.into_owned();
+            let leaf_node: Node = buffer
+                .try_into()
+                .map_err(MerkleTreeError::DeserializeError)?;
             let (path_nodes, side_nodes): (Vec<Node>, Vec<Node>) = self.path_set(leaf_node.clone());
             self.delete_with_path_set(&leaf_node, path_nodes.as_slice(), side_nodes.as_slice())?;
         }
@@ -282,7 +292,7 @@ where
 mod test {
     use crate::common::StorageMap;
     use crate::sparse::hash::sum;
-    use crate::sparse::MerkleTree;
+    use crate::sparse::{MerkleTree, MerkleTreeError};
     use hex;
 
     #[test]
@@ -665,10 +675,32 @@ mod test {
             tree.update(&sum(b"\x00\x00\x00\x04"), b"DATA").unwrap();
         }
 
-        {
-            let root = &sum(b"\xff\xff\xff\xff");
-            let tree = MerkleTree::load(&mut storage, root);
-            assert!(tree.is_err());
-        }
+        let root = &sum(b"\xff\xff\xff\xff");
+        let err = MerkleTree::load(&mut storage, root)
+            .expect_err("Expected load() to return Error; got Ok");
+        assert!(matches!(err, MerkleTreeError::LoadError(_)));
+    }
+
+    #[test]
+    fn test_load_returns_a_deserialize_error_if_the_storage_is_corrupted() {
+        use fuel_storage::StorageMutate;
+
+        let mut storage = StorageMap::new();
+
+        let mut tree = MerkleTree::new(&mut storage);
+        tree.update(&sum(b"\x00\x00\x00\x00"), b"DATA").unwrap();
+        tree.update(&sum(b"\x00\x00\x00\x01"), b"DATA").unwrap();
+        tree.update(&sum(b"\x00\x00\x00\x02"), b"DATA").unwrap();
+        tree.update(&sum(b"\x00\x00\x00\x03"), b"DATA").unwrap();
+        tree.update(&sum(b"\x00\x00\x00\x04"), b"DATA").unwrap();
+        let root = tree.root();
+
+        // Overwrite the root key-value with an invalid buffer to create a
+        // DeserializeError.
+        storage.insert(&root, &[255; 69]).unwrap();
+
+        let err = MerkleTree::load(&mut storage, &root)
+            .expect_err("Expected load() to return Error; got Ok");
+        assert!(matches!(err, MerkleTreeError::DeserializeError(_)));
     }
 }