tlsn_core/

merkle.rs

//! Merkle tree types.

use std::collections::HashSet;

use serde::{Deserialize, Serialize};

use crate::hash::{Hash, HashAlgId, HashAlgorithm, TypedHash};

/// Merkle tree error.
#[derive(Debug, thiserror::Error)]
#[error("merkle error: {0}")]
pub struct MerkleError(String);

impl MerkleError {
    fn new(msg: impl Into<String>) -> Self {
        Self(msg.into())
    }
}

/// Merkle proof.
#[derive(Clone, Serialize, Deserialize)]
pub struct MerkleProof {
    alg: HashAlgId,
    leaf_count: usize,
    proof: rs_merkle::MerkleProof<Hash>,
}

opaque_debug::implement!(MerkleProof);

impl MerkleProof {
    /// Checks if the counts of indices, hashes, and leaves are valid for the
    /// provided root.
    ///
    /// # Panics
    ///
    /// - If the length of `leaf_indices` and `leaf_hashes` does not match.
    /// - If `leaf_indices` contains duplicates.
    pub fn verify(
        &self,
        hasher: &dyn HashAlgorithm,
        root: &TypedHash,
        leaves: impl IntoIterator<Item = (usize, Hash)>,
    ) -> Result<(), MerkleError> {
        let mut leaves = leaves.into_iter().collect::<Vec<_>>();

        // Sort by index
        leaves.sort_by_key(|(idx, _)| *idx);

        let (indices, leaves): (Vec<_>, Vec<_>) = leaves.into_iter().unzip();

        if indices.len() != indices.iter().collect::<HashSet<_>>().len() {
            return Err(MerkleError::new("duplicate leaf indices provided"));
        }

        if !self.proof.verify(
            &RsMerkleHasher(hasher),
            root.value,
            &indices,
            &leaves,
            self.leaf_count,
        ) {
            return Err(MerkleError::new("invalid merkle proof"));
        }

        Ok(())
    }
}

#[derive(Clone)]
struct RsMerkleHasher<'a>(&'a dyn HashAlgorithm);

impl rs_merkle::Hasher for RsMerkleHasher<'_> {
    type Hash = Hash;

    fn hash(&self, data: &[u8]) -> Hash {
        self.0.hash(data)
    }
}

/// Merkle tree.
#[derive(Clone, Serialize, Deserialize)]
pub struct MerkleTree {
    alg: HashAlgId,
    tree: rs_merkle::MerkleTree<Hash>,
}

impl MerkleTree {
    /// Creates a new Merkle tree.
    pub fn new(alg: HashAlgId) -> Self {
        Self {
            alg,
            tree: Default::default(),
        }
    }

    /// Returns the hash algorithm used to create the tree.
    pub fn algorithm(&self) -> HashAlgId {
        self.alg
    }

    /// Returns the root of the tree.
    pub fn root(&self) -> TypedHash {
        TypedHash {
            alg: self.alg,
            value: self.tree.root().expect("tree should not be empty"),
        }
    }

    /// Inserts leaves into the tree.
    ///
    /// # Panics
    ///
    /// - If the provided hasher is not the same as the one used to create the
    ///   tree.
    pub fn insert(&mut self, hasher: &dyn HashAlgorithm, mut leaves: Vec<Hash>) {
        assert_eq!(self.alg, hasher.id(), "hash algorithm mismatch");

        self.tree.append(&mut leaves);
        self.tree.commit(&RsMerkleHasher(hasher))
    }

    /// Returns a Merkle proof for the provided indices.
    ///
    /// # Panics
    ///
    /// - If the provided indices are not unique and sorted.
    /// - If the provided indices are out of bounds.
    pub fn proof(&self, indices: &[usize]) -> MerkleProof {
        assert!(
            indices.windows(2).all(|w| w[0] < w[1]),
            "indices must be unique and sorted"
        );

        assert!(
            *indices.last().unwrap() < self.tree.leaves_len(),
            "one or more provided indices are out of bounds"
        );

        MerkleProof {
            alg: self.alg,
            leaf_count: self.tree.leaves_len(),
            proof: self.tree.proof(indices),
        }
    }
}

