celers_protocol/
crypto.rs

1//! Message encryption and decryption
2//!
3//! This module provides AES-256-GCM encryption for Celery protocol messages.
4//! It ensures message confidentiality by encrypting message bodies.
5//!
6//! # Example
7//!
8//! ```
9//! use celers_protocol::crypto::{MessageEncryptor, EncryptionError};
10//!
11//! # #[cfg(feature = "encryption")]
12//! # {
13//! let key = b"32-byte-secret-key-for-aes-256!!";
14//! let encryptor = MessageEncryptor::new(key).unwrap();
15//!
16//! let plaintext = b"secret task data";
17//! let (ciphertext, nonce) = encryptor.encrypt(plaintext).unwrap();
18//!
19//! // Decrypt
20//! let decrypted = encryptor.decrypt(&ciphertext, &nonce).unwrap();
21//! assert_eq!(decrypted, plaintext);
22//! # }
23//! ```
24
25#[cfg(feature = "encryption")]
26use aes_gcm::{
27    aead::{Aead, AeadCore, KeyInit, OsRng},
28    Aes256Gcm, Nonce,
29};
30
31use std::fmt;
32
33/// Error type for encryption operations
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub enum EncryptionError {
36    /// Invalid key length (must be 32 bytes for AES-256)
37    InvalidKeyLength,
38    /// Encryption failed
39    EncryptionFailed,
40    /// Decryption failed
41    DecryptionFailed,
42    /// Invalid nonce length (must be 12 bytes for GCM)
43    InvalidNonceLength,
44}
45
46impl fmt::Display for EncryptionError {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        match self {
49            EncryptionError::InvalidKeyLength => {
50                write!(f, "Invalid key length (must be 32 bytes for AES-256)")
51            }
52            EncryptionError::EncryptionFailed => write!(f, "Encryption failed"),
53            EncryptionError::DecryptionFailed => write!(f, "Decryption failed"),
54            EncryptionError::InvalidNonceLength => {
55                write!(f, "Invalid nonce length (must be 12 bytes)")
56            }
57        }
58    }
59}
60
61impl std::error::Error for EncryptionError {}
62
63/// Nonce size for AES-GCM (96 bits = 12 bytes)
64#[cfg(feature = "encryption")]
65pub const NONCE_SIZE: usize = 12;
66
67/// Key size for AES-256 (256 bits = 32 bytes)
68pub const KEY_SIZE: usize = 32;
69
70/// Message encryptor using AES-256-GCM
71///
72/// Provides authenticated encryption of messages using AES-256-GCM.
73/// GCM (Galois/Counter Mode) provides both confidentiality and authenticity.
74#[cfg(feature = "encryption")]
75pub struct MessageEncryptor {
76    cipher: Aes256Gcm,
77}
78
79#[cfg(feature = "encryption")]
80impl MessageEncryptor {
81    /// Create a new message encryptor with the given key
82    ///
83    /// # Arguments
84    ///
85    /// * `key` - 32-byte secret key for AES-256
86    ///
87    /// # Returns
88    ///
89    /// `Ok(MessageEncryptor)` if the key is valid, `Err(EncryptionError)` otherwise
90    pub fn new(key: &[u8]) -> Result<Self, EncryptionError> {
91        if key.len() != KEY_SIZE {
92            return Err(EncryptionError::InvalidKeyLength);
93        }
94
95        let cipher =
96            Aes256Gcm::new_from_slice(key).map_err(|_| EncryptionError::InvalidKeyLength)?;
97
98        Ok(Self { cipher })
99    }
100
101    /// Encrypt a message
102    ///
103    /// # Arguments
104    ///
105    /// * `plaintext` - The message bytes to encrypt
106    ///
107    /// # Returns
108    ///
109    /// A tuple of (ciphertext, nonce) on success, or `EncryptionError` on failure
110    pub fn encrypt(&self, plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>), EncryptionError> {
111        // Generate a random nonce
112        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
113
114        // Encrypt
115        let ciphertext = self
116            .cipher
117            .encrypt(&nonce, plaintext)
118            .map_err(|_| EncryptionError::EncryptionFailed)?;
119
120        Ok((ciphertext, nonce.to_vec()))
121    }
122
123    /// Decrypt a message
124    ///
125    /// # Arguments
126    ///
127    /// * `ciphertext` - The encrypted message bytes
128    /// * `nonce` - The nonce used during encryption (12 bytes)
129    ///
130    /// # Returns
131    ///
132    /// The decrypted plaintext on success, or `EncryptionError` on failure
133    pub fn decrypt(&self, ciphertext: &[u8], nonce: &[u8]) -> Result<Vec<u8>, EncryptionError> {
134        if nonce.len() != NONCE_SIZE {
135            return Err(EncryptionError::InvalidNonceLength);
136        }
137
138        let nonce = Nonce::from_slice(nonce);
139
140        let plaintext = self
141            .cipher
142            .decrypt(nonce, ciphertext)
143            .map_err(|_| EncryptionError::DecryptionFailed)?;
144
145        Ok(plaintext)
146    }
147
148    /// Encrypt a message and return hex-encoded ciphertext and nonce
149    ///
150    /// # Arguments
151    ///
152    /// * `plaintext` - The message bytes to encrypt
153    ///
154    /// # Returns
155    ///
156    /// A tuple of (ciphertext_hex, nonce_hex) on success
157    pub fn encrypt_hex(&self, plaintext: &[u8]) -> Result<(String, String), EncryptionError> {
158        let (ciphertext, nonce) = self.encrypt(plaintext)?;
159        Ok((hex::encode(ciphertext), hex::encode(nonce)))
160    }
161
162    /// Decrypt a hex-encoded message
163    ///
164    /// # Arguments
165    ///
166    /// * `ciphertext_hex` - The hex-encoded ciphertext
167    /// * `nonce_hex` - The hex-encoded nonce
168    ///
169    /// # Returns
170    ///
171    /// The decrypted plaintext on success, or `EncryptionError` on failure
172    pub fn decrypt_hex(
173        &self,
174        ciphertext_hex: &str,
175        nonce_hex: &str,
176    ) -> Result<Vec<u8>, EncryptionError> {
177        let ciphertext =
178            hex::decode(ciphertext_hex).map_err(|_| EncryptionError::DecryptionFailed)?;
179        let nonce = hex::decode(nonce_hex).map_err(|_| EncryptionError::InvalidNonceLength)?;
180
181        self.decrypt(&ciphertext, &nonce)
182    }
183}
184
185// Placeholder implementation when encryption feature is disabled
186#[cfg(not(feature = "encryption"))]
187pub struct MessageEncryptor;
188
189#[cfg(not(feature = "encryption"))]
190impl MessageEncryptor {
191    pub fn new(_key: &[u8]) -> Result<Self, EncryptionError> {
192        Err(EncryptionError::EncryptionFailed)
193    }
194}
195
196#[cfg(all(test, feature = "encryption"))]
197mod tests {
198    use super::*;
199
200    #[test]
201    fn test_encrypt_decrypt() {
202        let key = b"32-byte-secret-key-for-aes-256!!";
203        let encryptor = MessageEncryptor::new(key).unwrap();
204
205        let plaintext = b"secret message";
206        let (ciphertext, nonce) = encryptor.encrypt(plaintext).unwrap();
207
208        let decrypted = encryptor.decrypt(&ciphertext, &nonce).unwrap();
209        assert_eq!(decrypted, plaintext);
210    }
211
212    #[test]
213    fn test_invalid_key_length() {
214        let short_key = b"too-short";
215        assert!(MessageEncryptor::new(short_key).is_err());
216
217        let long_key = b"this-key-is-way-too-long-for-aes-256-encryption";
218        assert!(MessageEncryptor::new(long_key).is_err());
219    }
220
221    #[test]
222    fn test_decrypt_wrong_nonce() {
223        let key = b"32-byte-secret-key-for-aes-256!!";
224        let encryptor = MessageEncryptor::new(key).unwrap();
225
226        let plaintext = b"secret message";
227        let (ciphertext, _) = encryptor.encrypt(plaintext).unwrap();
228
229        let wrong_nonce = vec![0u8; NONCE_SIZE];
230        assert!(encryptor.decrypt(&ciphertext, &wrong_nonce).is_err());
231    }
232
233    #[test]
234    fn test_decrypt_tampered_ciphertext() {
235        let key = b"32-byte-secret-key-for-aes-256!!";
236        let encryptor = MessageEncryptor::new(key).unwrap();
237
238        let plaintext = b"secret message";
239        let (mut ciphertext, nonce) = encryptor.encrypt(plaintext).unwrap();
240
241        // Tamper with ciphertext
242        if !ciphertext.is_empty() {
243            ciphertext[0] ^= 1;
244        }
245
246        assert!(encryptor.decrypt(&ciphertext, &nonce).is_err());
247    }
248
249    #[test]
250    fn test_invalid_nonce_length() {
251        let key = b"32-byte-secret-key-for-aes-256!!";
252        let encryptor = MessageEncryptor::new(key).unwrap();
253
254        let plaintext = b"secret message";
255        let (ciphertext, _) = encryptor.encrypt(plaintext).unwrap();
256
257        let wrong_nonce = vec![0u8; 8]; // Wrong size
258        assert_eq!(
259            encryptor.decrypt(&ciphertext, &wrong_nonce),
260            Err(EncryptionError::InvalidNonceLength)
261        );
262    }
263
264    #[test]
265    fn test_encrypt_hex() {
266        let key = b"32-byte-secret-key-for-aes-256!!";
267        let encryptor = MessageEncryptor::new(key).unwrap();
268
269        let plaintext = b"secret message";
270        let (ciphertext_hex, nonce_hex) = encryptor.encrypt_hex(plaintext).unwrap();
271
272        // Verify hex encoding
273        assert!(hex::decode(&ciphertext_hex).is_ok());
274        assert!(hex::decode(&nonce_hex).is_ok());
275
276        // Decrypt using hex
277        let decrypted = encryptor.decrypt_hex(&ciphertext_hex, &nonce_hex).unwrap();
278        assert_eq!(decrypted, plaintext);
279    }
280
281    #[test]
282    fn test_different_nonces() {
283        let key = b"32-byte-secret-key-for-aes-256!!";
284        let encryptor = MessageEncryptor::new(key).unwrap();
285
286        let plaintext = b"secret message";
287        let (_, nonce1) = encryptor.encrypt(plaintext).unwrap();
288        let (_, nonce2) = encryptor.encrypt(plaintext).unwrap();
289
290        // Each encryption should use a different nonce
291        assert_ne!(nonce1, nonce2);
292    }
293
294    #[test]
295    fn test_encryption_error_display() {
296        assert_eq!(
297            EncryptionError::InvalidKeyLength.to_string(),
298            "Invalid key length (must be 32 bytes for AES-256)"
299        );
300        assert_eq!(
301            EncryptionError::EncryptionFailed.to_string(),
302            "Encryption failed"
303        );
304        assert_eq!(
305            EncryptionError::DecryptionFailed.to_string(),
306            "Decryption failed"
307        );
308        assert_eq!(
309            EncryptionError::InvalidNonceLength.to_string(),
310            "Invalid nonce length (must be 12 bytes)"
311        );
312    }
313
314    #[test]
315    fn test_empty_message() {
316        let key = b"32-byte-secret-key-for-aes-256!!";
317        let encryptor = MessageEncryptor::new(key).unwrap();
318
319        let plaintext = b"";
320        let (ciphertext, nonce) = encryptor.encrypt(plaintext).unwrap();
321
322        let decrypted = encryptor.decrypt(&ciphertext, &nonce).unwrap();
323        assert_eq!(decrypted, plaintext);
324    }
325
326    #[test]
327    fn test_large_message() {
328        let key = b"32-byte-secret-key-for-aes-256!!";
329        let encryptor = MessageEncryptor::new(key).unwrap();
330
331        let plaintext = vec![42u8; 10000];
332        let (ciphertext, nonce) = encryptor.encrypt(&plaintext).unwrap();
333
334        let decrypted = encryptor.decrypt(&ciphertext, &nonce).unwrap();
335        assert_eq!(decrypted, plaintext);
336    }
337}