fsverity_merkle/
builder.rs1use crate::FsVerityHash;
6use crate::tree::MerkleTree;
7use crate::util::FsVerityHasher;
8use std::cmp::min;
9use zerocopy::IntoBytes;
10
11#[derive(Clone, Debug)]
13pub struct MerkleTreeBuilder<D> {
14 block: Vec<u8>,
17 levels: Vec<Vec<D>>,
18 hasher: FsVerityHasher,
19}
20
21impl<D: FsVerityHash> MerkleTreeBuilder<D> {
22 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 pub fn write(&mut self, buf: &[u8]) {
40 let block_size = self.hasher.block_size();
41 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 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 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 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 pub fn finish(mut self) -> MerkleTree {
111 let hashes_per_block = self.hasher.block_size() / self.hasher.hash_size();
112
113 if !self.block.is_empty() || self.levels[0].is_empty() {
117 self.push_data_hash(self.hash_block(&self.block[..]));
118 }
119
120 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 #[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 #[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 #[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 #[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}