Skip to main content

streaming_crypto/core_api/crypto/
aead.rs

1// ## đź“‚ File: `src/crypto/aead.rs`
2
3//! src/crypto/aead.rs
4//! AEAD interface for AES-256-GCM and ChaCha20-Poly1305.
5//!
6//! Design notes:
7//! - Both ciphers use 32-byte keys and 12-byte nonces.
8//! - Tag verification is constant-time and must fail closed (no partial plaintext).
9//! - Caller provides nonce and AAD (built by aad module) per frame.
10//! - Cipher selection is driven by header.cipher (u16 registry).
11
12use crate::constants::{MASTER_KEY_LENGTHS, cipher_ids};
13use crate::headers::types::{HeaderV1};
14use crate::crypto::types::{KEY_LEN_32, NONCE_LEN_12, TAG_LEN, CryptoError};
15
16// Import AEAD traits from aes_gcm's re-export to avoid unresolved `aead` path and duplicates.
17use aes_gcm::aead::{Aead, Buffer, KeyInit, Payload};
18use aes_gcm::aead::AeadInOut;
19use hybrid_array::{Array, sizes::U12};
20
21// Concrete AEAD types
22use aes_gcm::{Aes256Gcm, Nonce as AesNonce};                 // 32-byte key, 12-byte nonce
23use chacha20poly1305::{ChaCha20Poly1305, Nonce as ChaNonce}; // 32-byte key, 12-byte nonce
24use tracing::debug; 
25
26/// Unified AEAD cipher implementation selected by header.cipher.
27#[derive(Clone)]
28pub enum AeadImpl {
29    AesGcm(Aes256Gcm),
30    ChaCha(ChaCha20Poly1305),
31}
32
33impl AeadImpl {
34    /// Construct AEAD implementation from header.cipher and derived session key.
35    pub fn from_header_and_key(header: &HeaderV1, session_key: &[u8]) -> Result<Self, CryptoError> {
36        if session_key.len() != KEY_LEN_32 {
37            return Err(CryptoError::InvalidKeyLen {
38                expected: &MASTER_KEY_LENGTHS,
39                actual: session_key.len(),
40            });
41        }
42
43        match header.cipher {
44            x if x == cipher_ids::AES256_GCM => {
45                let cipher = Aes256Gcm::new_from_slice(session_key)
46                    .map_err(|_| CryptoError::InvalidKeyLen {
47                        expected: &MASTER_KEY_LENGTHS,
48                        actual: session_key.len(),
49                    })?;
50                Ok(Self::AesGcm(cipher))
51            }
52            x if x == cipher_ids::CHACHA20_POLY1305 => {
53                let cipher = ChaCha20Poly1305::new_from_slice(session_key)
54                    .map_err(|_| CryptoError::InvalidKeyLen {
55                        expected: &MASTER_KEY_LENGTHS,
56                        actual: session_key.len(),
57                    })?;
58                Ok(Self::ChaCha(cipher))
59            }
60
61            other => Err(CryptoError::UnsupportedCipher { cipher_id: other }),
62        }
63    }
64        
65
66    fn extract_nonce(
67        &self,
68        nonce_12: &[u8],
69    ) -> Result<Array<u8, U12>, CryptoError> {
70        
71        if nonce_12.len() != NONCE_LEN_12 {
72            return Err(CryptoError::InvalidNonceLen {
73                expected: NONCE_LEN_12,
74                actual: nonce_12.len(),
75            });
76        }
77
78        match self {
79            AeadImpl::AesGcm(_) => {
80                let nonce = AesNonce::try_from(nonce_12).map_err(|e|CryptoError::Failure(e.to_string()))?;
81                Ok(nonce)
82            }
83            AeadImpl::ChaCha(_) => {
84                let nonce = ChaNonce::try_from(nonce_12).map_err(|e|CryptoError::Failure(e.to_string()))?;
85                Ok(nonce)
86            }
87        }
88    }
89
90    /// AEAD seal (encrypt) plaintext with nonce and AAD.
91    pub fn seal(
92        &self,
93        nonce_12: &[u8],
94        aad: &[u8],
95        plaintext: &[u8],
96    ) -> Result<Vec<u8>, CryptoError> {
97        if plaintext.is_empty() {
98            return Err(CryptoError::Failure("plaintext must not be empty".into()));
99        }
100        let nonce = self.extract_nonce(nonce_12)?;
101
102        // Debug information
103        debug!(
104            "[AEAD::seal] cipher={:?}, plaintext_len={}, aad_len={}, nonce={:02x?}",
105            match self {
106                AeadImpl::AesGcm(_) => "AES-GCM",
107                AeadImpl::ChaCha(_) => "ChaCha20-Poly1305",
108            },
109            plaintext.len(),
110            aad.len(),
111            nonce_12
112        );
113
114        match self {
115            AeadImpl::AesGcm(cipher) => {
116                debug!(
117                    "[AEAD::seal] AES-GCM sealing frame with {} bytes payload",
118                    plaintext.len()
119                );
120                cipher
121                    .encrypt(&nonce, Payload { msg: plaintext, aad })
122                    .map_err(|e| CryptoError::Failure(format!("AES-GCM seal failed: {e}")))
123            }
124            AeadImpl::ChaCha(cipher) => {
125                debug!(
126                    "[AEAD::seal] ChaCha20-Poly1305 sealing frame with {} bytes payload",
127                    plaintext.len()
128                );
129                cipher
130                    .encrypt(&nonce, Payload { msg: plaintext, aad })
131                    .map_err(|e| CryptoError::Failure(format!("ChaCha20-Poly1305 seal failed: {e}")))
132            }
133        }
134    }
135
136    /// AEAD open (decrypt) ciphertext with nonce and AAD.
137    pub fn open(
138        &self,
139        nonce_12: &[u8],
140        aad: &[u8],
141        ciphertext_and_tag: &[u8],
142    ) -> Result<Vec<u8>, CryptoError> {
143        if ciphertext_and_tag.len() < TAG_LEN {
144            return Err(CryptoError::Failure("ciphertext too short".into()));
145        }
146        let nonce = self.extract_nonce(nonce_12)?;
147
148        match self {
149            AeadImpl::AesGcm(cipher) => {
150                cipher
151                    .decrypt(&nonce, Payload { msg: ciphertext_and_tag, aad })
152                    .map_err(|e| CryptoError::Failure(format!("AES-GCM open failed: {e}")))
153            }
154            AeadImpl::ChaCha(cipher) => {
155                cipher
156                    .decrypt(&nonce, Payload { msg: ciphertext_and_tag, aad })
157                    .map_err(|e| CryptoError::Failure(format!("ChaCha20-Poly1305 open failed: {e}")))
158            }
159        }
160
161    }
162
163    // ### In‑Place AEAD Implementation
164
165    /// AEAD seal (encrypt) plaintext in place with nonce and AAD.
166    /// The buffer must have enough capacity for plaintext + tag.
167    pub fn seal_in_place(
168        &self,
169        nonce_12: &[u8],
170        aad: &[u8],
171        buf: &mut (dyn Buffer + 'static), // contains plaintext, will be replaced with ciphertext+tag
172    ) -> Result<(), CryptoError> {
173        if buf.is_empty() {
174            return Err(CryptoError::Failure("plaintext must not be empty".into()));
175        }
176
177        let nonce = self.extract_nonce(nonce_12)?;
178
179        match self {
180            AeadImpl::AesGcm(cipher) => cipher
181                .encrypt_in_place(&nonce, aad, buf)
182                .map_err(|e| CryptoError::Failure(format!("AES-GCM seal_in_place failed: {e}"))),
183
184            AeadImpl::ChaCha(cipher) => cipher
185                .encrypt_in_place(&nonce, aad, buf)
186                .map_err(|e| CryptoError::Failure(format!("ChaCha20-Poly1305 seal_in_place failed: {e}"))),
187        }
188    }
189
190    /// AEAD open (decrypt) ciphertext in place with nonce and AAD.
191    /// The buffer must contain ciphertext+tag, will be replaced with plaintext.
192    pub fn open_in_place(
193        &self,
194        nonce_12: &[u8],
195        aad: &[u8],
196        buf: &mut (dyn Buffer + 'static), // contains ciphertext+tag, will be replaced with plaintext
197    ) -> Result<(), CryptoError> {
198        if buf.len() < TAG_LEN {
199            return Err(CryptoError::Failure("ciphertext too short".into()));
200        }
201
202        let nonce = self.extract_nonce(nonce_12)?;
203        
204        match self {
205            AeadImpl::AesGcm(cipher) => cipher
206                .decrypt_in_place(&nonce, aad, buf)
207                .map_err(|e| CryptoError::Failure(format!("AES-GCM open_in_place failed: {e}"))),
208
209            AeadImpl::ChaCha(cipher) => cipher
210                .decrypt_in_place(&nonce, aad, buf)
211                .map_err(|e| CryptoError::Failure(format!("ChaCha20-Poly1305 open_in_place failed: {e}"))),
212        }
213    }
214
215    // ### Key Differences
216    // - **`encrypt_in_place` / `decrypt_in_place`**: These are the in‑place APIs provided by `aes-gcm`, `chacha20poly1305`, and similar crates. They mutate the buffer directly.
217    // - **Buffer ownership**: Instead of returning a new `Vec<u8>`, the caller provides a mutable buffer (`Vec<u8>` or `BytesMut`) that already contains plaintext (for seal) or ciphertext+tag (for open).
218    // - **No extra copies**: Encryption/decryption happens inside the buffer, then we can freeze into `Bytes` if needed.
219    // ### Usage Example
220    // ```rust
221    // let mut buf = BytesMut::from(&plaintext[..]);
222    // buf.reserve(TAG_LEN); // ensure space for tag
223
224    // aead.seal_in_place(&nonce, &aad, &mut buf)?;
225    // let ciphertext_and_tag = buf.freeze();
226
227    // let mut ct_buf = ciphertext_and_tag.to_vec();
228    // aead.open_in_place(&nonce, &aad, &mut ct_buf)?;
229    // let plaintext = Bytes::from(ct_buf);
230    // ```
231}
232    // This pattern ensures **true zero‑copy** across our encrypt/decrypt pipeline: one buffer, mutated in place, no redundant allocations.  
233