xmtp_attachments/
crypto.rs

1use aes::Aes256;
2use aes::cipher::{BlockEncrypt, KeyInit as _, StreamCipher, generic_array::GenericArray};
3use ctr::cipher::KeyIvInit as _;
4use ghash::GHash;
5use ghash::universal_hash::UniversalHash as _;
6use subtle::ConstantTimeEq as _;
7use xmtp_content_types::{encryption::derive_key, remote_attachment::RemoteAttachment};
8
9use crate::{AttachmentError, AttachmentFailureCause};
10
11const TAG_LEN: usize = 16;
12const BLOCK_LEN: usize = 16;
13/// GCM starts CTR at block 2, leaving at most 2^32 - 2 data blocks.
14const MAX_GCM_INPUT_BYTES: u64 = (u32::MAX as u64 - 1) * BLOCK_LEN as u64;
15
16/// The random values used by the attachment encryption scheme.
17#[derive(Clone, PartialEq, Eq)]
18pub struct KeyMaterial {
19    pub secret: [u8; 32],
20    pub salt: [u8; 32],
21    pub nonce: [u8; 12],
22}
23
24impl std::fmt::Debug for KeyMaterial {
25    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        formatter.write_str("KeyMaterial { .. }")
27    }
28}
29
30impl KeyMaterial {
31    pub fn random() -> Self {
32        Self {
33            secret: xmtp_common::rand_array(),
34            salt: xmtp_common::rand_array(),
35            nonce: xmtp_common::rand_array(),
36        }
37    }
38
39    pub fn from_remote(ra: &RemoteAttachment) -> Result<Self, AttachmentError> {
40        if ra.content_digest.len() != 64
41            || !ra
42                .content_digest
43                .bytes()
44                .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
45        {
46            return Err(AttachmentError::new(AttachmentFailureCause::Malformed));
47        }
48        Ok(Self {
49            secret: ra
50                .secret
51                .as_slice()
52                .try_into()
53                .map_err(|_| AttachmentError::new(AttachmentFailureCause::Malformed))?,
54            salt: ra
55                .salt
56                .as_slice()
57                .try_into()
58                .map_err(|_| AttachmentError::new(AttachmentFailureCause::Malformed))?,
59            nonce: ra
60                .nonce
61                .as_slice()
62                .try_into()
63                .map_err(|_| AttachmentError::new(AttachmentFailureCause::Malformed))?,
64        })
65    }
66}
67
68struct GcmCore {
69    cipher: ctr::Ctr32BE<Aes256>,
70    hash: GHash,
71    partial: [u8; BLOCK_LEN],
72    partial_len: usize,
73    ciphertext_len: u64,
74    tag_mask: [u8; TAG_LEN],
75}
76
77impl GcmCore {
78    fn new(material: &KeyMaterial) -> Self {
79        // HKDF expansion to 32 bytes cannot fail with SHA-256.
80        let key =
81            derive_key(&material.secret, &material.salt).expect("32-byte HKDF expansion is valid");
82        let aes = Aes256::new((&key).into());
83        let mut h = GenericArray::default();
84        aes.encrypt_block(&mut h);
85        let hash = GHash::new(&h);
86        let mut j0 = [0u8; 16];
87        j0[..12].copy_from_slice(&material.nonce);
88        j0[15] = 1;
89        let mut tag_mask = GenericArray::clone_from_slice(&j0);
90        aes.encrypt_block(&mut tag_mask);
91        let mut counter = j0;
92        counter[15] = 2;
93        Self {
94            cipher: ctr::Ctr32BE::<Aes256>::new((&key).into(), (&counter).into()),
95            hash,
96            partial: [0; BLOCK_LEN],
97            partial_len: 0,
98            ciphertext_len: 0,
99            tag_mask: tag_mask.into(),
100        }
101    }
102
103    fn check_input_len(&self, input_len: usize) -> Result<(), AttachmentError> {
104        let additional = u64::try_from(input_len)
105            .map_err(|_| AttachmentError::new(AttachmentFailureCause::TooLarge))?;
106        let remaining = MAX_GCM_INPUT_BYTES
107            .checked_sub(self.ciphertext_len)
108            .ok_or_else(|| AttachmentError::new(AttachmentFailureCause::TooLarge))?;
109        if additional > remaining {
110            return Err(AttachmentError::new(AttachmentFailureCause::TooLarge));
111        }
112        Ok(())
113    }
114
115    fn authenticate(&mut self, mut input: &[u8]) {
116        self.ciphertext_len += input.len() as u64;
117        if self.partial_len != 0 {
118            let take = (BLOCK_LEN - self.partial_len).min(input.len());
119            self.partial[self.partial_len..self.partial_len + take].copy_from_slice(&input[..take]);
120            self.partial_len += take;
121            input = &input[take..];
122            if self.partial_len == BLOCK_LEN {
123                self.hash.update(&[self.partial.into()]);
124                self.partial_len = 0;
125            } else {
126                return;
127            }
128        }
129        let blocks = input.len() / BLOCK_LEN;
130        self.hash.update_padded(&input[..blocks * BLOCK_LEN]);
131        let tail = &input[blocks * BLOCK_LEN..];
132        self.partial[..tail.len()].copy_from_slice(tail);
133        self.partial_len = tail.len();
134    }
135
136    fn tag(mut self) -> [u8; TAG_LEN] {
137        if self.partial_len != 0 {
138            self.partial[self.partial_len..].fill(0);
139            self.hash.update(&[self.partial.into()]);
140        }
141        let mut lengths = [0u8; BLOCK_LEN];
142        lengths[8..].copy_from_slice(&(self.ciphertext_len * 8).to_be_bytes());
143        self.hash.update(&[lengths.into()]);
144        let digest = self.hash.finalize();
145        let mut tag = self.tag_mask;
146        for (byte, hash_byte) in tag.iter_mut().zip(digest.iter()) {
147            *byte ^= hash_byte;
148        }
149        tag
150    }
151}
152
153/// Encrypts chunks and returns the tag after the final chunk.
154pub struct GcmEncryptor(GcmCore);
155
156impl GcmEncryptor {
157    pub fn new(material: &KeyMaterial) -> Self {
158        Self(GcmCore::new(material))
159    }
160
161    /// Returns `too_large` without changing `out` when the GCM counter limit is exceeded.
162    pub fn update(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), AttachmentError> {
163        self.0.check_input_len(input.len())?;
164        let start = out.len();
165        out.extend_from_slice(input);
166        self.0.cipher.apply_keystream(&mut out[start..]);
167        self.0.authenticate(&out[start..]);
168        Ok(())
169    }
170
171    pub fn finish(self) -> [u8; TAG_LEN] {
172        self.0.tag()
173    }
174}
175
176/// Decrypts chunks while retaining the final 16 bytes as the tag.
177///
178/// `update` releases unauthenticated plaintext. The caller must write it to a
179/// temporary file and make it readable only after `finish` succeeds.
180pub struct GcmDecryptor {
181    core: GcmCore,
182    tail: Vec<u8>,
183}
184
185impl GcmDecryptor {
186    pub fn new(material: &KeyMaterial) -> Self {
187        Self {
188            core: GcmCore::new(material),
189            tail: Vec::with_capacity(TAG_LEN),
190        }
191    }
192
193    /// Returns `too_large` without changing `out` when the GCM counter limit is exceeded.
194    pub fn update(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), AttachmentError> {
195        let release = input.len().saturating_sub(TAG_LEN - self.tail.len());
196        self.core.check_input_len(release)?;
197        let from_tail = release.min(self.tail.len());
198        if from_tail != 0 {
199            let old = self.tail.drain(..from_tail).collect::<Vec<_>>();
200            self.decrypt_part(&old, out);
201        }
202        let from_input = release - from_tail;
203        self.decrypt_part(&input[..from_input], out);
204        self.tail.extend_from_slice(&input[from_input..]);
205        Ok(())
206    }
207
208    fn decrypt_part(&mut self, ciphertext: &[u8], out: &mut Vec<u8>) {
209        self.core.authenticate(ciphertext);
210        let start = out.len();
211        out.extend_from_slice(ciphertext);
212        self.core.cipher.apply_keystream(&mut out[start..]);
213    }
214
215    pub fn finish(self) -> Result<(), AttachmentError> {
216        let Self { core, tail } = self;
217        if tail.len() != TAG_LEN || !bool::from(core.tag().as_slice().ct_eq(tail.as_slice())) {
218            return Err(AttachmentError::new(
219                AttachmentFailureCause::DecryptionFailed,
220            ));
221        }
222        Ok(())
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use aes_gcm::{Aes256Gcm, KeyInit as _, aead::Aead as _};
229
230    use super::*;
231
232    #[xmtp_common::test(unwrap_try = true)]
233    fn key_material_debug_hides_secrets() {
234        let material = KeyMaterial {
235            secret: [0xab; 32],
236            salt: [0xcd; 32],
237            nonce: [0xef; 12],
238        };
239        let debug = format!("{material:?}");
240        assert_eq!(debug, "KeyMaterial { .. }");
241        assert!(!debug.contains(&hex::encode(material.secret)));
242    }
243
244    fn check_lengths(material: &KeyMaterial, plaintext: &[u8], start: usize, end: usize) {
245        let key =
246            derive_key(&material.secret, &material.salt).expect("32-byte HKDF expansion is valid");
247        let cipher = Aes256Gcm::new((&key).into());
248        let mut seed = 0x91c3_7d52_u64 ^ start as u64;
249        for len in start..=end {
250            let expected = cipher
251                .encrypt((&material.nonce).into(), &plaintext[..len])
252                .expect("fixed key and nonce");
253            let mut encryptor = GcmEncryptor::new(material);
254            let mut actual = Vec::new();
255            let mut at = 0;
256            while at < len {
257                seed ^= seed << 13;
258                seed ^= seed >> 7;
259                seed ^= seed << 17;
260                let end = (at + 1 + (seed as usize % 32_768)).min(len);
261                encryptor
262                    .update(&plaintext[at..end], &mut actual)
263                    .expect("plaintext length is below the GCM limit");
264                at = end;
265            }
266            actual.extend_from_slice(&encryptor.finish());
267            assert_eq!(actual, expected, "length {len}");
268
269            let mut decryptor = GcmDecryptor::new(material);
270            let mut decoded = Vec::new();
271            let mut at = 0;
272            while at < actual.len() {
273                seed ^= seed << 13;
274                seed ^= seed >> 7;
275                seed ^= seed << 17;
276                let end = (at + 1 + (seed as usize % 24_576)).min(actual.len());
277                decryptor
278                    .update(&actual[at..end], &mut decoded)
279                    .expect("ciphertext length is below the GCM limit");
280                at = end;
281            }
282            decryptor.finish().expect("valid ciphertext and tag");
283            assert_eq!(decoded, plaintext[..len], "length {len}");
284
285            if [0, 1, 15, 16, 17, 1024, 70_000].contains(&len) {
286                let mut positions = vec![len, actual.len() - 1];
287                if len != 0 {
288                    positions.extend([0, len - 1]);
289                }
290                for pos in positions {
291                    for bit in 0..8 {
292                        let mut changed = actual.clone();
293                        changed[pos] ^= 1 << bit;
294                        let mut decryptor = GcmDecryptor::new(material);
295                        let mut out = Vec::new();
296                        decryptor
297                            .update(&changed, &mut out)
298                            .expect("ciphertext length is below the GCM limit");
299                        assert!(matches!(
300                            decryptor.finish(),
301                            Err(AttachmentError {
302                                cause: AttachmentFailureCause::DecryptionFailed,
303                                ..
304                            })
305                        ));
306                    }
307                }
308            }
309        }
310    }
311
312    // verifies: ATCH-012
313    #[xmtp_common::test(unwrap_try = true)]
314    async fn gcm_stream_matches_one_shot() {
315        let material = KeyMaterial {
316            secret: [7; 32],
317            salt: [9; 32],
318            nonce: [11; 12],
319        };
320        // The native run checks every required length. Wasm runs a shorter
321        // smoke pass because debug AES in the browser is much slower.
322        let max_len = if cfg!(target_arch = "wasm32") {
323            1024
324        } else {
325            70_000
326        };
327        let mut plaintext = vec![0u8; max_len];
328        let mut seed = 0x91c3_7d52_u64;
329        for byte in &mut plaintext {
330            seed ^= seed << 13;
331            seed ^= seed >> 7;
332            seed ^= seed << 17;
333            *byte = seed as u8;
334        }
335        #[cfg(target_arch = "wasm32")]
336        check_lengths(&material, &plaintext, 0, max_len);
337        #[cfg(not(target_arch = "wasm32"))]
338        std::thread::scope(|scope| {
339            const WORKERS: usize = 10;
340            let mut handles = Vec::with_capacity(WORKERS);
341            // Total work is proportional to the sum of plaintext lengths.
342            let boundary = |worker: usize| {
343                ((max_len + 1) as f64 * (worker as f64 / WORKERS as f64).sqrt()) as usize
344            };
345            for worker in 0..WORKERS {
346                let start = boundary(worker);
347                let end = boundary(worker + 1) - 1;
348                let material = &material;
349                let plaintext = &plaintext;
350                handles.push(scope.spawn(move || check_lengths(material, plaintext, start, end)));
351            }
352            for handle in handles {
353                handle.join().expect("GCM worker completed");
354            }
355        });
356    }
357
358    #[xmtp_common::test(unwrap_try = true)]
359    async fn gcm_counter_limit_rejects_before_output() {
360        let material = KeyMaterial {
361            secret: [7; 32],
362            salt: [9; 32],
363            nonce: [11; 12],
364        };
365
366        let mut encryptor = GcmEncryptor::new(&material);
367        encryptor.0.ciphertext_len = MAX_GCM_INPUT_BYTES - 1;
368        let mut encrypted = vec![0x5a];
369        assert_eq!(
370            encryptor.update(&[1, 2], &mut encrypted).unwrap_err().cause,
371            AttachmentFailureCause::TooLarge
372        );
373        assert_eq!(encrypted, [0x5a]);
374        assert_eq!(encryptor.0.ciphertext_len, MAX_GCM_INPUT_BYTES - 1);
375        encryptor.update(&[1], &mut encrypted)?;
376        assert_eq!(encrypted.len(), 2);
377        assert_eq!(encryptor.0.ciphertext_len, MAX_GCM_INPUT_BYTES);
378        assert_eq!(
379            encryptor.update(&[2], &mut encrypted).unwrap_err().cause,
380            AttachmentFailureCause::TooLarge
381        );
382        assert_eq!(encrypted.len(), 2);
383
384        let mut decryptor = GcmDecryptor::new(&material);
385        decryptor.core.ciphertext_len = MAX_GCM_INPUT_BYTES - 1;
386        decryptor.tail = vec![0x3c; TAG_LEN];
387        let mut decrypted = vec![0x5a];
388        assert_eq!(
389            decryptor.update(&[1, 2], &mut decrypted).unwrap_err().cause,
390            AttachmentFailureCause::TooLarge
391        );
392        assert_eq!(decrypted, [0x5a]);
393        assert_eq!(decryptor.tail, [0x3c; TAG_LEN]);
394        assert_eq!(decryptor.core.ciphertext_len, MAX_GCM_INPUT_BYTES - 1);
395        decryptor.update(&[1], &mut decrypted)?;
396        assert_eq!(decrypted.len(), 2);
397        assert_eq!(decryptor.core.ciphertext_len, MAX_GCM_INPUT_BYTES);
398        assert_eq!(
399            decryptor.update(&[2], &mut decrypted).unwrap_err().cause,
400            AttachmentFailureCause::TooLarge
401        );
402        assert_eq!(decrypted.len(), 2);
403    }
404}
xmtp_attachments/crypto.rs

xmtp_attachments/
crypto.rs
