Skip to main content

weavatrix_memory/codec/
encryption.rs

1use super::Codec;
2use crate::error::{MemoryError, Result};
3use chacha20poly1305::{
4    XChaCha20Poly1305, XNonce,
5    aead::{Aead, AeadInOut, KeyInit, Payload},
6};
7use zeroize::Zeroizing;
8
9const HEADER: &[u8; 8] = b"WMEMXE01";
10const NONCE_LEN: usize = 24;
11const TAG_LEN: usize = 16;
12const MAX_KEY_ID_BYTES: usize = 64;
13const FIXED_HEADER_LEN: usize = 9 + NONCE_LEN;
14
15/// Borrowed key material returned by an [`EncryptionKeys`] provider.
16pub struct EncryptionKey<'a> {
17    pub id: &'a str,
18    pub bytes: &'a [u8; 32],
19}
20
21/// Supplies the active encryption key and historical decryption keys.
22pub trait EncryptionKeys {
23    /// Returns the key used for new envelopes.
24    ///
25    /// # Errors
26    ///
27    /// Returns an error when the active key is unavailable.
28    fn active_key(&self) -> Result<EncryptionKey<'_>>;
29
30    /// Returns key material for an identifier stored in an envelope.
31    ///
32    /// # Errors
33    ///
34    /// Returns an error when the requested key is unavailable.
35    fn decryption_key(&self, id: &str) -> Result<&[u8; 32]>;
36}
37
38/// One zeroizing key suitable for applications without key rotation.
39pub struct StaticKey {
40    id: String,
41    bytes: Zeroizing<[u8; 32]>,
42}
43
44impl StaticKey {
45    /// Creates a key with a stable, envelope-visible identifier.
46    ///
47    /// # Errors
48    ///
49    /// Rejects empty, non-ASCII, or overlong identifiers.
50    pub fn new(id: impl Into<String>, bytes: [u8; 32]) -> Result<Self> {
51        let id = id.into();
52        validate_key_id(&id)?;
53        Ok(Self {
54            id,
55            bytes: Zeroizing::new(bytes),
56        })
57    }
58}
59
60impl EncryptionKeys for StaticKey {
61    fn active_key(&self) -> Result<EncryptionKey<'_>> {
62        Ok(EncryptionKey {
63            id: &self.id,
64            bytes: &self.bytes,
65        })
66    }
67
68    fn decryption_key(&self, id: &str) -> Result<&[u8; 32]> {
69        if id == self.id {
70            Ok(&self.bytes)
71        } else {
72            Err(codec("encryption key is unavailable"))
73        }
74    }
75}
76
77/// Produces a unique `XChaCha20` nonce for each encoded value.
78pub trait NonceSource {
79    /// Fills one 192-bit nonce.
80    ///
81    /// # Errors
82    ///
83    /// Returns an error when secure randomness is unavailable.
84    fn fill(&self, nonce: &mut [u8; NONCE_LEN]) -> Result<()>;
85}
86
87#[derive(Debug, Clone, Copy, Default)]
88pub struct OsNonce;
89
90impl NonceSource for OsNonce {
91    fn fill(&self, nonce: &mut [u8; NONCE_LEN]) -> Result<()> {
92        getrandom::fill(nonce).map_err(|_| codec("operating-system randomness is unavailable"))
93    }
94}
95
96/// Authenticated XChaCha20-Poly1305 wrapper for any existing codec.
97pub struct XChaCha20Codec<C, K, N = OsNonce> {
98    inner: C,
99    keys: K,
100    nonce_source: N,
101    context: Vec<u8>,
102    max_plaintext_bytes: usize,
103}
104
105impl<C, K> XChaCha20Codec<C, K, OsNonce> {
106    /// Uses OS randomness for nonces and binds ciphertext to `context`.
107    ///
108    /// # Errors
109    ///
110    /// Rejects an empty context or zero plaintext limit.
111    pub fn new(
112        inner: C,
113        keys: K,
114        context: impl Into<Vec<u8>>,
115        max_plaintext_bytes: usize,
116    ) -> Result<Self> {
117        Self::with_nonce_source(inner, keys, context, max_plaintext_bytes, OsNonce)
118    }
119}
120
121impl<C, K, N> XChaCha20Codec<C, K, N> {
122    /// Injects a nonce source, primarily for deterministic testing.
123    ///
124    /// # Errors
125    ///
126    /// Rejects an empty context or zero plaintext limit.
127    pub fn with_nonce_source(
128        inner: C,
129        keys: K,
130        context: impl Into<Vec<u8>>,
131        max_plaintext_bytes: usize,
132        nonce_source: N,
133    ) -> Result<Self> {
134        let context = context.into();
135        if context.is_empty() {
136            return Err(invalid("encryption.context", "must not be empty"));
137        }
138        if max_plaintext_bytes == 0 {
139            return Err(invalid("max_plaintext_bytes", "must be greater than zero"));
140        }
141        Ok(Self {
142            inner,
143            keys,
144            nonce_source,
145            context,
146            max_plaintext_bytes,
147        })
148    }
149}
150
151impl<T, C, K, N> Codec<T> for XChaCha20Codec<C, K, N>
152where
153    C: Codec<T>,
154    K: EncryptionKeys,
155    N: NonceSource,
156{
157    fn encode(&self, value: &T) -> Result<Vec<u8>> {
158        let plaintext = Zeroizing::new(self.inner.encode(value)?);
159        if plaintext.len() > self.max_plaintext_bytes {
160            return Err(invalid(
161                "plaintext",
162                "encoded value exceeds max_plaintext_bytes",
163            ));
164        }
165        let key = self.keys.active_key()?;
166        validate_key_id(key.id)?;
167        let mut nonce = [0_u8; NONCE_LEN];
168        self.nonce_source.fill(&mut nonce)?;
169        let header = envelope_header(key.id, &nonce)?;
170        let mut aad = header.clone();
171        aad.extend_from_slice(&self.context);
172        let cipher = XChaCha20Poly1305::new_from_slice(key.bytes)
173            .map_err(|_| codec("invalid encryption key"))?;
174        let nonce = XNonce::from(nonce);
175        let mut output = Vec::with_capacity(header.len() + plaintext.len() + TAG_LEN);
176        output.extend_from_slice(&header);
177        output.extend_from_slice(&plaintext);
178        let tag = cipher
179            .encrypt_inout_detached(&nonce, &aad, (&mut output[header.len()..]).into())
180            .map_err(|_| codec("encryption failed"))?;
181        output.extend_from_slice(&tag);
182        Ok(output)
183    }
184
185    fn decode(&self, bytes: &[u8]) -> Result<T> {
186        let envelope = parse_envelope(bytes)?;
187        if envelope.ciphertext.len() < TAG_LEN
188            || envelope.ciphertext.len() - TAG_LEN > self.max_plaintext_bytes
189        {
190            return Err(codec("ciphertext exceeds configured size limit"));
191        }
192        let mut aad = envelope.header.to_vec();
193        aad.extend_from_slice(&self.context);
194        let cipher = XChaCha20Poly1305::new_from_slice(self.keys.decryption_key(envelope.key_id)?)
195            .map_err(|_| codec("invalid encryption key"))?;
196        let nonce = XNonce::from(*envelope.nonce);
197        let plaintext = Zeroizing::new(
198            cipher
199                .decrypt(
200                    &nonce,
201                    Payload {
202                        msg: envelope.ciphertext,
203                        aad: &aad,
204                    },
205                )
206                .map_err(|_| codec("authentication failed"))?,
207        );
208        self.inner.decode(&plaintext)
209    }
210}
211
212fn envelope_header(key_id: &str, nonce: &[u8; NONCE_LEN]) -> Result<Vec<u8>> {
213    let key_len = u8::try_from(key_id.len()).map_err(|_| codec("key identifier is too long"))?;
214    let mut header = Vec::with_capacity(FIXED_HEADER_LEN + key_id.len());
215    header.extend_from_slice(HEADER);
216    header.push(key_len);
217    header.extend_from_slice(nonce);
218    header.extend_from_slice(key_id.as_bytes());
219    Ok(header)
220}
221
222struct Envelope<'a> {
223    header: &'a [u8],
224    key_id: &'a str,
225    nonce: &'a [u8; NONCE_LEN],
226    ciphertext: &'a [u8],
227}
228
229fn parse_envelope(bytes: &[u8]) -> Result<Envelope<'_>> {
230    if bytes.len() < FIXED_HEADER_LEN || &bytes[..8] != HEADER {
231        return Err(codec("unsupported encrypted envelope"));
232    }
233    let header_len = FIXED_HEADER_LEN
234        .checked_add(usize::from(bytes[8]))
235        .ok_or(MemoryError::CapacityOverflow)?;
236    if bytes.len() < header_len + TAG_LEN {
237        return Err(codec("truncated encrypted envelope"));
238    }
239    let nonce = bytes[9..FIXED_HEADER_LEN]
240        .try_into()
241        .map_err(|_| codec("invalid encryption nonce"))?;
242    let key_id = core::str::from_utf8(&bytes[FIXED_HEADER_LEN..header_len])
243        .map_err(|_| codec("invalid key identifier"))?;
244    validate_key_id(key_id)?;
245    Ok(Envelope {
246        header: &bytes[..header_len],
247        key_id,
248        nonce,
249        ciphertext: &bytes[header_len..],
250    })
251}
252
253fn validate_key_id(id: &str) -> Result<()> {
254    if id.is_empty() || id.len() > MAX_KEY_ID_BYTES || !id.is_ascii() {
255        return Err(invalid("encryption.key_id", "must be 1..=64 ASCII bytes"));
256    }
257    Ok(())
258}
259
260fn codec(message: &str) -> MemoryError {
261    MemoryError::Codec {
262        message: message.to_owned(),
263    }
264}
265
266fn invalid(field: &'static str, reason: &'static str) -> MemoryError {
267    MemoryError::InvalidValue { field, reason }
268}