Skip to main content

weavatrix_memory/codec/
encryption.rs

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