libsignal_protocol/
incremental_mac.rs

1//
2// Copyright 2023 Signal Messenger, LLC.
3// SPDX-License-Identifier: AGPL-3.0-only
4//
5
6use hmac::Mac;
7use sha2::digest::typenum::Unsigned;
8use sha2::digest::{FixedOutput, MacError, Output};
9
10#[derive(Clone)]
11pub struct Incremental<M: Mac + Clone> {
12    mac: M,
13    chunk_size: usize,
14    unused_length: usize,
15}
16
17#[derive(Clone)]
18pub struct Validating<M: Mac + Clone> {
19    incremental: Incremental<M>,
20    // Expected MACs in reversed order, to efficiently pop them from off the end
21    expected: Vec<Output<M>>,
22}
23
24const MINIMUM_CHUNK_SIZE: usize = 64 * 1024;
25const MAXIMUM_CHUNK_SIZE: usize = 2 * 1024 * 1024;
26const TARGET_TOTAL_DIGEST_SIZE: usize = 8 * 1024;
27
28pub const fn calculate_chunk_size<D>(data_size: usize) -> usize
29where
30    D: FixedOutput,
31{
32    assert!(
33        0 == TARGET_TOTAL_DIGEST_SIZE % D::OutputSize::USIZE,
34        "Target digest size should be a multiple of digest size"
35    );
36    let target_chunk_count = TARGET_TOTAL_DIGEST_SIZE / D::OutputSize::USIZE;
37    if data_size < target_chunk_count * MINIMUM_CHUNK_SIZE {
38        return MINIMUM_CHUNK_SIZE;
39    }
40    if data_size < target_chunk_count * MAXIMUM_CHUNK_SIZE {
41        return data_size.div_ceil(target_chunk_count);
42    }
43    MAXIMUM_CHUNK_SIZE
44}
45
46impl<M: Mac + Clone> Incremental<M> {
47    pub fn new(mac: M, chunk_size: usize) -> Self {
48        assert!(chunk_size > 0, "chunk size must be positive");
49        Self {
50            mac,
51            chunk_size,
52            unused_length: chunk_size,
53        }
54    }
55
56    pub fn validating<'a, A, I>(self, macs: I) -> Validating<M>
57    where
58        // This is a clunky way to spell an iterator over `&'a [u8; M::OutputSize]`, which requires
59        // feature(generic_const_exprs).
60        A: Into<&'a sha2::digest::array::Array<u8, M::OutputSize>>,
61        I: IntoIterator<Item = A, IntoIter: DoubleEndedIterator>,
62    {
63        let expected = macs
64            .into_iter()
65            .map(|mac| mac.into().to_owned())
66            .rev()
67            .collect();
68        Validating {
69            incremental: self,
70            expected,
71        }
72    }
73
74    pub fn update<'a>(&'a mut self, bytes: &'a [u8]) -> impl Iterator<Item = Output<M>> + 'a {
75        let split_point = std::cmp::min(bytes.len(), self.unused_length);
76        let (to_write, overflow) = bytes.split_at(split_point);
77
78        std::iter::once(to_write)
79            .chain(overflow.chunks(self.chunk_size))
80            .flat_map(move |chunk| self.update_chunk(chunk))
81    }
82
83    pub fn finalize(self) -> Output<M> {
84        self.mac.finalize().into_bytes()
85    }
86
87    fn update_chunk(&mut self, bytes: &[u8]) -> Option<Output<M>> {
88        assert!(bytes.len() <= self.unused_length);
89        self.mac.update(bytes);
90        self.unused_length -= bytes.len();
91        if self.unused_length == 0 {
92            self.unused_length = self.chunk_size;
93            let mac = self.mac.clone();
94            Some(mac.finalize().into_bytes())
95        } else {
96            None
97        }
98    }
99
100    fn pending_bytes_size(&self) -> usize {
101        self.chunk_size - self.unused_length
102    }
103}
104
105impl<M: Mac + Clone> Validating<M> {
106    pub fn update(&mut self, bytes: &[u8]) -> Result<usize, MacError> {
107        let mut result = Ok(0);
108        let macs = self.incremental.update(bytes);
109
110        let mut whole_chunks = 0;
111        for mac in macs {
112            match self.expected.last() {
113                Some(expected) if expected == &mac => {
114                    whole_chunks += 1;
115                    self.expected.pop();
116                }
117                _ => {
118                    result = Err(MacError);
119                }
120            }
121        }
122        let validated_bytes = whole_chunks * self.incremental.chunk_size;
123        result.map(|_| validated_bytes)
124    }
125
126    pub fn finalize(self) -> Result<usize, MacError> {
127        let pending_bytes_size = self.incremental.pending_bytes_size();
128        let mac = self.incremental.finalize();
129        match &self.expected[..] {
130            [expected] if expected == &mac => Ok(pending_bytes_size),
131            _ => Err(MacError),
132        }
133    }
134}
135
136#[cfg(test)]
137mod test {
138    use const_str::hex;
139    use hmac::{Hmac, KeyInit as _};
140    use proptest::prelude::*;
141    use rand::distr::uniform::{UniformSampler as _, UniformUsize};
142    use rand::prelude::{Rng, ThreadRng};
143    use sha2::Sha256;
144    use sha2::digest::OutputSizeUser;
145
146    use super::*;
147    use crate::crypto::hmac_sha256;
148
149    const TEST_HMAC_KEY: &[u8] =
150        &hex!("a83481457efecc69ad1342e21d9c0297f71debbf5c9304b4c1b2e433c1a78f98");
151
152    const TEST_CHUNK_SIZE: usize = 32;
153
154    fn new_incremental(key: &[u8], chunk_size: usize) -> Incremental<Hmac<Sha256>> {
155        let hmac = Hmac::<Sha256>::new_from_slice(key)
156            .expect("Should be able to create a new HMAC instance");
157        Incremental::new(hmac, chunk_size)
158    }
159
160    #[test]
161    #[should_panic]
162    fn chunk_size_zero() {
163        new_incremental(&[], 0);
164    }
165
166    #[test]
167    fn simple_test() {
168        let key = TEST_HMAC_KEY;
169        let input = "this is a simple test input string which is longer than the chunk";
170
171        let bytes = input.as_bytes();
172        let expected = hmac_sha256(key, bytes);
173        let mut incremental = new_incremental(key, TEST_CHUNK_SIZE);
174        let _ = incremental.update(bytes).collect::<Vec<_>>();
175        let digest = incremental.finalize();
176        let actual: [u8; 32] = digest.into();
177        assert_eq!(actual, expected);
178    }
179
180    #[test]
181    fn final_result_should_be_equal_to_non_incremental_hmac() {
182        let key = TEST_HMAC_KEY;
183        proptest!(|(input in ".{0,100}")| {
184            let bytes = input.as_bytes();
185            let expected = hmac_sha256(key, bytes);
186            let mut incremental = new_incremental(key, TEST_CHUNK_SIZE);
187            let _ = incremental.update(bytes).collect::<Vec<_>>();
188            let actual: [u8; 32] = incremental.finalize().into();
189            assert_eq!(actual, expected);
190        });
191    }
192
193    #[test]
194    fn incremental_macs_are_valid() {
195        let key = TEST_HMAC_KEY;
196
197        proptest!(|(input in ".{50,100}")| {
198            let bytes = input.as_bytes();
199            let mut incremental = new_incremental(key, TEST_CHUNK_SIZE);
200
201            // Manually breaking the input in buffer-sized chunks and calculating the HMACs on the
202            // ever-increasing input prefix.
203            let expected: Vec<_> = bytes
204                .chunks(incremental.chunk_size)
205                .scan(Vec::new(), |acc, chunk| {
206                    acc.extend(chunk.iter());
207                    Some(hmac_sha256(key, acc).to_vec())
208                })
209                .collect();
210
211            let mut actual: Vec<Vec<u8>> = bytes
212                .random_chunks(incremental.chunk_size)
213                .flat_map(|chunk| incremental.update(chunk).collect::<Vec<_>>())
214                .map(|out| out.into())
215                .map(|bs: [u8; 32]| bs.to_vec())
216                .collect();
217            // If the input is not an exact multiple of the chunk_size, there are some leftovers in
218            // the incremental that need to be accounted for.
219            if bytes.len() % incremental.chunk_size != 0 {
220                let last_hmac: [u8; 32] = incremental.finalize().into();
221                actual.push(last_hmac.to_vec());
222            }
223            assert_eq!(actual, expected);
224        });
225    }
226
227    #[test]
228    fn validating_simple_test() {
229        let key = TEST_HMAC_KEY;
230        let input = "this is a simple test input string";
231
232        let bytes = input.as_bytes();
233        let mut incremental = new_incremental(key, TEST_CHUNK_SIZE);
234        let mut expected_macs: Vec<_> = incremental.update(bytes).collect();
235        expected_macs.push(incremental.finalize());
236
237        let expected_bytes: Vec<[u8; 32]> =
238            expected_macs.into_iter().map(|mac| mac.into()).collect();
239
240        {
241            let mut validating = new_incremental(key, TEST_CHUNK_SIZE).validating(&expected_bytes);
242            validating
243                .update(bytes)
244                .expect("update: validation should succeed");
245            validating
246                .finalize()
247                .expect("finalize: validation should succeed");
248        }
249
250        {
251            let mut failing_first_update = expected_bytes.clone();
252            failing_first_update
253                .first_mut()
254                .expect("there must be at least one mac")[0] ^= 0xff;
255            let mut validating =
256                new_incremental(key, TEST_CHUNK_SIZE).validating(&failing_first_update);
257            validating.update(bytes).expect_err("MacError");
258        }
259
260        {
261            let mut failing_finalize = expected_bytes.clone();
262            failing_finalize
263                .last_mut()
264                .expect("there must be at least one mac")[0] ^= 0xff;
265            let mut validating =
266                new_incremental(key, TEST_CHUNK_SIZE).validating(&failing_finalize);
267            validating.update(bytes).expect("update should succeed");
268            validating.finalize().expect_err("MacError");
269        }
270
271        {
272            let missing_last_mac = &expected_bytes[0..expected_bytes.len() - 1];
273            let mut validating = new_incremental(key, TEST_CHUNK_SIZE).validating(missing_last_mac);
274            validating.update(bytes).expect("update should succeed");
275            validating.finalize().expect_err("MacError");
276        }
277
278        {
279            let missing_first_mac: Vec<_> = expected_bytes.clone().into_iter().skip(1).collect();
280            let mut validating =
281                new_incremental(key, TEST_CHUNK_SIZE).validating(&missing_first_mac);
282            validating.update(bytes).expect_err("MacError");
283        }
284        // To make clippy happy and allow extending the test in the future
285        std::hint::black_box(expected_bytes);
286    }
287
288    #[test]
289    fn validating_returns_right_size() {
290        let key = TEST_HMAC_KEY;
291        let input = "this is a simple test input string";
292
293        let bytes = input.as_bytes();
294        let mut incremental = new_incremental(key, TEST_CHUNK_SIZE);
295        let mut expected_macs: Vec<_> = incremental.update(bytes).collect();
296        expected_macs.push(incremental.finalize());
297
298        let expected_bytes: Vec<[u8; 32]> =
299            expected_macs.into_iter().map(|mac| mac.into()).collect();
300
301        let mut validating = new_incremental(key, TEST_CHUNK_SIZE).validating(&expected_bytes);
302
303        // Splitting input into chunks of 16 will give us one full incremental chunk + 3 bytes
304        // authenticated by call to finalize.
305        let input_chunks = bytes.chunks(16).collect::<Vec<_>>();
306        assert_eq!(3, input_chunks.len());
307        let expected_remainder = bytes.len() - TEST_CHUNK_SIZE;
308
309        for (expected_size, input) in std::iter::zip([0, TEST_CHUNK_SIZE, 0], input_chunks) {
310            assert_eq!(
311                expected_size,
312                validating
313                    .update(input)
314                    .expect("update: validation should succeed")
315            );
316        }
317        assert_eq!(
318            expected_remainder,
319            validating
320                .finalize()
321                .expect("finalize: validation should succeed")
322        );
323    }
324
325    #[test]
326    fn produce_and_validate() {
327        let key = TEST_HMAC_KEY;
328
329        proptest!(|(input in ".{0,100}")| {
330            let bytes = input.as_bytes();
331            let mut incremental = new_incremental(key, TEST_CHUNK_SIZE);
332            let input_chunks = bytes.random_chunks(incremental.chunk_size*2);
333
334            let mut produced: Vec<[u8; 32]> = input_chunks.clone()
335                .flat_map(|chunk| incremental.update(chunk).collect::<Vec<_>>())
336                .map(|out| out.into())
337                .collect();
338            produced.push(incremental.finalize().into());
339
340            let mut validating = new_incremental(key, TEST_CHUNK_SIZE).validating(&produced);
341            for chunk in input_chunks.clone() {
342                validating.update(chunk).expect("update: validation should succeed");
343            }
344            validating.finalize().expect("finalize: validation should succeed");
345        });
346    }
347
348    const KIBIBYTES: usize = 1024;
349    const MEBIBYTES: usize = 1024 * KIBIBYTES;
350    const GIBIBYTES: usize = 1024 * MEBIBYTES;
351
352    #[test]
353    fn chunk_sizes_sha256() {
354        for (data_size, expected) in [
355            (0, MINIMUM_CHUNK_SIZE),
356            (KIBIBYTES, MINIMUM_CHUNK_SIZE),
357            (10 * KIBIBYTES, MINIMUM_CHUNK_SIZE),
358            (100 * KIBIBYTES, MINIMUM_CHUNK_SIZE),
359            (MEBIBYTES, MINIMUM_CHUNK_SIZE),
360            (10 * MEBIBYTES, MINIMUM_CHUNK_SIZE),
361            (20 * MEBIBYTES, 80 * KIBIBYTES),
362            (100 * MEBIBYTES, 400 * KIBIBYTES),
363            (200 * MEBIBYTES, 800 * KIBIBYTES),
364            (256 * MEBIBYTES, MEBIBYTES),
365            (512 * MEBIBYTES, 2 * MEBIBYTES),
366            (GIBIBYTES, 2 * MEBIBYTES),
367            (2 * GIBIBYTES, 2 * MEBIBYTES),
368        ] {
369            let actual = calculate_chunk_size::<Sha256>(data_size);
370            assert_eq!(actual, expected);
371        }
372    }
373
374    #[test]
375    fn chunk_sizes_sha512() {
376        for (data_size, expected) in [
377            (0, MINIMUM_CHUNK_SIZE),
378            (KIBIBYTES, MINIMUM_CHUNK_SIZE),
379            (10 * KIBIBYTES, MINIMUM_CHUNK_SIZE),
380            (100 * KIBIBYTES, MINIMUM_CHUNK_SIZE),
381            (MEBIBYTES, MINIMUM_CHUNK_SIZE),
382            (10 * MEBIBYTES, 80 * KIBIBYTES),
383            (20 * MEBIBYTES, 160 * KIBIBYTES),
384            (100 * MEBIBYTES, 800 * KIBIBYTES),
385            (200 * MEBIBYTES, 1600 * KIBIBYTES),
386            (256 * MEBIBYTES, 2 * MEBIBYTES),
387            (512 * MEBIBYTES, 2 * MEBIBYTES),
388            (GIBIBYTES, 2 * MEBIBYTES),
389        ] {
390            let actual = calculate_chunk_size::<sha2::Sha512>(data_size);
391            assert_eq!(actual, expected);
392        }
393    }
394
395    #[test]
396    fn total_digest_size_is_never_too_big() {
397        fn total_digest_size(data_size: usize) -> usize {
398            let chunk_size = calculate_chunk_size::<Sha256>(data_size);
399            let num_chunks = std::cmp::max(1, data_size.div_ceil(chunk_size));
400            num_chunks * <Sha256 as OutputSizeUser>::OutputSize::USIZE
401        }
402        let config = ProptestConfig::with_cases(10_000);
403        proptest!(config, |(data_size in 256..256*MEBIBYTES)| {
404            assert!(total_digest_size(data_size) <= 8*KIBIBYTES)
405        });
406        proptest!(|(data_size_mib in 256_usize..2048)| {
407            assert!(total_digest_size(data_size_mib*MEBIBYTES) <= 32*KIBIBYTES)
408        });
409    }
410
411    #[derive(Clone)]
412    struct RandomChunks<'a, T, R: Rng> {
413        base: &'a [T],
414        distribution: UniformUsize,
415        rng: R,
416    }
417
418    impl<'a, T, R: Rng> Iterator for RandomChunks<'a, T, R> {
419        type Item = &'a [T];
420
421        fn next(&mut self) -> Option<Self::Item> {
422            if self.base.is_empty() {
423                None
424            } else {
425                let candidate = self.distribution.sample(&mut self.rng);
426                let chunk_size = std::cmp::min(candidate, self.base.len());
427                let (before, after) = self.base.split_at(chunk_size);
428                self.base = after;
429                Some(before)
430            }
431        }
432    }
433
434    trait RandomChunksIterator<T> {
435        fn random_chunks(&self, max_size: usize) -> RandomChunks<'_, T, ThreadRng>;
436    }
437
438    impl<T> RandomChunksIterator<T> for [T] {
439        fn random_chunks(&self, max_size: usize) -> RandomChunks<'_, T, ThreadRng> {
440            assert!(max_size > 0, "Maximal chunk size should be positive");
441            RandomChunks {
442                base: self,
443                distribution: UniformUsize::new_inclusive(0, max_size).expect("valid range"),
444                rng: Default::default(),
445            }
446        }
447    }
448}
libsignal_protocol/incremental_mac.rs

libsignal_protocol/
incremental_mac.rs
