xmtp_content_types/
encryption.rs

1use aes_gcm::{Aes256Gcm, KeyInit, aead::Aead};
2use hkdf::Hkdf;
3use sha2::Sha256;
4
5use crate::CodecError;
6
7/// Size of the HKDF salt in bytes (256-bit)
8pub const HKDF_SALT_SIZE: usize = 32;
9
10/// Size of the AES-GCM nonce in bytes (96-bit)
11pub const AES_GCM_NONCE_SIZE: usize = 12;
12
13/// Size of the AES-GCM authentication tag in bytes (128-bit)
14pub const AES_GCM_TAG_SIZE: usize = 16;
15
16/// Size of the encryption secret in bytes (256-bit)
17pub const SECRET_SIZE: usize = 32;
18
19/// Encrypted payload containing the ciphertext and encryption parameters.
20#[derive(Debug, Clone)]
21pub struct EncryptedPayload {
22    /// The encrypted content (ciphertext + 16-byte auth tag)
23    pub payload: Vec<u8>,
24    /// The 32-byte salt used for HKDF key derivation
25    pub salt: Vec<u8>,
26    /// The 12-byte nonce used for AES-GCM encryption
27    pub nonce: Vec<u8>,
28}
29
30/// Encrypts plaintext using AES-256-GCM with HKDF-SHA256 key derivation.
31pub fn encrypt(plaintext: &[u8], secret: &[u8]) -> Result<EncryptedPayload, CodecError> {
32    // Generate random salt and nonce
33    let salt: [u8; HKDF_SALT_SIZE] = xmtp_common::rand_array();
34    let nonce: [u8; AES_GCM_NONCE_SIZE] = xmtp_common::rand_array();
35
36    // Derive AES-256 key using HKDF-SHA256
37    let key = derive_key(secret, &salt).map_err(CodecError::Encode)?;
38
39    // Create cipher
40    let cipher = Aes256Gcm::new_from_slice(&key)
41        .map_err(|e| CodecError::Encode(format!("failed to create cipher: {e}")))?;
42
43    let ciphertext = cipher
44        .encrypt((&nonce).into(), plaintext)
45        .map_err(|e| CodecError::Encode(format!("encryption failed: {e}")))?;
46
47    Ok(EncryptedPayload {
48        payload: ciphertext,
49        salt: salt.to_vec(),
50        nonce: nonce.to_vec(),
51    })
52}
53
54/// Decrypts ciphertext that was encrypted with [`encrypt`].
55pub fn decrypt(encrypted: &EncryptedPayload, secret: &[u8]) -> Result<Vec<u8>, CodecError> {
56    // Validate salt and nonce lengths
57    if encrypted.salt.len() != HKDF_SALT_SIZE {
58        return Err(CodecError::Decode(format!(
59            "invalid salt length: expected {}, got {}",
60            HKDF_SALT_SIZE,
61            encrypted.salt.len()
62        )));
63    }
64    if encrypted.nonce.len() != AES_GCM_NONCE_SIZE {
65        return Err(CodecError::Decode(format!(
66            "invalid nonce length: expected {}, got {}",
67            AES_GCM_NONCE_SIZE,
68            encrypted.nonce.len()
69        )));
70    }
71
72    // Derive AES-256 key using HKDF-SHA256
73    let key = derive_key(secret, &encrypted.salt).map_err(CodecError::Decode)?;
74
75    // Create cipher
76    let cipher = Aes256Gcm::new_from_slice(&key)
77        .map_err(|e| CodecError::Decode(format!("failed to create cipher: {e}")))?;
78
79    let nonce: &[u8; AES_GCM_NONCE_SIZE] = encrypted
80        .nonce
81        .as_slice()
82        .try_into()
83        .expect("nonce length already validated");
84
85    cipher
86        .decrypt(nonce.into(), encrypted.payload.as_slice())
87        .map_err(|e| CodecError::Decode(format!("decryption failed: {e}")))
88}
89
90/// Computes the SHA-256 hash of the given bytes.
91pub fn sha256(bytes: &[u8]) -> Vec<u8> {
92    use sha2::Digest;
93    let mut hasher = Sha256::new();
94    hasher.update(bytes);
95    hasher.finalize().to_vec()
96}
97
98/// Derives an AES-256 key from a secret and salt using HKDF-SHA256.
99fn derive_key(secret: &[u8], salt: &[u8]) -> Result<[u8; 32], String> {
100    let hkdf = Hkdf::<Sha256>::new(Some(salt), secret);
101
102    let mut key = [0u8; 32];
103    // Empty info, matching the TypeScript implementation
104    hkdf.expand(&[], &mut key)
105        .map_err(|e| format!("HKDF key derivation failed: {e}"))?;
106
107    Ok(key)
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
115    #[cfg_attr(not(target_arch = "wasm32"), test)]
116    fn test_encrypt_decrypt_roundtrip() {
117        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
118        let plaintext = b"Hello, XMTP remote attachments!";
119
120        let encrypted = encrypt(plaintext, &secret).unwrap();
121        let decrypted = decrypt(&encrypted, &secret).unwrap();
122
123        assert_eq!(plaintext.as_slice(), decrypted.as_slice());
124    }
125
126    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
127    #[cfg_attr(not(target_arch = "wasm32"), test)]
128    fn test_decrypt_wrong_secret_fails() {
129        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
130        let wrong_secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
131        let plaintext = b"Secret message";
132
133        let encrypted = encrypt(plaintext, &secret).unwrap();
134        let result = decrypt(&encrypted, &wrong_secret);
135
136        assert!(result.is_err());
137    }
138
139    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
140    #[cfg_attr(not(target_arch = "wasm32"), test)]
141    fn test_encrypt_produces_different_output_each_time() {
142        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
143        let plaintext = b"Same message";
144
145        let encrypted1 = encrypt(plaintext, &secret).unwrap();
146        let encrypted2 = encrypt(plaintext, &secret).unwrap();
147
148        // Due to random salt and nonce, encrypted output should differ
149        assert_ne!(encrypted1.payload, encrypted2.payload);
150        assert_ne!(encrypted1.salt, encrypted2.salt);
151        assert_ne!(encrypted1.nonce, encrypted2.nonce);
152    }
153
154    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
155    #[cfg_attr(not(target_arch = "wasm32"), test)]
156    fn test_encrypted_payload_sizes() {
157        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
158        let plaintext = b"Test message";
159
160        let encrypted = encrypt(plaintext, &secret).unwrap();
161
162        assert_eq!(encrypted.salt.len(), HKDF_SALT_SIZE);
163        assert_eq!(encrypted.nonce.len(), AES_GCM_NONCE_SIZE);
164        // Ciphertext is plaintext + authentication tag
165        assert_eq!(encrypted.payload.len(), plaintext.len() + AES_GCM_TAG_SIZE);
166    }
167
168    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
169    #[cfg_attr(not(target_arch = "wasm32"), test)]
170    fn test_invalid_salt_length() {
171        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
172        let encrypted = EncryptedPayload {
173            payload: vec![0u8; 32],
174            salt: vec![0u8; 16], // Wrong size
175            nonce: vec![0u8; AES_GCM_NONCE_SIZE],
176        };
177
178        let result = decrypt(&encrypted, &secret);
179        assert!(result.is_err());
180        assert!(
181            result
182                .unwrap_err()
183                .to_string()
184                .contains("invalid salt length")
185        );
186    }
187
188    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
189    #[cfg_attr(not(target_arch = "wasm32"), test)]
190    fn test_invalid_nonce_length() {
191        let secret: [u8; SECRET_SIZE] = xmtp_common::rand_array();
192        let encrypted = EncryptedPayload {
193            payload: vec![0u8; 32],
194            salt: vec![0u8; HKDF_SALT_SIZE],
195            nonce: vec![0u8; 8], // Wrong size
196        };
197
198        let result = decrypt(&encrypted, &secret);
199        assert!(result.is_err());
200        assert!(
201            result
202                .unwrap_err()
203                .to_string()
204                .contains("invalid nonce length")
205        );
206    }
207}
xmtp_content_types/encryption.rs

xmtp_content_types/
encryption.rs
