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
15pub struct EncryptionKey<'a> {
17 pub id: &'a str,
18 pub bytes: &'a [u8; 32],
19}
20
21pub trait EncryptionKeys {
23 fn active_key(&self) -> Result<EncryptionKey<'_>>;
29
30 fn decryption_key(&self, id: &str) -> Result<&[u8; 32]>;
36}
37
38pub struct StaticKey {
40 id: String,
41 bytes: Zeroizing<[u8; 32]>,
42}
43
44impl StaticKey {
45 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
77pub trait NonceSource {
79 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
96pub 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 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 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}