mega/
file_validator.rs

1use crate::FileKey;
2use cbc::cipher::BlockModeEncrypt;
3use cbc::cipher::KeyIvInit;
4
5type Aes128CbcEnc = cbc::Encryptor<aes::Aes128>;
6
7/// An error that occurs when a file fails validation.
8#[derive(Debug, Clone)]
9pub struct FileValidationError {
10    /// The actual created mac
11    pub actual_mac: [u8; 8],
12
13    /// The expected mac
14    pub expected_mac: [u8; 8],
15}
16
17impl std::fmt::Display for FileValidationError {
18    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
19        write!(
20            f,
21            "file mac mismatch, expected {} but got {}",
22            HexSlice(&self.expected_mac),
23            HexSlice(&self.actual_mac)
24        )
25    }
26}
27
28impl std::error::Error for FileValidationError {}
29
30/// A helper to format byte slices as hex
31struct HexSlice<'a>(&'a [u8]);
32
33impl std::fmt::Display for HexSlice<'_> {
34    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
35        write!(f, "0x")?;
36        for byte in self.0.iter() {
37            write!(f, "{byte:X}")?;
38        }
39        Ok(())
40    }
41}
42
43/// A struct to validate files.
44pub struct FileValidator {
45    file_key: FileKey,
46    chunk_iter: ChunkIter,
47    left_in_chunk: usize,
48    file_mac: u128,
49    chunk_mac: u128,
50    buffer: [u8; 16],
51    buffer_end: usize,
52    done: bool,
53}
54
55impl FileValidator {
56    /// Make a new file validator.
57    pub fn new(file_key: FileKey) -> Self {
58        let mut chunk_iter = ChunkIter::new();
59        // ChunkIter is infinite.
60        let (_, left_in_chunk) = chunk_iter.next().unwrap();
61        // This can only fail when a usize is a u16.
62        let left_in_chunk = usize::try_from(left_in_chunk).unwrap();
63        let chunk_mac = create_chunk_mac(&file_key);
64
65        Self {
66            file_key,
67            chunk_iter,
68            left_in_chunk,
69            file_mac: 0,
70            chunk_mac,
71            buffer: [0; 16],
72            buffer_end: 0,
73            done: false,
74        }
75    }
76
77    /// Process a block
78    fn process_block(&mut self, block: [u8; 16]) {
79        self.chunk_mac ^= u128::from_be_bytes(block);
80        let mut chunk_mac_bytes = self.chunk_mac.to_be_bytes();
81        aes_cbc_encrypt_u128(self.file_key.key, &mut chunk_mac_bytes);
82        self.chunk_mac = u128::from_be_bytes(chunk_mac_bytes);
83
84        self.left_in_chunk -= 16;
85        if self.left_in_chunk == 0 {
86            self.begin_new_chunk();
87        }
88    }
89
90    /// Begin a new chunk.
91    fn begin_new_chunk(&mut self) {
92        self.file_mac ^= self.chunk_mac;
93        let mut file_mac_bytes = self.file_mac.to_be_bytes();
94        aes_cbc_encrypt_u128(self.file_key.key, &mut file_mac_bytes);
95        self.file_mac = u128::from_be_bytes(file_mac_bytes);
96
97        // Reset chunk state.
98        self.chunk_mac = create_chunk_mac(&self.file_key);
99        // ChunkIter is infinite.
100        let (_, left_in_chunk) = self.chunk_iter.next().unwrap();
101        // This can only fail when a usize is a u16.
102        self.left_in_chunk = usize::try_from(left_in_chunk).unwrap();
103    }
104
105    /// Feed data.
106    ///
107    /// This should be fed decrypted bytes.
108    ///
109    /// # Panics
110    /// Panics if finish has already been called.
111    pub fn feed(&mut self, mut input: &[u8]) {
112        assert!(!self.done, "cannot feed a completed file validator");
113
114        if self.buffer_end != 0 {
115            // Try to complete the buffer and process the block.
116            let need_to_consume = self.buffer.len() - self.buffer_end;
117            let consume_input_len = std::cmp::min(need_to_consume, input.len());
118
119            self.buffer[self.buffer_end..(self.buffer_end + consume_input_len)]
120                .copy_from_slice(&input[..consume_input_len]);
121            if consume_input_len < need_to_consume {
122                self.buffer_end += consume_input_len;
123                // The input was too small to create even 1 block.
124                return;
125            }
126
127            self.process_block(self.buffer);
128            input = &input[need_to_consume..];
129            self.buffer_end = 0;
130        }
131
132        // Split the input into blocks and process up until the last whole block.
133        let (block_iter, remainder) = input.as_chunks::<16>();
134        for block in block_iter {
135            self.process_block(*block);
136        }
137
138        // Save remaining input into the buffer.
139        let remainder_len = remainder.len();
140        if remainder_len != 0 {
141            self.buffer[..remainder_len].copy_from_slice(remainder);
142            self.buffer_end = remainder_len;
143        }
144    }
145
146    /// Finish feeding this data and validate the file.
147    ///
148    /// # Panics
149    /// Panics if finish has already been called.
150    pub fn finish(&mut self) -> Result<(), FileValidationError> {
151        assert!(!self.done, "cannot finish a completed file validator");
152        self.done = true;
153
154        if self.buffer_end != 0 {
155            // Pad block with zeros.
156            self.buffer[self.buffer_end..].fill(0);
157            self.process_block(self.buffer);
158        }
159
160        let mut file_mac = self.file_mac ^ self.chunk_mac;
161        let mut file_mac_bytes = file_mac.to_be_bytes();
162        aes_cbc_encrypt_u128(self.file_key.key, &mut file_mac_bytes);
163        file_mac = u128::from_be_bytes(file_mac_bytes);
164
165        let file_mac_bytes = file_mac.to_be_bytes();
166        let file_mac_u32_0 = u32::from_be_bytes(file_mac_bytes[..4].try_into().unwrap());
167        let file_mac_u32_1 = u32::from_be_bytes(file_mac_bytes[4..8].try_into().unwrap());
168        let file_mac_u32_2 = u32::from_be_bytes(file_mac_bytes[8..12].try_into().unwrap());
169        let file_mac_u32_3 = u32::from_be_bytes(file_mac_bytes[12..].try_into().unwrap());
170
171        let final_file_mac_u32_0 = file_mac_u32_0 ^ file_mac_u32_1;
172        let final_file_mac_u32_1 = file_mac_u32_2 ^ file_mac_u32_3;
173
174        let mut final_file_mac_bytes = [0; 8];
175        final_file_mac_bytes[..4].copy_from_slice(&final_file_mac_u32_0.to_be_bytes());
176        final_file_mac_bytes[4..].copy_from_slice(&final_file_mac_u32_1.to_be_bytes());
177        let final_file_mac = u64::from_be_bytes(final_file_mac_bytes);
178
179        if final_file_mac != self.file_key.meta_mac {
180            return Err(FileValidationError {
181                expected_mac: self.file_key.meta_mac.to_be_bytes(),
182                actual_mac: final_file_mac.to_be_bytes(),
183            });
184        }
185
186        Ok(())
187    }
188
189    /// Check if finish has already been called.
190    pub fn is_finished(&self) -> bool {
191        self.done
192    }
193}
194
195fn create_chunk_mac(file_key: &FileKey) -> u128 {
196    let mut chunk_mac_bytes = [0; 16];
197    let iv_bytes = file_key.iv.to_be_bytes();
198    chunk_mac_bytes[..8].copy_from_slice(&iv_bytes[..8]);
199    chunk_mac_bytes[8..].copy_from_slice(&iv_bytes[..8]);
200    u128::from_be_bytes(chunk_mac_bytes)
201}
202
203fn aes_cbc_encrypt_u128(key: u128, data: &mut [u8; 16]) {
204    let mut cipher = Aes128CbcEnc::new(&key.to_be_bytes().into(), &[0; 16].into());
205    cipher.encrypt_block((data).into());
206}
207
208/// An iterator over chunks
209#[derive(Debug)]
210struct ChunkIter {
211    /// The offset into the file
212    offset: u64,
213    delta: u64,
214}
215
216impl ChunkIter {
217    fn new() -> Self {
218        Self {
219            delta: 0,
220            offset: 0,
221        }
222    }
223}
224
225impl Iterator for ChunkIter {
226    type Item = (u64, u64);
227
228    fn next(&mut self) -> Option<Self::Item> {
229        self.delta += 128 * 1024;
230        self.delta = std::cmp::min(self.delta, 1024 * 1024);
231
232        let old_offset = self.offset;
233        self.offset += self.delta;
234
235        Some((old_offset, self.delta))
236    }
237}
238
239#[cfg(test)]
240mod test {
241    use super::*;
242    use crate::test::*;
243
244    #[test]
245    #[expect(clippy::erasing_op, clippy::identity_op)]
246    fn chunk_iter() {
247        let mut iter = ChunkIter::new();
248        assert!(iter.next() == Some((128 * 0 * 2014, 128 * 1 * 1024)));
249        assert!(iter.next() == Some((128 * 1 * 1024, 128 * 2 * 1024)));
250        assert!(iter.next() == Some((128 * 3 * 1024, 128 * 3 * 1024)));
251        assert!(iter.next() == Some((128 * 6 * 1024, 128 * 4 * 1024)));
252        assert!(iter.next() == Some((128 * 10 * 1024, 128 * 5 * 1024)));
253        assert!(iter.next() == Some((128 * 15 * 1024, 128 * 6 * 1024)));
254        assert!(iter.next() == Some((128 * 21 * 1024, 128 * 7 * 1024)));
255        assert!(iter.next() == Some((128 * 28 * 1024, 128 * 8 * 1024)));
256        assert!(iter.next() == Some((128 * 36 * 1024, 128 * 8 * 1024)));
257        assert!(iter.next() == Some((128 * 44 * 1024, 128 * 8 * 1024)));
258    }
259
260    #[test]
261    fn file_validator() {
262        let file_key = FileKey {
263            key: TEST_FILE_KEY_KEY_DECODED,
264            iv: TEST_FILE_KEY_IV_DECODED,
265            meta_mac: TEST_FILE_META_MAC_DECODED,
266        };
267
268        let mut validator = FileValidator::new(file_key);
269        validator.feed(TEST_FILE_BYTES);
270        validator.finish().expect("invalid mac");
271    }
272}
mega/file_validator.rs

mega/
file_validator.rs