#[cfg(test)]
mod test {
    use crate::hash::{Blake3, Keccak256, Sha256};

    use super::*;
    use rstest::*;

    #[derive(Serialize)]
    struct T(u64);

    fn leaves<H: HashAlgorithm>(hasher: &H, leaves: impl IntoIterator<Item = T>) -> Vec<Hash> {
        leaves
            .into_iter()
            .map(|x| hasher.hash(&x.0.to_be_bytes()))
            .collect()
    }

    fn choose_leaves(
        indices: impl IntoIterator<Item = usize>,
        leaves: &[Hash],
    ) -> Vec<(usize, Hash)> {
        indices.into_iter().map(|i| (i, leaves[i])).collect()
    }

    // Expect Merkle proof verification to succeed.
    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    fn test_verify_success<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        let proof = tree.proof(&[2, 3, 4]);

        assert!(proof
            .verify(&hasher, &tree.root(), choose_leaves([2, 3, 4], &leaves))
            .is_ok());
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    fn test_verify_fail_wrong_leaf<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        let proof = tree.proof(&[2, 3, 4]);

        let mut choices = choose_leaves([2, 3, 4], &leaves);

        choices[1].1 = leaves[0];

        // Fail because the leaf is wrong.
        assert!(proof.verify(&hasher, &tree.root(), choices).is_err());
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    #[should_panic]
    fn test_proof_fail_length_unsorted<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        _ = tree.proof(&[2, 4, 3]);
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    #[should_panic]
    fn test_proof_fail_index_out_of_bounds<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        _ = tree.proof(&[2, 3, 4, 6]);
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    #[should_panic]
    fn test_proof_fail_length_duplicates<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        _ = tree.proof(&[2, 2, 3]);
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    fn test_verify_fail_duplicates<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        let proof = tree.proof(&[2, 3, 4]);

        assert!(proof
            .verify(&hasher, &tree.root(), choose_leaves([2, 2, 3], &leaves))
            .is_err());
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    fn test_verify_fail_incorrect_leaf_count<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        let mut proof1 = tree.proof(&[2, 3, 4]);
        let mut proof2 = proof1.clone();

        // Increment the `leaf_count` field.
        proof1.leaf_count += 1;
        // Decrement the `leaf_count` field.
        proof2.leaf_count -= 1;

        // Fail because leaf count is wrong.
        assert!(proof1
            .verify(&hasher, &tree.root(), choose_leaves([2, 3, 4], &leaves))
            .is_err());

        assert!(proof2
            .verify(&hasher, &tree.root(), choose_leaves([2, 3, 4], &leaves))
            .is_err());
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    fn test_verify_fail_incorrect_indices<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        let proof = tree.proof(&[2, 3, 4]);

        let mut choices = choose_leaves([2, 3, 4], &leaves);
        choices[1].0 = 1;

        // Fail because leaf index is wrong.
        assert!(proof.verify(&hasher, &tree.root(), choices).is_err());
    }

    #[rstest]
    #[case::sha2(Sha256::default())]
    #[case::blake3(Blake3::default())]
    #[case::keccak(Keccak256::default())]
    fn test_verify_fail_fewer_indices<H: HashAlgorithm>(#[case] hasher: H) {
        let mut tree = MerkleTree::new(hasher.id());

        let leaves = leaves(&hasher, [T(0), T(1), T(2), T(3), T(4)]);

        tree.insert(&hasher, leaves.clone());

        let proof = tree.proof(&[2, 3, 4]);

        // Trying to verify less leaves than what was included in the proof.
        assert!(proof
            .verify(&hasher, &tree.root(), choose_leaves([2, 3], &leaves))
            .is_err());
    }
}