blob: 94b9718f3273d9a51ceca3d937edadb97f5e34b6 [file] [log] [blame]
Victor Hsiehdde17902021-02-26 12:35:31 -08001/*
2 * Copyright (C) 2021 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use super::common::{build_fsverity_digest, merkle_tree_height, FsverityError};
18use crate::common::CHUNK_SIZE;
19use crate::crypto::{CryptoError, Sha256Hash, Sha256Hasher};
20
21const HASH_SIZE: usize = Sha256Hasher::HASH_SIZE;
22const HASH_PER_PAGE: usize = CHUNK_SIZE as usize / HASH_SIZE;
23
24/// MerkleLeaves can be used by the class' customer for bookkeeping integrity data for their bytes.
25/// It can also be used to generate the standard fs-verity digest for the source data.
26///
27/// It's in-memory because for the initial use cases, we don't need to read back an existing file,
28/// and only need to deal with new files. Also, considering that the output file won't be large at
29/// the moment, it is sufficient to simply keep the Merkle tree in memory in the trusted world. To
30/// further simplify the initial implementation, we only need to keep the leaf nodes in memory, and
31/// generate the tree / root hash when requested.
32pub struct MerkleLeaves {
33 leaves: Vec<Sha256Hash>,
34 file_size: u64,
35}
36
37fn hash_all_pages(source: &[Sha256Hash]) -> Result<Vec<Sha256Hash>, CryptoError> {
38 source
39 .chunks(HASH_PER_PAGE)
40 .map(|chunk| {
41 let padding_bytes = (HASH_PER_PAGE - chunk.len()) * HASH_SIZE;
42 Ok(Sha256Hasher::new()?
43 .update_from(chunk)?
44 .update(&vec![0u8; padding_bytes])?
45 .finalize()?)
46 })
47 .collect()
48}
49
Victor Hsiehdde17902021-02-26 12:35:31 -080050impl MerkleLeaves {
51 /// Creates a `MerkleLeaves` instance with empty data.
52 pub fn new() -> Self {
53 Self { leaves: Vec::new(), file_size: 0 }
54 }
55
Victor Hsieh6a47e7f2021-03-03 15:53:49 -080056 /// Gets size of the file represented by `MerkleLeaves`.
57 pub fn file_size(&self) -> u64 {
58 self.file_size
59 }
60
Victor Hsiehdde17902021-02-26 12:35:31 -080061 /// Updates the hash of the `index`-th leaf, and increase the size to `size_at_least` if the
62 /// current size is smaller.
63 pub fn update_hash(&mut self, index: usize, hash: &Sha256Hash, size_at_least: u64) {
64 // +1 since index is zero-based.
65 if self.leaves.len() < index + 1 {
66 // When resizing, fill in hash of zeros by default. This makes it easy to handle holes
67 // in a file.
68 self.leaves.resize(index + 1, Sha256Hasher::HASH_OF_4096_ZEROS);
69 }
70 self.leaves[index].clone_from_slice(hash);
71
72 if size_at_least > self.file_size {
73 self.file_size = size_at_least;
74 }
75 }
76
77 /// Returns whether `index` is within the bound of leaves.
78 pub fn is_index_valid(&self, index: usize) -> bool {
79 index < self.leaves.len()
80 }
81
82 /// Returns whether the `index`-th hash is consistent to `hash`.
83 pub fn is_consistent(&self, index: usize, hash: &Sha256Hash) -> bool {
84 if let Some(element) = self.leaves.get(index) {
85 element == hash
86 } else {
87 false
88 }
89 }
90
91 fn calculate_root_hash(&self) -> Result<Sha256Hash, FsverityError> {
92 match self.leaves.len() {
93 // Special cases per fs-verity digest definition.
94 0 => {
95 debug_assert_eq!(self.file_size, 0);
96 Ok([0u8; HASH_SIZE])
97 }
98 1 => {
99 debug_assert!(self.file_size <= CHUNK_SIZE && self.file_size > 0);
100 Ok(self.leaves[0])
101 }
102 n => {
103 debug_assert_eq!((self.file_size - 1) / CHUNK_SIZE, n as u64);
104 let size_for_equivalent = n as u64 * CHUNK_SIZE;
105 let level = merkle_tree_height(size_for_equivalent).unwrap(); // safe since n > 0
106
107 // `leaves` is owned and can't be the initial state below. Here we manually hash it
108 // first to avoid a copy and to get the type right.
109 let second_level = hash_all_pages(&self.leaves)?;
110 let hashes =
111 (1..=level).try_fold(second_level, |source, _| hash_all_pages(&source))?;
112 if hashes.len() != 1 {
113 Err(FsverityError::InvalidState)
114 } else {
115 Ok(hashes.into_iter().next().unwrap())
116 }
117 }
118 }
119 }
120
121 /// Returns the fs-verity digest based on the current tree and file size.
122 pub fn calculate_fsverity_digest(&self) -> Result<Sha256Hash, FsverityError> {
123 let root_hash = self.calculate_root_hash()?;
124 Ok(build_fsverity_digest(&root_hash, self.file_size)?)
125 }
126}
127
128#[cfg(test)]
129mod tests {
130 // Test data below can be generated by:
131 // $ perl -e 'print "\x{00}" x 6000' > foo
132 // $ perl -e 'print "\x{01}" x 5000' >> foo
133 // $ fsverity digest foo
134 use super::*;
135 use anyhow::Result;
136
137 #[test]
138 fn merkle_tree_empty_file() -> Result<()> {
139 assert_eq!(
140 to_u8_vec("3d248ca542a24fc62d1c43b916eae5016878e2533c88238480b26128a1f1af95"),
141 generate_fsverity_digest_sequentially(&Vec::new())?
142 );
143 Ok(())
144 }
145
146 #[test]
147 fn merkle_tree_file_size_less_than_or_equal_to_4k() -> Result<()> {
148 // Test a file that contains 4096 '\01's.
149 assert_eq!(
150 to_u8_vec("cd0875ca59c7d37e962c5e8f5acd3770750ac80225e2df652ce5672fd34500af"),
151 generate_fsverity_digest_sequentially(&vec![1; 4096])?
152 );
153 Ok(())
154 }
155
156 #[test]
157 fn merkle_tree_more_sizes() -> Result<()> {
158 // Test files that contains >4096 '\01's.
159
160 assert_eq!(
161 to_u8_vec("2901b849fda2d91e3929524561c4a47e77bb64734319759507b2029f18b9cc52"),
162 generate_fsverity_digest_sequentially(&vec![1; 4097])?
163 );
164
165 assert_eq!(
166 to_u8_vec("2a476d58eb80394052a3a783111e1458ac3ecf68a7878183fed86ca0ff47ec0d"),
167 generate_fsverity_digest_sequentially(&vec![1; 8192])?
168 );
169
170 // Test with max size that still fits in 2 levels.
171 assert_eq!(
172 to_u8_vec("26b7c190a34e19f420808ee7ec233b09fa6c34543b5a9d2950530114c205d14f"),
173 generate_fsverity_digest_sequentially(&vec![1; 524288])?
174 );
175
176 // Test with data that requires 3 levels.
177 assert_eq!(
178 to_u8_vec("316835d9be1c95b5cd55d07ae7965d651689efad186e26cbf680e40b683a3262"),
179 generate_fsverity_digest_sequentially(&vec![1; 524289])?
180 );
181 Ok(())
182 }
183
184 #[test]
185 fn merkle_tree_non_sequential() -> Result<()> {
186 let mut tree = MerkleLeaves::new();
187 let hash = Sha256Hasher::new()?.update(&vec![1u8; CHUNK_SIZE as usize])?.finalize()?;
188
189 // Update hashes of 4 1-blocks.
190 tree.update_hash(1, &hash, CHUNK_SIZE * 2);
191 tree.update_hash(3, &hash, CHUNK_SIZE * 4);
192 tree.update_hash(0, &hash, CHUNK_SIZE);
193 tree.update_hash(2, &hash, CHUNK_SIZE * 3);
194
195 assert_eq!(
196 to_u8_vec("7d3c0d2e1dc54230b20ed875f5f3a4bd3f9873df601936b3ca8127d4db3548f3"),
197 tree.calculate_fsverity_digest()?
198 );
199 Ok(())
200 }
201
202 fn generate_fsverity_digest_sequentially(test_data: &[u8]) -> Result<Sha256Hash> {
203 let mut tree = MerkleLeaves::new();
204 for (index, chunk) in test_data.chunks(CHUNK_SIZE as usize).enumerate() {
205 let hash = Sha256Hasher::new()?
206 .update(&chunk)?
207 .update(&vec![0u8; CHUNK_SIZE as usize - chunk.len()])?
208 .finalize()?;
209
210 tree.update_hash(index, &hash, CHUNK_SIZE * index as u64 + chunk.len() as u64);
211 }
212 Ok(tree.calculate_fsverity_digest()?)
213 }
214
215 fn to_u8_vec(hex_str: &str) -> Vec<u8> {
216 assert!(hex_str.len() % 2 == 0);
217 (0..hex_str.len())
218 .step_by(2)
219 .map(|i| u8::from_str_radix(&hex_str[i..i + 2], 16).unwrap())
220 .collect()
221 }
222}