merkle_root/

builder.rs

// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::cmp::min;

use crate::tree::MerkleTree;
use crate::util::{hash_block, hash_hashes, hashes_per_block};
use crate::{Hash, HashAlgorithm, BLOCK_SIZE};

/// A `MerkleTreeBuilder` generates a [`MerkleTree`] from one or more write calls.
///
/// # Examples
/// ```
/// # use merkle_root::*;
/// let data = vec![0xff; 8192];
/// let mut builder = MerkleTreeBuilder::new();
/// for i in 0..8 {
///     builder.write(&data[..]);
/// }
/// assert_eq!(
///     builder.finish().root(),
///     &hex::decode("f75f59a944d2433bc6830ec243bfefa457704d2aed12f30539cd4f18bf1d62cf")
///         .unwrap()
/// );
/// ```
#[derive(Clone, Debug)]
pub struct MerkleTreeBuilder {
    /// Buffer to hold a partial block of data between [`MerkleTreeBuilder::write`] calls.
    /// `block.len()` will never exceed [`BLOCK_SIZE`].
    block: Vec<u8>,
    levels: Vec<Vec<Hash>>,
    algorithm: HashAlgorithm,
}

impl Default for MerkleTreeBuilder {
    fn default() -> Self {
        Self {
            levels: vec![Vec::new()],
            block: Vec::with_capacity(BLOCK_SIZE),
            algorithm: HashAlgorithm::SHA256,
        }
    }
}

impl MerkleTreeBuilder {
    /// Creates a new, empty `MerkleTreeBuilder`.
    pub fn new() -> Self {
        Self::default()
    }

    /// Creates a new, empty `MerkleTreeBuilder` with a specific algorithm.
    pub fn with_algorithm(algorithm: HashAlgorithm) -> Self {
        Self {
            levels: vec![Vec::new()],
            block: Vec::with_capacity(BLOCK_SIZE),
            algorithm,
        }
    }

    /// Append a buffer of bytes to the merkle tree.
    ///
    /// No internal buffering is required if all writes are [`BLOCK_SIZE`] aligned.
    pub fn write(&mut self, buf: &[u8]) {
        // Fill the current partial block, if it exists.
        let buf = if self.block.is_empty() {
            buf
        } else {
            let left = BLOCK_SIZE - self.block.len();
            let prefix = min(buf.len(), left);
            let (buf, rest) = buf.split_at(prefix);
            self.block.extend_from_slice(buf);
            if self.block.len() == BLOCK_SIZE {
                self.push_data_hash(self.hash_block(&self.block[..]));
            }
            rest
        };

        // Write full blocks, saving any final partial block for later writes.
        for block in buf.chunks(BLOCK_SIZE) {
            if block.len() == BLOCK_SIZE {
                self.push_data_hash(self.hash_block(block));
            } else {
                self.block.extend_from_slice(block);
            }
        }
    }

    /// Hash a block of data (level 0), using an offset based on the current number of level 0
    /// hashes.
    fn hash_block(&self, block: &[u8]) -> Hash {
        hash_block(block, self.levels[0].len() * BLOCK_SIZE, self.algorithm)
    }

    /// Save a data block hash, propagating full blocks of hashes to higher layers. Also clear a
    /// stored data block.
    pub fn push_data_hash(&mut self, hash: Hash) {
        self.block.clear();
        self.levels[0].push(hash);
        if self.levels[0].len() % hashes_per_block(self.algorithm) == 0 {
            self.commit_tail_block(0);
        }
    }

    /// Hash a complete (or final partial) block of hashes, chaining to higher levels as needed.
    fn commit_tail_block(&mut self, level: usize) {
        let len = self.levels[level].len();
        let next_level = level + 1;

        if next_level >= self.levels.len() {
            self.levels.push(Vec::new());
        }

        let per_block = hashes_per_block(self.algorithm);
        let first_hash = if len % per_block == 0 {
            len - per_block
        } else {
            len - (len % per_block)
        };

        let hash = hash_hashes(
            &self.levels[level][first_hash..],
            next_level,
            self.levels[next_level].len() * BLOCK_SIZE,
            self.algorithm,
        );

        self.levels[next_level].push(hash);
        if self.levels[next_level].len() % per_block == 0 {
            self.commit_tail_block(next_level);
        }
    }

