public/rust/fuchsia-merkle/src/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::hash::Hash;
use crate::tree::MerkleTree;
use crate::util::{hash_block, hash_hashes, HASHES_PER_BLOCK};
use crate::BLOCK_SIZE;

/// A `MerkleTreeBuilder` generates a [`MerkleTree`] from one or more write calls.
///
/// # Examples
/// ```
/// # use fuchsia_merkle::*;
/// let data = vec![0xff; 8192];
/// let mut builder = MerkleTreeBuilder::new();
/// for i in 0..8 {
///     builder.write(&data[..]);
/// }
/// assert_eq!(
///     builder.finish().root(),
///     "f75f59a944d2433bc6830ec243bfefa457704d2aed12f30539cd4f18bf1d62cf"
///         .parse()
///         .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>>,
}

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

    /// 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)
    }

    /// Save a data block hash, propagating full blocks of hashes to higher layers. Also clear a
    /// stored data block.
    fn push_data_hash(&mut self, hash: Hash) {
        self.block.clear();
        self.levels[0].push(hash);
        if self.levels[0].len() % HASHES_PER_BLOCK == 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 first_hash = if len % HASHES_PER_BLOCK == 0 {
            len - HASHES_PER_BLOCK
        } else {
            len - (len % HASHES_PER_BLOCK)
        };

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

        self.levels[next_level].push(hash);
        if self.levels[next_level].len() % HASHES_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 != 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;

    macro_rules! test_case {
        ($name:ident, $input:expr, $output:expr) => {
            #[test]
            fn $name() {
                let input = $input;
                let mut tree = MerkleTreeBuilder::new();
                tree.write(input.as_slice());
                let actual = tree.finish().root();
                let expected: Hash = $output.parse().unwrap();
                assert_eq!(expected, actual);
            }
        };
    }

    test_case!(
        test_empty,
        vec![],
        "15ec7bf0b50732b49f8228e07d24365338f9e3ab994b00af08e5a3bffe55fd8b"
    );

    test_case!(
        test_oneblock,
        vec![0xFF; 8192],
        "68d131bc271f9c192d4f6dcd8fe61bef90004856da19d0f2f514a7f4098b0737"
    );

    test_case!(
        test_small,
        vec![0xFF; 65536],
        "f75f59a944d2433bc6830ec243bfefa457704d2aed12f30539cd4f18bf1d62cf"
    );

    test_case!(
        test_large,
        vec![0xFF; 2105344],
        "7d75dfb18bfd48e03b5be4e8e9aeea2f89880cb81c1551df855e0d0a0cc59a67"
    );

    test_case!(
        test_unaligned,
        vec![0xFF; 2109440],
        "7577266aa98ce587922fdc668c186e27f3c742fb1b732737153b70ae46973e43"
    );

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

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

        let expected = "68d131bc271f9c192d4f6dcd8fe61bef90004856da19d0f2f514a7f4098b0737"
            .parse()
            .unwrap();
        assert_eq!(root, expected);
    }

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

        for chunk_size in vec![1, 100, 1024, 8193] {
            let mut tree = MerkleTreeBuilder::new();
            for block in data.as_slice().chunks(chunk_size) {
                tree.write(block);
            }
            let root = tree.finish().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 actual = t.finish().root();
        let expected: Hash = "2feb488cffc976061998ac90ce7292241dfa86883c0edc279433b5c4370d0f30"
            .parse()
            .unwrap();
        assert_eq!(expected, actual);
    }
}