xmtp_content_types/
file_encryption.rs

1//! Native file encryption compatible with one-tag AES-256-GCM attachments.
2
3use std::{
4    fs::File,
5    io::{Read, Write},
6    path::{Path, PathBuf},
7};
8
9use aes::{
10    Aes256,
11    cipher::{Block, BlockEncrypt, KeyInit, KeyIvInit, StreamCipher},
12};
13use ctr::Ctr32BE;
14use ghash::{GHash, universal_hash::UniversalHash};
15use sha2::{Digest, Sha256};
16use subtle::ConstantTimeEq;
17
18use crate::{
19    CodecError,
20    encryption::{
21        self, AES_GCM_NONCE_SIZE, AES_GCM_TAG_SIZE, EncryptionKeys, HKDF_SALT_SIZE, SECRET_SIZE,
22    },
23};
24
25const FILE_CHUNK_BYTES: usize = 64 * 1024;
26// NIST SP 800-38D limits one GCM plaintext to 2^36 - 32 bytes.
27const MAX_GCM_BYTES: u64 = (1u64 << 36) - 32;
28
29struct GcmStream {
30    ctr: Ctr32BE<Aes256>,
31    ghash: GHash,
32    mask: Block<Aes256>,
33}
34
35impl GcmStream {
36    fn new(secret: &[u8], salt: &[u8], nonce: &[u8]) -> Result<Self, CodecError> {
37        if salt.len() != HKDF_SALT_SIZE || nonce.len() != AES_GCM_NONCE_SIZE {
38            return Err(CodecError::Decode(
39                "invalid file encryption salt or nonce".into(),
40            ));
41        }
42        let key = encryption::derive_key(secret, salt).map_err(CodecError::Decode)?;
43        let aes = Aes256::new_from_slice(&key)
44            .map_err(|e| CodecError::Decode(format!("invalid AES key: {e}")))?;
45        let mut h = Block::<Aes256>::default();
46        aes.encrypt_block(&mut h);
47        let ghash = GHash::new(ghash::Key::from_slice(&h));
48        let mut mask = Block::<Aes256>::default();
49        mask[..AES_GCM_NONCE_SIZE].copy_from_slice(nonce);
50        mask[15] = 1;
51        aes.encrypt_block(&mut mask);
52        let mut iv = Block::<Aes256>::default();
53        iv[..AES_GCM_NONCE_SIZE].copy_from_slice(nonce);
54        iv[15] = 2;
55        let ctr = Ctr32BE::<Aes256>::new((&key).into(), &iv);
56        Ok(Self { ctr, ghash, mask })
57    }
58
59    fn encrypt_chunk(&mut self, chunk: &mut [u8]) {
60        self.ctr.apply_keystream(chunk);
61        self.ghash.update_padded(chunk);
62    }
63
64    fn decrypt_chunk(&mut self, chunk: &mut [u8]) {
65        self.ghash.update_padded(chunk);
66        self.ctr.apply_keystream(chunk);
67    }
68
69    fn finish(mut self, ciphertext_len: u64) -> [u8; AES_GCM_TAG_SIZE] {
70        let mut lengths = ghash::Block::default();
71        lengths[8..].copy_from_slice(&(ciphertext_len * 8).to_be_bytes());
72        self.ghash.update(&[lengths]);
73        let mut tag = [0u8; AES_GCM_TAG_SIZE];
74        tag.copy_from_slice(&self.ghash.finalize());
75        for (byte, mask) in tag.iter_mut().zip(self.mask.iter()) {
76            *byte ^= mask;
77        }
78        tag
79    }
80}
81
82fn output_parent(output: &Path) -> PathBuf {
83    output
84        .parent()
85        .filter(|parent| !parent.as_os_str().is_empty())
86        .unwrap_or_else(|| Path::new("."))
87        .to_path_buf()
88}
89
90fn encrypt_stream(
91    input: &mut impl Read,
92    output: &mut impl Write,
93    secret: &[u8],
94    salt: &[u8],
95    nonce: &[u8],
96) -> Result<(String, u64), CodecError> {
97    let mut gcm = GcmStream::new(secret, salt, nonce)?;
98    let mut digest = Sha256::new();
99    let mut chunk = [0u8; FILE_CHUNK_BYTES];
100    let mut length = 0u64;
101    loop {
102        let mut read = 0;
103        while read < FILE_CHUNK_BYTES {
104            let count = input
105                .read(&mut chunk[read..])
106                .map_err(|e| CodecError::Encode(e.to_string()))?;
107            if count == 0 {
108                break;
109            }
110            read += count;
111        }
112        if read == 0 {
113            break;
114        }
115        length = length
116            .checked_add(read as u64)
117            .filter(|value| *value <= MAX_GCM_BYTES)
118            .ok_or_else(|| CodecError::Encode("file exceeds AES-GCM size limit".into()))?;
119        gcm.encrypt_chunk(&mut chunk[..read]);
120        output
121            .write_all(&chunk[..read])
122            .map_err(|e| CodecError::Encode(e.to_string()))?;
123        digest.update(&chunk[..read]);
124    }
125    let tag = gcm.finish(length);
126    output
127        .write_all(&tag)
128        .map_err(|e| CodecError::Encode(e.to_string()))?;
129    digest.update(tag);
130    Ok((
131        hex::encode(digest.finalize()),
132        length + AES_GCM_TAG_SIZE as u64,
133    ))
134}
135
136/// Encrypts a file in 64 KiB chunks and writes one attachment-compatible tag.
137pub fn encrypt_file(input: &Path, output: &Path) -> Result<EncryptionKeys, CodecError> {
138    let mut source = File::open(input).map_err(|e| CodecError::Encode(e.to_string()))?;
139    let mut target = tempfile::NamedTempFile::new_in(output_parent(output))
140        .map_err(|e| CodecError::Encode(e.to_string()))?;
141    let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
142    let salt: [u8; HKDF_SALT_SIZE] = xmtp_common::rand_array();
143    let nonce: [u8; AES_GCM_NONCE_SIZE] = xmtp_common::rand_array();
144    let (digest, length) = encrypt_stream(&mut source, &mut target, &secret, &salt, &nonce)?;
145    target
146        .persist(output)
147        .map_err(|e| CodecError::Encode(e.to_string()))?;
148    Ok(EncryptionKeys {
149        secret: secret.to_vec(),
150        salt: salt.to_vec(),
151        nonce: nonce.to_vec(),
152        digest,
153        length,
154    })
155}
156
157/// Verifies the digest and tag before moving the plaintext to `output`.
158pub fn decrypt_file(input: &Path, output: &Path, keys: &EncryptionKeys) -> Result<(), CodecError> {
159    let mut source = File::open(input).map_err(|e| CodecError::Decode(e.to_string()))?;
160    let total_len = source
161        .metadata()
162        .map_err(|e| CodecError::Decode(e.to_string()))?
163        .len();
164    if total_len != keys.length
165        || total_len < AES_GCM_TAG_SIZE as u64
166        || total_len - AES_GCM_TAG_SIZE as u64 > MAX_GCM_BYTES
167    {
168        return Err(CodecError::Decode("invalid encrypted file length".into()));
169    }
170    let mut target = tempfile::NamedTempFile::new_in(output_parent(output))
171        .map_err(|e| CodecError::Decode(e.to_string()))?;
172    let mut gcm = GcmStream::new(&keys.secret, &keys.salt, &keys.nonce)?;
173    let mut digest = Sha256::new();
174    let mut remaining = total_len - AES_GCM_TAG_SIZE as u64;
175    let mut chunk = [0u8; FILE_CHUNK_BYTES];
176    while remaining > 0 {
177        let size = remaining.min(FILE_CHUNK_BYTES as u64) as usize;
178        source
179            .read_exact(&mut chunk[..size])
180            .map_err(|e| CodecError::Decode(e.to_string()))?;
181        digest.update(&chunk[..size]);
182        gcm.decrypt_chunk(&mut chunk[..size]);
183        target
184            .write_all(&chunk[..size])
185            .map_err(|e| CodecError::Decode(e.to_string()))?;
186        remaining -= size as u64;
187    }
188    let mut tag = [0u8; AES_GCM_TAG_SIZE];
189    source
190        .read_exact(&mut tag)
191        .map_err(|e| CodecError::Decode(e.to_string()))?;
192    digest.update(tag);
193    if hex::encode(digest.finalize()) != keys.digest {
194        return Err(CodecError::Decode("content digest mismatch".into()));
195    }
196    let expected_tag = gcm.finish(total_len - AES_GCM_TAG_SIZE as u64);
197    if !bool::from(expected_tag.ct_eq(&tag)) {
198        return Err(CodecError::Decode("AES-GCM authentication failed".into()));
199    }
200    target
201        .persist(output)
202        .map_err(|e| CodecError::Decode(e.to_string()))?;
203    Ok(())
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209    use crate::{
210        ContentCodec,
211        attachment::{Attachment, AttachmentCodec},
212        remote_attachment::{RemoteAttachment, decrypt_attachment},
213    };
214    use prost::Message;
215
216    // verifies: CTYPE-015
217    #[xmtp_common::test(unwrap_try = true)]
218    async fn decrypt_file_wrong_key_leaves_no_output() {
219        let dir = tempfile::tempdir()?;
220        let source = dir.path().join("plain");
221        let encrypted = dir.path().join("encrypted");
222        let output = dir.path().join("decrypted");
223        std::fs::write(&source, b"secret file")?;
224        let mut keys = encrypt_file(&source, &encrypted)?;
225        keys.secret[0] ^= 1;
226        assert!(decrypt_file(&encrypted, &output, &keys).is_err());
227        assert!(!output.exists());
228        assert_eq!(std::fs::read_dir(dir.path())?.count(), 2);
229        let mut keys = encrypt_file(&source, &encrypted)?;
230        keys.digest = "bad digest".into();
231        assert!(decrypt_file(&encrypted, &output, &keys).is_err());
232        assert!(!output.exists());
233    }
234
235    struct BoundedReader<'a> {
236        bytes: &'a [u8],
237        largest: usize,
238    }
239    impl Read for BoundedReader<'_> {
240        fn read(&mut self, output: &mut [u8]) -> std::io::Result<usize> {
241            self.largest = self.largest.max(output.len());
242            self.bytes.read(output)
243        }
244    }
245
246    #[xmtp_common::test(unwrap_try = true)]
247    async fn file_encryption_reads_bounded_chunks() {
248        let bytes = vec![42u8; FILE_CHUNK_BYTES * 2 + 19];
249        let mut reader = BoundedReader {
250            bytes: &bytes,
251            largest: 0,
252        };
253        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
254        let salt: [u8; HKDF_SALT_SIZE] = xmtp_common::rand_array();
255        let nonce: [u8; AES_GCM_NONCE_SIZE] = xmtp_common::rand_array();
256        let mut ciphertext = Vec::new();
257        let (digest, length) =
258            encrypt_stream(&mut reader, &mut ciphertext, &secret, &salt, &nonce)?;
259        assert_eq!(reader.largest, FILE_CHUNK_BYTES);
260        let keys = EncryptionKeys {
261            secret: secret.to_vec(),
262            salt: salt.to_vec(),
263            nonce: nonce.to_vec(),
264            digest,
265            length,
266        };
267        assert_eq!(encryption::decrypt_bytes(&ciphertext, &keys)?, bytes);
268
269        let dir = tempfile::tempdir()?;
270        let source = dir.path().join("plain");
271        let encrypted = dir.path().join("encrypted");
272        let output = dir.path().join("decrypted");
273        let content = AttachmentCodec::encode(Attachment {
274            filename: None,
275            mime_type: "application/octet-stream".into(),
276            content: bytes,
277        })?;
278        std::fs::write(&source, content.encode_to_vec())?;
279        let keys = encrypt_file(&source, &encrypted)?;
280        decrypt_file(&encrypted, &output, &keys)?;
281        assert_eq!(std::fs::read(&source)?, std::fs::read(&output)?);
282        let remote = RemoteAttachment {
283            filename: None,
284            content_length: Some(keys.length as u32),
285            url: String::new(),
286            content_digest: keys.digest,
287            secret: keys.secret,
288            salt: keys.salt,
289            nonce: keys.nonce,
290            scheme: String::new(),
291        };
292        assert_eq!(
293            decrypt_attachment(&std::fs::read(encrypted)?, &remote)?
294                .content
295                .len(),
296            FILE_CHUNK_BYTES * 2 + 19
297        );
298    }
299}
xmtp_content_types/file_encryption.rs

xmtp_content_types/
file_encryption.rs
