fsverity_merkle/
builder.rs

Help
1// Copyright 2023 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::{FsVerityHasher, MerkleTree};
6use std::cmp::min;
7
8/// A `MerkleTreeBuilder` generates a [`MerkleTree`] from one or more write calls.
9///
10/// # Examples
11/// ```
12/// # use fsverity_merkle::*;
13/// let data = vec![0xff; 8192];
14/// let hasher = FsVerityHasherSha256(FsVerityHasherOptions::new(vec![0xFF; 8], 4096));
15/// let mut builder = MerkleTreeBuilder::new(hasher);
16/// for i in 0..8 {
17///     builder.write(&data[..]);
18/// }
19/// assert_eq!(
20///     builder.finish().root().bytes(),
21///     FromHex::from_hex("ec6b4dc183833a5665b8d804c6e900f2543b54914f153e4139cb77b261f59615")
22///         .unwrap()
23/// );
24/// ```
25#[derive(Clone, Debug)]
26pub struct MerkleTreeBuilder {
27    /// Buffer to hold a partial block of data between [`MerkleTreeBuilder::write`] calls.
28    /// `block.len()` will never exceed `hasher.block_size()`.
29    block: Vec<u8>,
30    levels: Vec<Vec<Vec<u8>>>,
31    hasher: FsVerityHasher,
32}
33
34impl MerkleTreeBuilder {
35    /// Creates a new, empty `MerkleTreeBuilder`.
36    pub fn new(hasher: FsVerityHasher) -> Self {
37        MerkleTreeBuilder {
38            block: Vec::with_capacity(hasher.block_size().into()),
39            levels: vec![Vec::new()],
40            hasher,
41        }
42    }
43
44    /// Append a buffer of bytes to the merkle tree.
45    ///
46    /// No internal buffering is required if all writes are `self.hasher.block_size()` aligned.
47    pub fn write(&mut self, buf: &[u8]) {
48        let block_size = self.hasher.block_size();
49        // Fill the current partial block, if it exists.
50        let buf = if self.block.is_empty() {
51            buf
52        } else {
53            let left = block_size - self.block.len();
54            let prefix = min(buf.len(), left.into());
55            let (buf, rest) = buf.split_at(prefix);
56            self.block.extend_from_slice(buf);
57            if self.block.len() == block_size {
58                self.push_data_hash(self.hasher.hash_block(&self.block[..]));
59            }
60            rest
61        };
62
63        // Write full blocks, saving any final partial block for later writes.
64        for block in buf.chunks(block_size) {
65            if block.len() == block_size {
66                self.push_data_hash(self.hasher.hash_block(block));
67            } else {
68                self.block.extend_from_slice(block);
69            }
70        }
71    }
72
73    /// Save a data block hash, propagating full blocks of hashes to higher layers. Also clear a
74    /// stored data block.
75    pub fn push_data_hash(&mut self, hash: Vec<u8>) {
76        let hashes_per_block = self.hasher.block_size() / self.hasher.hash_size();
77        self.block.clear();
78        self.levels[0].push(hash);
79        if self.levels[0].len() % hashes_per_block == 0 {
80            self.commit_tail_block(0);
81        }
82    }
83
84    /// Hash a complete (or final partial) block of hashes, chaining to higher levels as needed.
85    fn commit_tail_block(&mut self, level: usize) {
86        let hashes_per_block = self.hasher.block_size() / self.hasher.hash_size();
87
88        let len = self.levels[level].len();
89        let next_level = level + 1;
90
91        if next_level >= self.levels.len() {
92            self.levels.push(Vec::new());
93        }
94
95        let first_hash = if len % hashes_per_block == 0 {
96            len - hashes_per_block
97        } else {
98            if !self.hasher.fsverity() {
99                for _ in 0..(hashes_per_block - (len % hashes_per_block)) {
100                    self.levels[level].push(self.hasher.hash_block(&[]))
101                }
102            }
103            len - (len % hashes_per_block)
104        };
105
106        let hash = self.hasher.hash_hashes(&self.levels[level][first_hash..]);
107
108        self.levels[next_level].push(hash);
109        if self.levels[next_level].len() % hashes_per_block == 0 {
110            self.commit_tail_block(next_level);
111        }
112    }
113
114    /// Finalize all levels of the merkle tree, converting this `MerkleTreeBuilder` instance to a
115    /// [`MerkleTree`].
116    pub fn finish(mut self) -> MerkleTree {
117        let hashes_per_block = self.hasher.block_size() / self.hasher.hash_size();
118
119        // The data protected by the tree may not be `hasher.block_size()` aligned. Commit a partial
120        // data block before finalizing the hash levels.
121        // Also, an empty tree consists of a single, empty block. Handle that case now as well.
122        if !self.block.is_empty() || self.levels[0].is_empty() {
123            self.push_data_hash(self.hasher.hash_block(&self.block[..]));
124        }
125
126        // Enumerate the hash levels, finalizing any that have a partial block of hashes.
127        // `commit_tail_block` may add new levels to the tree, so don't assume a length up front.
128        for level in 0.. {
129            if level >= self.levels.len() {
130                break;
131            }
132
133            let len = self.levels[level].len();
134            if len > 1 && len % hashes_per_block != 0 {
135                self.commit_tail_block(level);
136            }
137        }
138
139        MerkleTree::from_levels(self.levels)
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use crate::FsVerityHasherOptions;
147    use hex::FromHex;
148    use test_case::test_case;
149
150    /// Output produced via:
151    /// fsverity digest foo.txt --out-descriptor=/tmp/descriptor
152    /// hexdump /tmp/descriptor -e "16/1 \"%02x\" \"\n\"" -v
153    #[allow(clippy::unused_unit)]
154    #[test_case(vec![], "0000000000000000000000000000000000000000000000000000000000000000" ; "test_empty")]
155    #[test_case(vec![0xFF; 8192], "e95eba0e6902ce10c80029e06051080479c696c21b63c3fffa4d7a01aa15e8cb" ; "test_oneblock")]
156    #[test_case(vec![0xFF; 65536], "b4d0d21943e744d999df91cf5efe432744f41d1763a04a75ef6559e04a143857"; "test_small")]
157    #[test_case(vec![0xFF; 2105344], "b4050a226383d94c09c004d59a81b08bed17726b79cf9bd0994931f13213652d"; "test_large")]
158    #[test_case(vec![0xFF; 2109440], "1a07efa041afdf78b86df2c580ec6f8446eb6e802321252996563c14334a5342"; "test_unaligned")]
159    fn sha256_tests_no_salt(input: Vec<u8>, output: &str) {
160        let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha256(
161            FsVerityHasherOptions::new(vec![], 4096),
162        ));
163        builder.write(input.as_slice());
164        let tree = builder.finish();
165        let expected: [u8; 32] = FromHex::from_hex(output).unwrap();
166        assert_eq!(expected, tree.root());
167    }
168
169    /// Output produced via:
170    /// fsverity digest foo.txt --out-descriptor=/tmp/descriptor --salt="ffffffffffffffff"
171    /// hexdump /tmp/descriptor -e "16/1 \"%02x\" \"\n\"" -v
172    #[test_case(vec![], "0000000000000000000000000000000000000000000000000000000000000000" ; "test_empty")]
173    #[test_case(vec![0xFF; 8192], "e9c09b505561b9509f93b5c7990ed41427f708480c56306453d505e94076d600" ; "test_oneblock")]
174    #[test_case(vec![0xFF; 65536], "ec6b4dc183833a5665b8d804c6e900f2543b54914f153e4139cb77b261f59615"; "test_small")]
175    #[test_case(vec![0xFF; 2105344], "b433c8b632c79ca9fc2c04913541aa38970ae9da04a43269f67770221e79fe37"; "test_large")]
176    #[test_case(vec![0xFF; 2109440], "fbd261c306f522aba5ac0c70229870594d236634f5afe68fe9656ea04eb4a4fe"; "test_unaligned")]
177    fn sha256_tests_with_salt(input: Vec<u8>, output: &str) {
178        let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha256(
179            FsVerityHasherOptions::new(vec![0xFF; 8], 4096),
180        ));
181        builder.write(input.as_slice());
182        let tree = builder.finish();
183        let expected: [u8; 32] = FromHex::from_hex(output).unwrap();
184        assert_eq!(expected, tree.root());
185    }
186
187    /// Output produced via:
188    /// fsverity digest foo.txt --out-descriptor=/tmp/descriptor --hash-alg=sha512
189    /// hexdump /tmp/descriptor -e "16/1 \"%02x\" \"\n\"" -v
190    #[test_case(vec![], "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000" ; "test_empty")]
191    #[test_case(vec![0xFF; 8192], "9606a93b78555a72b49e583fca582bc5b0411decd879286378d86b5f42190f33b560071063622b3d8e0abc9118e3838dbe77301870743ffc80bd0c910ab3522e" ; "test_oneblock")]
192    #[test_case(vec![0xFF; 65536], "524238523bf3c88f78fba612223d322b3c4290a9ecbd9c61aeb8f1c293b0740083a61a0b2e344b8dc6020ece806ea1d048885db1e77ee9fbaa63c3589f6e403b"; "test_small")]
193    #[test_case(vec![0xFF; 2105344], "51977ac06edd17d32761e27d384f6c437ead6922f0a3fbabc3390d8f6e929bc1d9ff9e4ee34fb060484e8eff272f9cc36fa1cf26361c3258b5d8b87d8144b497"; "test_large")]
194    #[test_case(vec![0xFF; 2109440], "f6e821f7cdd1306031080ff99c4c2d7270c6d6bbaa07f4e3040a5d20a1178af1e4f6377f898166d5835ec22b2fcca6d364711cf0c20862d40f3580b6b6276683"; "test_unaligned")]
195    fn sha512_tests_no_salt(input: Vec<u8>, output: &str) {
196        let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha512(
197            FsVerityHasherOptions::new(vec![], 4096),
198        ));
199        builder.write(input.as_slice());
200        let tree = builder.finish();
201        let expected: [u8; 64] = FromHex::from_hex(output).unwrap();
202        assert_eq!(expected, tree.root());
203    }
204
205    /// Output produced via:
206    /// fsverity digest foo.txt --out-descriptor=/tmp/descriptor --salt="ffffffffffffffff" --hash-alg=sha512
207    /// hexdump /tmp/descriptor -e "16/1 \"%02x\" \"\n\"" -v
208    #[test_case(vec![], "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000" ; "test_empty")]
209    #[test_case(vec![0xFF; 8192], "22750472f522bf68a1fe2a66ee1ac57759b322c634d931097b3751e3cd9fe9dd2d8f551631922bf8f675e4b5e3a38e6db11c7df0e5053e80ffbac2c2d7a0105b" ; "test_oneblock")]
210    #[test_case(vec![0xFF; 65536], "4f6a2e16dabf6347b9ae88d5c298befcff0cc71abe1905fa6aefcee14fa5acb89ecbf949daef002d11a9dbb51f211f0eb3e2f7f5e2911b0af2e9fb68c7799a94"; "test_small")]
211    #[test_case(vec![0xFF; 2105344], "a92ddf722dfcf679a64b6364de7f823850f8f856e0ba2c53d66f75cf72d5572bf1d525b3c185e5c39818e2d29997d259f81363daab80a902f86291a71514f891"; "test_large")]
212    #[test_case(vec![0xFF; 2109440], "b6913e8c1d3bb84b467e24667aedad0491ad86f548e849741969688b2526919a380946bebf481ec1ee1bdda86631e10c4a82e7329afdd84db2ac43994a524785"; "test_unaligned")]
213    fn sha512_tests_with_salt(input: Vec<u8>, output: &str) {
214        let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha512(
215            FsVerityHasherOptions::new(vec![0xFF; 8], 4096),
216        ));
217        builder.write(input.as_slice());
218        let tree = builder.finish();
219        let expected: [u8; 64] = FromHex::from_hex(output).unwrap();
220        assert_eq!(expected, tree.root());
221    }
222
223    #[test]
224    fn test_unaligned_single_block_sha256() {
225        let data = vec![0xFF; 8192];
226        let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha256(
227            FsVerityHasherOptions::new(vec![0xFF; 8], 4096),
228        ));
229        let (first, second) = &data[..].split_at(1024);
230        builder.write(first);
231        builder.write(second);
232        let tree = builder.finish();
233        let expected: [u8; 32] =
234            FromHex::from_hex("e9c09b505561b9509f93b5c7990ed41427f708480c56306453d505e94076d600")
235                .unwrap();
236        assert_eq!(tree.root(), expected);
237    }
238
239    #[test]
240    fn test_unaligned_single_block_sha512() {
241        let data = vec![0xFF; 8192];
242        let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha512(
243            FsVerityHasherOptions::new(vec![0xFF; 8], 4096),
244        ));
245        let (first, second) = &data[..].split_at(1024);
246        builder.write(first);
247        builder.write(second);
248        let tree = builder.finish();
249        let expected: [u8; 64] = FromHex::from_hex("22750472f522bf68a1fe2a66ee1ac57759b322c634d931097b3751e3cd9fe9dd2d8f551631922bf8f675e4b5e3a38e6db11c7df0e5053e80ffbac2c2d7a0105b").unwrap();
250        assert_eq!(tree.root(), expected);
251    }
252
253    #[test]
254    fn test_unaligned_n_block_sha256() {
255        let data = vec![0xFF; 65536];
256        let expected: [u8; 32] =
257            FromHex::from_hex("ec6b4dc183833a5665b8d804c6e900f2543b54914f153e4139cb77b261f59615")
258                .unwrap();
259
260        for chunk_size in &[1, 100, 1024, 8193] {
261            let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha256(
262                FsVerityHasherOptions::new(vec![0xFF; 8], 4096),
263            ));
264            for block in data.as_slice().chunks(*chunk_size) {
265                builder.write(block);
266            }
267            let tree = builder.finish();
268
269            assert_eq!(tree.root(), expected);
270        }
271    }
272
273    #[test]
274    fn test_unaligned_n_block_sha512() {
275        let data = vec![0xFF; 65536];
276        let expected: [u8; 64] = FromHex::from_hex("4f6a2e16dabf6347b9ae88d5c298befcff0cc71abe1905fa6aefcee14fa5acb89ecbf949daef002d11a9dbb51f211f0eb3e2f7f5e2911b0af2e9fb68c7799a94").unwrap();
277
278        for chunk_size in &[1, 100, 1024, 8193] {
279            let mut builder = MerkleTreeBuilder::new(FsVerityHasher::Sha512(
280                FsVerityHasherOptions::new(vec![0xFF; 8], 4096),
281            ));
282            for block in data.as_slice().chunks(*chunk_size) {
283                builder.write(block);
284            }
285            let tree = builder.finish();
286
287            assert_eq!(tree.root(), expected);
288        }
289    }
290}
fsverity_merkle/builder.rs

fsverity_merkle/
builder.rs
