Skip to main content

xmtp_content_types/
encryption.rs

1use aes_gcm::{Aes256Gcm, KeyInit, aead::Aead};
2use hkdf::Hkdf;
3use prost::Message;
4use sha2::Sha256;
5use xmtp_proto::xmtp::mls::message_contents::EncodedContent;
6
7use crate::CodecError;
8
9/// Size of the HKDF salt in bytes (256-bit)
10pub const HKDF_SALT_SIZE: usize = 32;
11
12/// Size of the AES-GCM nonce in bytes (96-bit)
13pub const AES_GCM_NONCE_SIZE: usize = 12;
14
15/// Size of the AES-GCM authentication tag in bytes (128-bit)
16pub const AES_GCM_TAG_SIZE: usize = 16;
17
18/// Size of the encryption secret in bytes (256-bit)
19pub const SECRET_SIZE: usize = 32;
20
21/// Encrypted payload containing the ciphertext and encryption parameters.
22#[derive(Debug, Clone)]
23pub struct EncryptedPayload {
24    /// The encrypted content (ciphertext + 16-byte auth tag)
25    pub payload: Vec<u8>,
26    /// The 32-byte salt used for HKDF key derivation
27    pub salt: Vec<u8>,
28    /// The 12-byte nonce used for AES-GCM encryption
29    pub nonce: Vec<u8>,
30}
31
32/// Parameters for the existing remote-attachment encryption scheme.
33#[derive(Clone)]
34pub struct EncryptionKeys {
35    pub secret: Vec<u8>,
36    pub salt: Vec<u8>,
37    pub nonce: Vec<u8>,
38    /// Lowercase hexadecimal SHA-256 of ciphertext including the tag.
39    pub digest: String,
40    /// Ciphertext length including the 16-byte tag.
41    pub length: u64,
42}
43
44impl std::fmt::Debug for EncryptionKeys {
45    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        formatter
47            .debug_struct("EncryptionKeys")
48            .field("secret", &"<redacted>")
49            .field("salt", &self.salt)
50            .field("nonce", &self.nonce)
51            .field("digest", &self.digest)
52            .field("length", &self.length)
53            .finish()
54    }
55}
56
57/// Ciphertext and the parameters needed to decrypt it.
58#[derive(Debug, Clone)]
59pub struct EncryptedEncodedContent {
60    pub ciphertext: Vec<u8>,
61    pub keys: EncryptionKeys,
62}
63
64/// Encrypts arbitrary bytes with the remote-attachment scheme.
65pub fn encrypt_bytes(bytes: &[u8]) -> Result<EncryptedEncodedContent, CodecError> {
66    let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
67    let encrypted = encrypt(bytes, &secret)?;
68    Ok(EncryptedEncodedContent {
69        keys: EncryptionKeys {
70            secret: secret.to_vec(),
71            salt: encrypted.salt,
72            nonce: encrypted.nonce,
73            digest: hex::encode(sha256(&encrypted.payload)),
74            length: encrypted.payload.len() as u64,
75        },
76        ciphertext: encrypted.payload,
77    })
78}
79
80/// Checks the length and digest before decrypting arbitrary bytes.
81pub fn decrypt_bytes(ciphertext: &[u8], keys: &EncryptionKeys) -> Result<Vec<u8>, CodecError> {
82    if ciphertext.len() as u64 != keys.length {
83        return Err(CodecError::Decode(
84            "encrypted content length mismatch".into(),
85        ));
86    }
87    let actual_digest = hex::encode(sha256(ciphertext));
88    if actual_digest != keys.digest {
89        return Err(CodecError::Decode(format!(
90            "content digest mismatch: expected {}, got {}",
91            keys.digest, actual_digest
92        )));
93    }
94    decrypt(
95        &EncryptedPayload {
96            payload: ciphertext.to_vec(),
97            salt: keys.salt.clone(),
98            nonce: keys.nonce.clone(),
99        },
100        &keys.secret,
101    )
102}
103
104/// Encrypts the protobuf envelope, not only its `content` field.
105pub fn encrypt_encoded_content(
106    content: EncodedContent,
107) -> Result<EncryptedEncodedContent, CodecError> {
108    encrypt_bytes(&content.encode_to_vec())
109}
110
111/// Decrypts a protobuf envelope encrypted with [`encrypt_encoded_content`].
112pub fn decrypt_encoded_content(
113    ciphertext: &[u8],
114    keys: &EncryptionKeys,
115) -> Result<EncodedContent, CodecError> {
116    let bytes = decrypt_bytes(ciphertext, keys)?;
117    EncodedContent::decode(bytes.as_slice())
118        .map_err(|error| CodecError::Decode(format!("failed to decode EncodedContent: {error}")))
119}
120
121/// Encrypts plaintext using AES-256-GCM with HKDF-SHA256 key derivation.
122// implements: CTYPE-015
123pub fn encrypt(plaintext: &[u8], secret: &[u8]) -> Result<EncryptedPayload, CodecError> {
124    // Generate random salt and nonce
125    let salt: [u8; HKDF_SALT_SIZE] = xmtp_common::rand_array();
126    let nonce: [u8; AES_GCM_NONCE_SIZE] = xmtp_common::rand_array();
127
128    // Derive AES-256 key using HKDF-SHA256
129    let key = derive_key(secret, &salt).map_err(CodecError::Encode)?;
130
131    // Create cipher
132    let cipher = Aes256Gcm::new_from_slice(&key)
133        .map_err(|e| CodecError::Encode(format!("failed to create cipher: {e}")))?;
134
135    let ciphertext = cipher
136        .encrypt((&nonce).into(), plaintext)
137        .map_err(|e| CodecError::Encode(format!("encryption failed: {e}")))?;
138
139    Ok(EncryptedPayload {
140        payload: ciphertext,
141        salt: salt.to_vec(),
142        nonce: nonce.to_vec(),
143    })
144}
145
146/// Decrypts ciphertext that was encrypted with [`encrypt`].
147// implements: CTYPE-015
148pub fn decrypt(encrypted: &EncryptedPayload, secret: &[u8]) -> Result<Vec<u8>, CodecError> {
149    // Validate salt and nonce lengths
150    if encrypted.salt.len() != HKDF_SALT_SIZE {
151        return Err(CodecError::Decode(format!(
152            "invalid salt length: expected {}, got {}",
153            HKDF_SALT_SIZE,
154            encrypted.salt.len()
155        )));
156    }
157    if encrypted.nonce.len() != AES_GCM_NONCE_SIZE {
158        return Err(CodecError::Decode(format!(
159            "invalid nonce length: expected {}, got {}",
160            AES_GCM_NONCE_SIZE,
161            encrypted.nonce.len()
162        )));
163    }
164
165    // Derive AES-256 key using HKDF-SHA256
166    let key = derive_key(secret, &encrypted.salt).map_err(CodecError::Decode)?;
167
168    // Create cipher
169    let cipher = Aes256Gcm::new_from_slice(&key)
170        .map_err(|e| CodecError::Decode(format!("failed to create cipher: {e}")))?;
171
172    let nonce: &[u8; AES_GCM_NONCE_SIZE] = encrypted
173        .nonce
174        .as_slice()
175        .try_into()
176        .expect("nonce length already validated");
177
178    cipher
179        .decrypt(nonce.into(), encrypted.payload.as_slice())
180        .map_err(|e| CodecError::Decode(format!("decryption failed: {e}")))
181}
182
183/// Computes the SHA-256 hash of the given bytes.
184pub fn sha256(bytes: &[u8]) -> Vec<u8> {
185    use sha2::Digest;
186    let mut hasher = Sha256::new();
187    hasher.update(bytes);
188    hasher.finalize().to_vec()
189}
190
191/// Derives an AES-256 key from a secret and salt using HKDF-SHA256.
192pub fn derive_key(secret: &[u8], salt: &[u8]) -> Result<[u8; 32], String> {
193    let hkdf = Hkdf::<Sha256>::new(Some(salt), secret);
194
195    let mut key = [0u8; 32];
196    // Empty info, matching the TypeScript implementation
197    hkdf.expand(&[], &mut key)
198        .map_err(|e| format!("HKDF key derivation failed: {e}"))?;
199
200    Ok(key)
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::{
207        ContentCodec,
208        attachment::{Attachment, AttachmentCodec},
209        remote_attachment::{RemoteAttachment, decrypt_attachment},
210    };
211
212    #[xmtp_common::test(unwrap_try = true)]
213    async fn encoded_content_encryption_matches_attachment_scheme() {
214        let attachment = Attachment {
215            filename: Some("file.txt".into()),
216            mime_type: "text/plain".into(),
217            content: b"encrypted content".to_vec(),
218        };
219        let encoded = AttachmentCodec::encode(attachment.clone())?;
220        let encrypted = encrypt_encoded_content(encoded.clone())?;
221        assert_eq!(
222            decrypt_encoded_content(&encrypted.ciphertext, &encrypted.keys)?,
223            encoded
224        );
225        let remote = RemoteAttachment {
226            filename: attachment.filename,
227            content_length: Some(encrypted.keys.length as u32),
228            url: String::new(),
229            content_digest: encrypted.keys.digest,
230            secret: encrypted.keys.secret,
231            salt: encrypted.keys.salt,
232            nonce: encrypted.keys.nonce,
233            scheme: String::new(),
234        };
235        assert_eq!(
236            decrypt_attachment(&encrypted.ciphertext, &remote)?.content,
237            attachment.content
238        );
239
240        let bytes = encrypt_bytes(b"raw bytes")?;
241        assert_eq!(decrypt_bytes(&bytes.ciphertext, &bytes.keys)?, b"raw bytes");
242    }
243
244    /// Fixed fields from the original XMTP JavaScript ciphertext fixture.
245    #[xmtp_common::test(unwrap_try = true)]
246    async fn decrypt_bytes_reads_the_javascript_ciphertext_fixture() {
247        let keys = EncryptionKeys {
248            secret: vec![1, 2, 3, 4],
249            salt: vec![
250                23, 10, 217, 190, 235, 216, 145, 38, 49, 224, 165, 169, 22, 55, 152, 150, 176, 65,
251                207, 91, 45, 45, 16, 171, 146, 125, 143, 60, 152, 128, 0, 120,
252            ],
253            nonce: vec![219, 247, 207, 184, 141, 179, 171, 100, 251, 171, 120, 137],
254            digest: "68bd4ef27c192622cb66b1da180ddd78e7f928a9b31c319257ab297bbf186da4".into(),
255            length: 19,
256        };
257        let ciphertext = [
258            216, 215, 152, 167, 118, 59, 93, 177, 53, 242, 147, 10, 87, 143, 27, 245, 154, 169, 109,
259        ];
260        assert_eq!(decrypt_bytes(&ciphertext, &keys)?, [5, 5, 5]);
261    }
262
263    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
264    #[cfg_attr(not(target_arch = "wasm32"), test)]
265    fn test_encrypt_decrypt_roundtrip() {
266        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
267        let plaintext = b"Hello, XMTP remote attachments!";
268
269        let encrypted = encrypt(plaintext, &secret).unwrap();
270        let decrypted = decrypt(&encrypted, &secret).unwrap();
271
272        assert_eq!(plaintext.as_slice(), decrypted.as_slice());
273    }
274
275    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
276    #[cfg_attr(not(target_arch = "wasm32"), test)]
277    fn test_decrypt_wrong_secret_fails() {
278        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
279        let wrong_secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
280        let plaintext = b"Secret message";
281
282        let encrypted = encrypt(plaintext, &secret).unwrap();
283        let result = decrypt(&encrypted, &wrong_secret);
284
285        assert!(result.is_err());
286    }
287
288    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
289    #[cfg_attr(not(target_arch = "wasm32"), test)]
290    // verifies: CTYPE-015
291    fn test_encrypt_produces_different_output_each_time() {
292        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
293        let plaintext = b"Same message";
294
295        let encrypted1 = encrypt(plaintext, &secret).unwrap();
296        let encrypted2 = encrypt(plaintext, &secret).unwrap();
297
298        // Due to random salt and nonce, encrypted output should differ
299        assert_ne!(encrypted1.payload, encrypted2.payload);
300        assert_ne!(encrypted1.salt, encrypted2.salt);
301        assert_ne!(encrypted1.nonce, encrypted2.nonce);
302    }
303
304    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
305    #[cfg_attr(not(target_arch = "wasm32"), test)]
306    // verifies: CTYPE-015
307    fn test_encrypted_payload_sizes() {
308        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
309        let plaintext = b"Test message";
310
311        let encrypted = encrypt(plaintext, &secret).unwrap();
312
313        assert_eq!(encrypted.salt.len(), HKDF_SALT_SIZE);
314        assert_eq!(encrypted.nonce.len(), AES_GCM_NONCE_SIZE);
315        // Ciphertext is plaintext + authentication tag
316        assert_eq!(encrypted.payload.len(), plaintext.len() + AES_GCM_TAG_SIZE);
317    }
318
319    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
320    #[cfg_attr(not(target_arch = "wasm32"), test)]
321    fn test_invalid_salt_length() {
322        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
323        let encrypted = EncryptedPayload {
324            payload: vec![0u8; 32],
325            salt: vec![0u8; 16], // Wrong size
326            nonce: vec![0u8; AES_GCM_NONCE_SIZE],
327        };
328
329        let result = decrypt(&encrypted, &secret);
330        assert!(result.is_err());
331        assert!(
332            result
333                .unwrap_err()
334                .to_string()
335                .contains("invalid salt length")
336        );
337    }
338
339    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
340    #[cfg_attr(not(target_arch = "wasm32"), test)]
341    fn test_invalid_nonce_length() {
342        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
343        let encrypted = EncryptedPayload {
344            payload: vec![0u8; 32],
345            salt: vec![0u8; HKDF_SALT_SIZE],
346            nonce: vec![0u8; 8], // Wrong size
347        };
348
349        let result = decrypt(&encrypted, &secret);
350        assert!(result.is_err());
351        assert!(
352            result
353                .unwrap_err()
354                .to_string()
355                .contains("invalid nonce length")
356        );
357    }
358}