Sparse Merkle Tree (#123)

This PR will change how the sparse merkle tree is optimized. Currently
it implements certain optimizations that are beneficial in some regard,
but there are trade-offs. [Other
literature](https://github.com/google/trillian/blob/master/docs/papers/RevocationTransparency.pdf)
references this other optimization, which after discussion, we find has
benefits that we would like. Specifically, it will come in handy with
respect to the implementation of [this
issue](#66). [These
notes](https://hackmd.io/6RI9ZyGqR1ykRITggTb-pg) may be helpful for
understanding the refactor.

---------

Signed-off-by: Daniel Macovei <[email protected]>
Co-authored-by: Kyle Brown <[email protected]>

crates/client/src/api.rs

@@ -401,7 +401,7 @@ impl Client {
             }
         }
 
-        let map_proof_bundle: MapProofBundle<Sha256, MapLeaf> =
+        let map_proof_bundle: MapProofBundle<Sha256, LogId, MapLeaf> =
             MapProofBundle::decode(response.map.as_slice())?;
         let map_inclusions = map_proof_bundle.unbundle();
         for (leaf, proof) in leafs.iter().zip(map_inclusions.iter()) {

crates/client/src/lib.rs

@@ -366,7 +366,6 @@ impl<R: RegistryStorage, C: ContentStorage> Client<R, C> {
             })
             .inspect(|(_, p)| tracing::info!("package `{id}` will be updated", id = p.id))
             .collect::<HashMap<_, _>>();
-
         if packages.is_empty() {
             return Ok(());
         }

crates/crypto/Cargo.toml

@@ -18,6 +18,7 @@ secrecy = { workspace = true }
 signature = { workspace = true }
 thiserror = { workspace = true }
 anyhow = { workspace = true }
+once_cell.workspace = true
 
 [dev-dependencies]
 pretty_assertions = { workspace = true }

crates/crypto/src/hash/mod.rs

@@ -1,4 +1,5 @@
 use anyhow::Error;
+use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 use std::{fmt, str::FromStr};
 use thiserror::Error;
@@ -11,6 +12,8 @@ pub use dynamic::{AnyHash, AnyHashError};
 pub use r#static::Hash;
 pub use sha2::Sha256;
 
+use crate::VisitBytes;
+
 use self::r#static::IncorrectLengthError;
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
@@ -39,12 +42,54 @@ impl FromStr for HashAlgorithm {
     }
 }
 
-pub trait SupportedDigest: Digest + private::Sealed {
+static EMPTY_TREE_HASH: Lazy<Vec<Hash<Sha256>>> = Lazy::new(|| {
+    let mut v: Vec<Hash<Sha256>> = Vec::with_capacity(257);
+    fn empty_tree_hash<D: SupportedDigest>(v: &mut Vec<Hash<D>>, height: u32) -> Hash<D> {
+        let hash: Hash<D> = if height == 0 {
+            hash_empty()
+        } else {
+            let last_hash = empty_tree_hash(v, height - 1);
+            hash_branch(&last_hash, &last_hash)
+        };
+        v.push(hash.clone());
+        hash
+    }
+    empty_tree_hash(&mut v, 256);
+    v
+});
+
+// If updating this function, also update `hash_empty` in transparency map
+pub(crate) fn hash_empty<D: SupportedDigest>() -> Hash<D> {
+    hash_leaf(())
+}
+
+// If updating this function, also update `hash_leaf` in transparency map
+pub(crate) fn hash_leaf<D, V>(value: V) -> Hash<D>
+where
+    D: SupportedDigest,
+    V: VisitBytes,
+{
+    Hash::of(&(0b0, value))
+}
+
+// If updating this function, also update `hash_branch` in transparency map
+pub(crate) fn hash_branch<D>(lhs: &Hash<D>, rhs: &Hash<D>) -> Hash<D>
+where
+    D: SupportedDigest,
+{
+    Hash::of((0b1, lhs, rhs))
+}
+
+pub trait SupportedDigest: Digest + private::Sealed + Sized + 'static {
     const ALGORITHM: HashAlgorithm;
+    fn empty_tree_hash(height: usize) -> &'static Hash<Self>;
 }
 
 impl SupportedDigest for Sha256 {
     const ALGORITHM: HashAlgorithm = HashAlgorithm::Sha256;
+    fn empty_tree_hash(height: usize) -> &'static Hash<Sha256> {
+        &EMPTY_TREE_HASH[height]
+    }
 }
 
 mod private {

crates/crypto/src/hash/static.rs

@@ -54,6 +54,10 @@ impl<D: SupportedDigest> Hash<D> {
     pub fn len(&self) -> usize {
         self.bytes().len()
     }
+
+    pub fn bit_len(&self) -> usize {
+        self.bytes().len() * 8
+    }
 }
 
 impl<D: SupportedDigest> VisitBytes for Hash<D> {

crates/protocol/src/registry.rs

@@ -245,3 +245,21 @@ impl From<AnyHash> for RecordId {
         Self(value)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use warg_crypto::hash::Sha256;
+    use warg_transparency::map::Map;
+
+    #[test]
+    fn log_id() {
+        let first = Map::<Sha256, LogId, &'static str>::default();
+        let second = first.insert(LogId::operator_log::<Sha256>(), "foobar");
+        let proof = second.prove(LogId::operator_log::<Sha256>()).unwrap();
+        assert_eq!(
+            second.root().clone(),
+            proof.evaluate(&LogId::operator_log::<Sha256>(), &"foobar")
+        );
+    }
+}

crates/server/src/services/core.rs

@@ -203,7 +203,6 @@ impl CoreService {
         while let Some(res) = initial.next().await {
             let InitialLeaf { leaf, checkpoint } = res?;
             data.log.push(&leaf);
-
             data.map = data.map.insert(
                 leaf.log_id.clone(),
                 MapLeaf {

crates/server/src/services/data/map.rs

@@ -7,7 +7,7 @@ use tokio::{
 };
 use tokio_util::sync::CancellationToken;
 use warg_crypto::hash::{Hash, Sha256};
-use warg_protocol::registry::{LogLeaf, MapLeaf};
+use warg_protocol::registry::{LogId, LogLeaf, MapLeaf};
 use warg_transparency::map::MapProofBundle;
 
 pub struct Input {
@@ -35,16 +35,15 @@ impl MapData {
         &self,
         root: &Hash<Sha256>,
         leaves: &[LogLeaf],
-    ) -> Result<MapProofBundle<Sha256, MapLeaf>, DataServiceError> {
+    ) -> Result<MapProofBundle<Sha256, LogId, MapLeaf>, DataServiceError> {
         let map = self
             .map_index
             .get(root)
             .ok_or_else(|| DataServiceError::RootNotFound(root.clone()))?;
-
         let mut proofs = Vec::new();
         for LogLeaf { log_id, record_id } in leaves {
             let proof = map
-                .prove(log_id)
+                .prove(log_id.clone())
                 .ok_or_else(|| DataServiceError::PackageNotIncluded(log_id.clone()))?;
             let leaf = MapLeaf {
                 record_id: record_id.clone(),

crates/server/src/services/transparency/map.rs

@@ -49,7 +49,8 @@ pub fn spawn(input: Input) -> Output {
                 summary = log_summary_rx.recv() => {
                     if let Some(summary) = summary {
                         let leaf = summary.leaf;
-                        map = map.insert(leaf.log_id.clone(), MapLeaf { record_id: leaf.record_id.clone() });
+                        let hash = leaf.log_id.clone();
+                        map = map.insert(hash, MapLeaf { record_id: leaf.record_id.clone() });
                         leaves.push(leaf);
 
                         current = Some(MapCheckpoint {

crates/transparency/src/map/fork.rs

@@ -6,16 +6,21 @@ use warg_crypto::hash::{Hash, SupportedDigest};
 
 use super::{link::Link, map::hash_branch, path::Side};
 
+#[derive(Debug)]
 pub struct Fork<D: SupportedDigest> {
-    left: Option<Arc<Link<D>>>,
-    right: Option<Arc<Link<D>>>,
+    left: Arc<Link<D>>,
+    right: Arc<Link<D>>,
 }
 
 impl<D: SupportedDigest> Fork<D> {
+    pub fn new(left: Arc<Link<D>>, right: Arc<Link<D>>) -> Self {
+        Self { left, right }
+    }
+
     pub fn hash(&self) -> Hash<D> {
-        let lhs = self.left.as_ref().map(|left| left.hash().clone());
-        let rhs = self.right.as_ref().map(|right| right.hash().clone());
-        hash_branch(lhs, rhs)
+        let lhs = self.left.as_ref().hash().clone();
+        let rhs = self.right.as_ref().hash().clone();
+        hash_branch(&lhs, &rhs)
     }
 }
 
@@ -28,17 +33,8 @@ impl<D: SupportedDigest> Clone for Fork<D> {
     }
 }
 
-impl<D: SupportedDigest> Default for Fork<D> {
-    fn default() -> Self {
-        Self {
-            left: None,
-            right: None,
-        }
-    }
-}
-
 impl<D: SupportedDigest> Index<Side> for Fork<D> {
-    type Output = Option<Arc<Link<D>>>;
+    type Output = Arc<Link<D>>;
 
     fn index(&self, index: Side) -> &Self::Output {
         match index {

crates/transparency/src/map/link.rs

@@ -1,17 +1,20 @@
+use std::sync::Arc;
+
 use warg_crypto::hash::{Hash, SupportedDigest};
 
 use super::node::Node;
 
+#[derive(Debug)]
 pub struct Link<D: SupportedDigest> {
     hash: Hash<D>,
-    node: Node<D>,
+    node: Arc<Node<D>>,
 }
 
 impl<D: SupportedDigest> Link<D> {
     pub fn new(node: Node<D>) -> Self {
         Self {
             hash: node.hash(),
-            node,
+            node: Arc::new(node),
         }
     }
 
@@ -32,14 +35,3 @@ impl<D: SupportedDigest> Clone for Link<D> {
         }
     }
 }
-
-impl<D: SupportedDigest> Default for Link<D> {
-    fn default() -> Self {
-        let node = Node::default();
-
-        Link {
-            hash: node.hash(),
-            node,
-        }
-    }
-}