Skip to main content

fsverity_merkle/
builder.rs

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