    /// Finalize all levels of the merkle tree, converting this `MerkleTreeBuilder` instance to a
    /// [`MerkleTree`].
    pub fn finish(mut self) -> MerkleTree {
        // The data protected by the tree may not be BLOCK_SIZE aligned. Commit a partial data
        // block before finalizing the hash levels.
        // Also, an empty tree consists of a single, empty block. Handle that case now as well.
        if !self.block.is_empty() || self.levels[0].is_empty() {
            self.push_data_hash(self.hash_block(&self.block[..]));
        }

        // Enumerate the hash levels, finalizing any that have a partial block of hashes.
        // `commit_tail_block` may add new levels to the tree, so don't assume a length up front.
        for level in 0.. {
            if level >= self.levels.len() {
                break;
            }

            let len = self.levels[level].len();
            if len > 1 && len % hashes_per_block(self.algorithm) != 0 {
                self.commit_tail_block(level);
            }
        }

        MerkleTree::from_levels(self.levels)
    }
}

impl From<MerkleTreeBuilder> for MerkleTree {
    fn from(builder: MerkleTreeBuilder) -> Self {
        builder.finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::cmp::min;
    use test_case::test_case;

    #[allow(clippy::unused_unit)]
    #[test_case(vec![], "15ec7bf0b50732b49f8228e07d24365338f9e3ab994b00af08e5a3bffe55fd8b" ; "test_empty")]
    #[test_case(vec![0xFF; 8192], "68d131bc271f9c192d4f6dcd8fe61bef90004856da19d0f2f514a7f4098b0737"; "test_oneblock")]
    #[test_case(vec![0xFF; 65536], "f75f59a944d2433bc6830ec243bfefa457704d2aed12f30539cd4f18bf1d62cf"; "test_small")]
    #[test_case(vec![0xFF; 2105344], "7d75dfb18bfd48e03b5be4e8e9aeea2f89880cb81c1551df855e0d0a0cc59a67"; "test_large")]
    #[test_case(vec![0xFF; 2109440], "7577266aa98ce587922fdc668c186e27f3c742fb1b732737153b70ae46973e43"; "test_unaligned")]
    fn tests(input: Vec<u8>, output: &str) {
        let mut builder = MerkleTreeBuilder::new();
        builder.write(input.as_slice());

        let tree = builder.finish();
        let actual = tree.root();
        let expected: Hash = hex::decode(output).unwrap();
        assert_eq!(&expected, actual);
    }

    #[test]
    fn test_unaligned_single_block() {
        let data = vec![0xFF; 8192];
        let mut builder = MerkleTreeBuilder::new();
        let (first, second) = &data[..].split_at(1024);
        builder.write(first);
        builder.write(second);

        let tree = builder.finish();
        let root = tree.root();

        let expected =
            hex::decode("68d131bc271f9c192d4f6dcd8fe61bef90004856da19d0f2f514a7f4098b0737")
                .unwrap();
        assert_eq!(root, &expected);
    }

    #[test]
    fn test_unaligned_n_block() {
        let data = vec![0xFF; 65536];
        let expected =
            hex::decode("f75f59a944d2433bc6830ec243bfefa457704d2aed12f30539cd4f18bf1d62cf")
                .unwrap();

        for chunk_size in &[1, 100, 1024, 8193] {
            let mut builder = MerkleTreeBuilder::new();
            for block in data.as_slice().chunks(*chunk_size) {
                builder.write(block);
            }
            let tree = builder.finish();
            let root = tree.root();

            assert_eq!(root, &expected);
        }
    }

    #[test]
    fn test_fuchsia() {
        let fuchsia: Vec<_> = vec![0xff, 0x00, 0x80]
            .into_iter()
            .cycle()
            .take(3 * BLOCK_SIZE)
            .collect();

        let mut t = MerkleTreeBuilder::new();

        let mut remaining = 0xff0080;
        while remaining > 0 {
            let n = min(remaining, fuchsia.len());
            t.write(&fuchsia[..n]);
            remaining -= n;
        }

        let tree = t.finish();
        let root = tree.root();

        let expected: Hash =
            hex::decode("2feb488cffc976061998ac90ce7292241dfa86883c0edc279433b5c4370d0f30")
                .unwrap();
        assert_eq!(&expected, root);
    }
}