mega/
file_validator.rs

1use crate::FileKey;
2use cbc::cipher::BlockEncryptMut;
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, "{:X}", byte)?;
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 mut block_iter = input.chunks_exact(16);
134        for block in block_iter.by_ref() {
135            // The iter will always produce blocks of the right size.
136            let block = block.try_into().unwrap();
137            self.process_block(block);
138        }
139
140        // Save remaining input into the buffer.
141        let remainder = block_iter.remainder();
142        let remainder_len = remainder.len();
143        if remainder_len != 0 {
144            self.buffer[..remainder_len].copy_from_slice(remainder);
145            self.buffer_end = remainder_len;
146        }
147    }
148
149    /// Finish feeding this data and validate the file.
150    ///
151    /// # Panics
152    /// Panics if finish has already been called.
153    pub fn finish(&mut self) -> Result<(), FileValidationError> {
154        assert!(!self.done, "cannot finish a completed file validator");
155        self.done = true;
156
157        if self.buffer_end != 0 {
158            // Pad block with zeros.
159            self.buffer[self.buffer_end..].fill(0);
160            self.process_block(self.buffer);
161        }
162
163        let mut file_mac = self.file_mac ^ self.chunk_mac;
164        let mut file_mac_bytes = file_mac.to_be_bytes();
165        aes_cbc_encrypt_u128(self.file_key.key, &mut file_mac_bytes);
166        file_mac = u128::from_be_bytes(file_mac_bytes);
167
168        let file_mac_bytes = file_mac.to_be_bytes();
169        let file_mac_u32_0 = u32::from_be_bytes(file_mac_bytes[..4].try_into().unwrap());
170        let file_mac_u32_1 = u32::from_be_bytes(file_mac_bytes[4..8].try_into().unwrap());
171        let file_mac_u32_2 = u32::from_be_bytes(file_mac_bytes[8..12].try_into().unwrap());
172        let file_mac_u32_3 = u32::from_be_bytes(file_mac_bytes[12..].try_into().unwrap());
173
174        let final_file_mac_u32_0 = file_mac_u32_0 ^ file_mac_u32_1;
175        let final_file_mac_u32_1 = file_mac_u32_2 ^ file_mac_u32_3;
176
177        let mut final_file_mac_bytes = [0; 8];
178        final_file_mac_bytes[..4].copy_from_slice(&final_file_mac_u32_0.to_be_bytes());
179        final_file_mac_bytes[4..].copy_from_slice(&final_file_mac_u32_1.to_be_bytes());
180        let final_file_mac = u64::from_be_bytes(final_file_mac_bytes);
181
182        if final_file_mac != self.file_key.meta_mac {
183            return Err(FileValidationError {
184                expected_mac: self.file_key.meta_mac.to_be_bytes(),
185                actual_mac: final_file_mac.to_be_bytes(),
186            });
187        }
188
189        Ok(())
190    }
191
192    /// Check if finish has already been called.
193    pub fn is_finished(&self) -> bool {
194        self.done
195    }
196}
197
198fn create_chunk_mac(file_key: &FileKey) -> u128 {
199    let mut chunk_mac_bytes = [0; 16];
200    let iv_bytes = file_key.iv.to_be_bytes();
201    chunk_mac_bytes[..8].copy_from_slice(&iv_bytes[..8]);
202    chunk_mac_bytes[8..].copy_from_slice(&iv_bytes[..8]);
203    u128::from_be_bytes(chunk_mac_bytes)
204}
205
206fn aes_cbc_encrypt_u128(key: u128, data: &mut [u8; 16]) {
207    let mut cipher = Aes128CbcEnc::new(&key.to_be_bytes().into(), &[0; 16].into());
208    cipher.encrypt_block_mut((data).into());
209}
210
211/// An iterator over chunks
212#[derive(Debug)]
213struct ChunkIter {
214    /// The offset into the file
215    offset: u64,
216    delta: u64,
217}
218
219impl ChunkIter {
220    fn new() -> Self {
221        Self {
222            delta: 0,
223            offset: 0,
224        }
225    }
226}
227
228impl Iterator for ChunkIter {
229    type Item = (u64, u64);
230
231    fn next(&mut self) -> Option<Self::Item> {
232        self.delta += 128 * 1024;
233        self.delta = std::cmp::min(self.delta, 1024 * 1024);
234
235        let old_offset = self.offset;
236        self.offset += self.delta;
237
238        Some((old_offset, self.delta))
239    }
240}
241
242#[cfg(test)]
243mod test {
244    use super::*;
245    use crate::test::*;
246
247    #[test]
248    #[expect(clippy::erasing_op, clippy::identity_op)]
249    fn chunk_iter() {
250        let mut iter = ChunkIter::new();
251        assert!(iter.next() == Some((128 * 0 * 2014, 128 * 1 * 1024)));
252        assert!(iter.next() == Some((128 * 1 * 1024, 128 * 2 * 1024)));
253        assert!(iter.next() == Some((128 * 3 * 1024, 128 * 3 * 1024)));
254        assert!(iter.next() == Some((128 * 6 * 1024, 128 * 4 * 1024)));
255        assert!(iter.next() == Some((128 * 10 * 1024, 128 * 5 * 1024)));
256        assert!(iter.next() == Some((128 * 15 * 1024, 128 * 6 * 1024)));
257        assert!(iter.next() == Some((128 * 21 * 1024, 128 * 7 * 1024)));
258        assert!(iter.next() == Some((128 * 28 * 1024, 128 * 8 * 1024)));
259        assert!(iter.next() == Some((128 * 36 * 1024, 128 * 8 * 1024)));
260        assert!(iter.next() == Some((128 * 44 * 1024, 128 * 8 * 1024)));
261    }
262
263    #[test]
264    fn file_validator() {
265        let file_key = FileKey {
266            key: TEST_FILE_KEY_KEY_DECODED,
267            iv: TEST_FILE_KEY_IV_DECODED,
268            meta_mac: TEST_FILE_META_MAC_DECODED,
269        };
270
271        let mut validator = FileValidator::new(file_key);
272        validator.feed(TEST_FILE_BYTES);
273        validator.finish().expect("invalid mac");
274    }
275}
mega/file_validator.rs

mega/
file_validator.rs
