Implement proof generation and verification

firewood/src/db.rs

@@ -30,7 +30,6 @@ pub enum DbError {
     InvalidParams,
     Merkle(MerkleError),
     System(nix::Error),
-    KeyNotFound,
     CreateError,
     IO(std::io::Error),
     InvalidProposal,
@@ -42,7 +41,6 @@ impl fmt::Display for DbError {
             DbError::InvalidParams => write!(f, "invalid parameters provided"),
             DbError::Merkle(e) => write!(f, "merkle error: {e:?}"),
             DbError::System(e) => write!(f, "system error: {e:?}"),
-            DbError::KeyNotFound => write!(f, "not found"),
             DbError::CreateError => write!(f, "database create error"),
             DbError::IO(e) => write!(f, "I/O error: {e:?}"),
             DbError::InvalidProposal => write!(f, "invalid proposal"),
@@ -79,7 +77,7 @@ impl<T: ReadLinearStore> api::DbView for HistoricalRev<T> {
     async fn single_key_proof<K: api::KeyType>(
         &self,
         _key: K,
-    ) -> Result<Option<Proof<Vec<u8>>>, api::Error> {
+    ) -> Result<Option<Proof>, api::Error> {
         todo!()
     }
 
@@ -124,14 +122,14 @@ impl<T: ReadLinearStore> HistoricalRev<T> {
         todo!()
     }
 
-    pub fn prove(&self, _key: &[u8]) -> Result<Proof<Vec<u8>>, MerkleError> {
+    pub fn prove(&self, _key: &[u8]) -> Result<Proof, MerkleError> {
         todo!()
     }
 
     /// Verifies a range proof is valid for a set of keys.
-    pub fn verify_range_proof<N: AsRef<[u8]> + Send, V: AsRef<[u8]>>(
+    pub fn verify_range_proof<V: AsRef<[u8]>>(
         &self,
-        _proof: Proof<N>,
+        _proof: Proof,
         _first_key: &[u8],
         _last_key: &[u8],
         _keys: Vec<&[u8]>,
@@ -177,7 +175,7 @@ impl<T: WriteLinearStore> api::DbView for Proposal<T> {
         todo!()
     }
 
-    async fn single_key_proof<K>(&self, _key: K) -> Result<Option<Proof<Vec<u8>>>, api::Error>
+    async fn single_key_proof<K>(&self, _key: K) -> Result<Option<Proof>, api::Error>
     where
         K: api::KeyType,
     {

firewood/src/hashednode.rs

@@ -18,6 +18,7 @@ const MAX_VARINT_SIZE: usize = 10;
 const BITS_PER_NIBBLE: u64 = 4;
 
 use crate::merkle::MerkleError;
+use crate::proof::ProofNode;
 use storage::PathIterItem;
 
 /// A [HashedNodeStore] keeps track of nodes as they change when they are backed by a LinearStore.
@@ -283,13 +284,11 @@ pub fn hash_preimage(node: &Node, path_prefix: &Path) -> Box<[u8]> {
             }
             .write(&mut buf);
         }
-        Node::Leaf(node) => {
-            NodeAndPrefix {
-                node,
-                prefix: path_prefix,
-            }
-            .write(&mut buf);
+        Node::Leaf(node) => NodeAndPrefix {
+            node,
+            prefix: path_prefix,
         }
+        .write(&mut buf),
     }
     buf.into_boxed_slice()
 }
@@ -310,43 +309,44 @@ impl HasUpdate for Vec<u8> {
     }
 }
 
-pub(super) enum ValueDigest<'a> {
-    /// A node's value.
-    Value(&'a [u8]),
-    /// The hash of a a node's value.
-    /// TODO danlaine: Use this variant when implement ProofNode
-    _Hash(Box<[u8]>),
+#[derive(Clone, Debug)]
+/// A ValueDigest is either a node's value or the hash of its value.
+pub enum ValueDigest<T> {
+    Value(T),
+    /// TODO this variant will be used when we deserialize a proof node
+    /// from a remote Firewood instance. The serialized proof node they
+    /// send us may the hash of the value, not the value itself.
+    _Hash(T),
 }
 
-trait Hashable {
+pub(crate) trait Hashable {
+    type ValueDigestType: AsRef<[u8]>;
+
     /// The key of the node where each byte is a nibble.
     fn key(&self) -> impl Iterator<Item = u8> + Clone;
     /// The node's value or hash.
-    fn value_digest(&self) -> Option<ValueDigest>;
+    fn value_digest(&self) -> Option<ValueDigest<Self::ValueDigestType>>;
     /// Each element is a child's index and hash.
     /// Yields 0 elements if the node is a leaf.
     fn children(&self) -> impl Iterator<Item = (usize, &TrieHash)> + Clone;
 }
 
 pub(super) trait Preimage {
     /// Returns the hash of this preimage.
-    fn to_hash(self) -> TrieHash;
+    fn to_hash(&self) -> TrieHash;
     /// Write this hash preimage to `buf`.
-    fn write(self, buf: &mut impl HasUpdate);
+    fn write(&self, buf: &mut impl HasUpdate);
 }
 
 // Implement Preimage for all types that implement Hashable
-impl<T> Preimage for T
-where
-    T: Hashable,
-{
-    fn to_hash(self) -> TrieHash {
+impl<T: Hashable> Preimage for T {
+    fn to_hash(&self) -> TrieHash {
         let mut hasher = Sha256::new();
         self.write(&mut hasher);
         hasher.finalize().into()
     }
 
-    fn write(self, buf: &mut impl HasUpdate) {
+    fn write(&self, buf: &mut impl HasUpdate) {
         let children = self.children();
 
         let num_children = children.clone().count() as u64;
@@ -421,6 +421,8 @@ impl<'a, N: HashableNode> From<NodeAndPrefix<'a, N>> for TrieHash {
 }
 
 impl<'a, N: HashableNode> Hashable for NodeAndPrefix<'a, N> {
+    type ValueDigestType = &'a [u8];
+
     fn key(&self) -> impl Iterator<Item = u8> + Clone {
         self.prefix
             .0
@@ -429,7 +431,7 @@ impl<'a, N: HashableNode> Hashable for NodeAndPrefix<'a, N> {
             .chain(self.node.partial_path())
     }
 
-    fn value_digest(&self) -> Option<ValueDigest> {
+    fn value_digest(&self) -> Option<ValueDigest<&'a [u8]>> {
         self.node.value().map(ValueDigest::Value)
     }
 
@@ -438,7 +440,32 @@ impl<'a, N: HashableNode> Hashable for NodeAndPrefix<'a, N> {
     }
 }
 
-fn add_value_digest_to_buf<H: HasUpdate>(buf: &mut H, value_digest: Option<ValueDigest>) {
+impl<'a> Hashable for &'a ProofNode {
+    type ValueDigestType = &'a [u8];
+
+    fn key(&self) -> impl Iterator<Item = u8> + Clone {
+        self.key.as_ref().iter().copied()
+    }
+
+    fn value_digest(&self) -> Option<ValueDigest<&'a [u8]>> {
+        self.value_digest.as_ref().map(|vd| match vd {
+            ValueDigest::Value(v) => ValueDigest::Value(v.as_ref()),
+            ValueDigest::_Hash(h) => ValueDigest::_Hash(h.as_ref()),
+        })
+    }
+
+    fn children(&self) -> impl Iterator<Item = (usize, &TrieHash)> + Clone {
+        self.child_hashes
+            .iter()
+            .enumerate()
+            .filter_map(|(i, hash)| hash.as_ref().map(|h| (i, h)))
+    }
+}
+
+fn add_value_digest_to_buf<H: HasUpdate, T: AsRef<[u8]>>(
+    buf: &mut H,
+    value_digest: Option<ValueDigest<T>>,
+) {
     let Some(value_digest) = value_digest else {
         let value_exists: u8 = 0;
         buf.update([value_exists]);
@@ -451,21 +478,21 @@ fn add_value_digest_to_buf<H: HasUpdate>(buf: &mut H, value_digest: Option<Value
     match value_digest {
         ValueDigest::Value(value) if value.as_ref().len() >= 32 => {
             let hash = Sha256::digest(value);
-            add_len_and_value_to_buf(buf, hash.as_ref());
+            add_len_and_value_to_buf(buf, hash);
         }
         ValueDigest::Value(value) => {
             add_len_and_value_to_buf(buf, value);
         }
         ValueDigest::_Hash(hash) => {
-            add_len_and_value_to_buf(buf, hash.as_ref());
+            add_len_and_value_to_buf(buf, hash);
         }
     }
 }
 
 #[inline]
 /// Writes the length of `value` and `value` to `buf`.
-fn add_len_and_value_to_buf<H: HasUpdate>(buf: &mut H, value: &[u8]) {
-    let value_len = value.len();
+fn add_len_and_value_to_buf<H: HasUpdate, V: AsRef<[u8]>>(buf: &mut H, value: V) {
+    let value_len = value.as_ref().len();
     buf.update([value_len as u8]);
     buf.update(value);
 }