1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
// 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::{crypto_library_init, hash_block, hash_hashes, HASHES_PER_BLOCK};
use crate::{Hash, 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 Default for MerkleTreeBuilder {
    fn default() -> Self {
        crypto_library_init();
        Self { levels: vec![Vec::new()], block: Vec::with_capacity(BLOCK_SIZE) }
    }
}

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

    /// 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.
    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 == 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;
    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 tree = MerkleTreeBuilder::new();
        tree.write(input.as_slice());
        let actual = tree.finish().root();
        let expected: Hash = output.parse().unwrap();
        assert_eq!(expected, actual);
    }

    #[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 &[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);
    }
}