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
14pub struct EncryptionKey<'a> {
16 pub id: &'a str,
17 pub bytes: &'a [u8; 32],
18}
19
20pub trait EncryptionKeys {
22 fn active_key(&self) -> Result<EncryptionKey<'_>>;
28
29 fn decryption_key(&self, id: &str) -> Result<&[u8; 32]>;
35}
36
37pub struct StaticKey {
39 id: String,
40 bytes: Zeroizing<[u8; 32]>,
41}
42
43impl StaticKey {
44 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
76pub trait NonceSource {
78 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
95pub 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 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 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}