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
9pub const HKDF_SALT_SIZE: usize = 32;
11
12pub const AES_GCM_NONCE_SIZE: usize = 12;
14
15pub const AES_GCM_TAG_SIZE: usize = 16;
17
18pub const SECRET_SIZE: usize = 32;
20
21#[derive(Debug, Clone)]
23pub struct EncryptedPayload {
24 pub payload: Vec<u8>,
26 pub salt: Vec<u8>,
28 pub nonce: Vec<u8>,
30}
31
32#[derive(Clone)]
34pub struct EncryptionKeys {
35 pub secret: Vec<u8>,
36 pub salt: Vec<u8>,
37 pub nonce: Vec<u8>,
38 pub digest: String,
40 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#[derive(Debug, Clone)]
59pub struct EncryptedEncodedContent {
60 pub ciphertext: Vec<u8>,
61 pub keys: EncryptionKeys,
62}
63
64pub 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
80pub 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
104pub fn encrypt_encoded_content(
106 content: EncodedContent,
107) -> Result<EncryptedEncodedContent, CodecError> {
108 encrypt_bytes(&content.encode_to_vec())
109}
110
111pub 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
121pub fn encrypt(plaintext: &[u8], secret: &[u8]) -> Result<EncryptedPayload, CodecError> {
124 let salt: [u8; HKDF_SALT_SIZE] = xmtp_common::rand_array();
126 let nonce: [u8; AES_GCM_NONCE_SIZE] = xmtp_common::rand_array();
127
128 let key = derive_key(secret, &salt).map_err(CodecError::Encode)?;
130
131 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
146pub fn decrypt(encrypted: &EncryptedPayload, secret: &[u8]) -> Result<Vec<u8>, CodecError> {
149 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 let key = derive_key(secret, &encrypted.salt).map_err(CodecError::Decode)?;
167
168 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
183pub 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
191pub 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 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 #[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 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 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 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 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], 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], };
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}