1use std::collections::HashMap;
45use std::fmt;
46
47use ring::aead::{Aad, CHACHA20_POLY1305, LessSafeKey, NONCE_LEN, Nonce, UnboundKey};
48use ring::rand::{SecureRandom, SystemRandom};
49use zeroize::Zeroizing;
50
51pub const KEK_LEN: usize = 32;
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum CryptoError {
61 InvalidKek,
63 UnknownKekId,
65 SealFailed,
67 OpenFailed,
70}
71
72impl fmt::Display for CryptoError {
76 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77 match self {
78 Self::InvalidKek => {
79 write!(f, "invalid key-encryption key: must be {KEK_LEN} raw bytes")
80 }
81 Self::UnknownKekId => {
82 write!(f, "no key-encryption key configured for this ciphertext")
83 }
84 Self::SealFailed => write!(f, "could not seal secret"),
85 Self::OpenFailed => write!(f, "could not open secret"),
86 }
87 }
88}
89
90impl std::error::Error for CryptoError {}
91
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub struct SecretBinding<'a> {
98 pub namespace: &'a str,
100 pub provider: &'a str,
102}
103
104impl SecretBinding<'_> {
105 fn to_aad_bytes(self) -> Vec<u8> {
112 let ns = self.namespace.as_bytes();
113 let pr = self.provider.as_bytes();
114 let mut out = Vec::with_capacity(16 + ns.len() + pr.len());
115 out.extend_from_slice(&(ns.len() as u64).to_be_bytes());
120 out.extend_from_slice(ns);
121 out.extend_from_slice(&(pr.len() as u64).to_be_bytes());
122 out.extend_from_slice(pr);
123 out
124 }
125}
126
127#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
138pub struct SealedSecret {
139 pub kek_id: String,
141 #[serde(with = "hex_array")]
143 pub nonce: [u8; NONCE_LEN],
144 #[serde(with = "hex_vec")]
146 pub ciphertext: Vec<u8>,
147}
148
149mod hex_array {
151 use super::NONCE_LEN;
152 use serde::{Deserialize, Deserializer, Serializer};
153
154 pub(super) fn serialize<S: Serializer>(v: &[u8; NONCE_LEN], s: S) -> Result<S::Ok, S::Error> {
155 s.serialize_str(&super::to_hex(v))
156 }
157
158 pub(super) fn deserialize<'de, D: Deserializer<'de>>(
159 d: D,
160 ) -> Result<[u8; NONCE_LEN], D::Error> {
161 let raw = String::deserialize(d)?;
162 let bytes = super::decode_hex(&raw).map_err(serde::de::Error::custom)?;
163 <[u8; NONCE_LEN]>::try_from(bytes.as_slice())
164 .map_err(|_| serde::de::Error::custom(format!("nonce must be {NONCE_LEN} bytes")))
165 }
166}
167
168mod hex_vec {
170 use serde::{Deserialize, Deserializer, Serializer};
171
172 pub(super) fn serialize<S: Serializer>(v: &[u8], s: S) -> Result<S::Ok, S::Error> {
173 s.serialize_str(&super::to_hex(v))
174 }
175
176 pub(super) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
177 let raw = String::deserialize(d)?;
178 super::decode_hex(&raw)
179 .map(|z| z.to_vec())
180 .map_err(serde::de::Error::custom)
181 }
182}
183
184fn to_hex(bytes: &[u8]) -> String {
186 use std::fmt::Write as _;
187 bytes
188 .iter()
189 .fold(String::with_capacity(bytes.len() * 2), |mut acc, b| {
190 let _ = write!(acc, "{b:02x}");
191 acc
192 })
193}
194
195impl fmt::Debug for SealedSecret {
198 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
199 f.debug_struct("SealedSecret")
200 .field("kek_id", &self.kek_id)
201 .field("ciphertext_len", &self.ciphertext.len())
202 .finish_non_exhaustive()
203 }
204}
205
206pub struct KeyRing {
212 keys: HashMap<String, LessSafeKey>,
213 current: String,
214 rng: SystemRandom,
215}
216
217impl fmt::Debug for KeyRing {
220 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
221 let mut ids: Vec<&str> = self.keys.keys().map(String::as_str).collect();
222 ids.sort_unstable();
223 f.debug_struct("KeyRing")
226 .field("kek_ids", &ids)
227 .field("current", &self.current)
228 .finish_non_exhaustive()
229 }
230}
231
232impl KeyRing {
233 pub fn new<I>(keys: I, current: impl Into<String>) -> Result<Self, CryptoError>
241 where
242 I: IntoIterator<Item = (String, Zeroizing<Vec<u8>>)>,
243 {
244 let mut map = HashMap::new();
245 for (id, raw) in keys {
246 if raw.len() != KEK_LEN {
247 return Err(CryptoError::InvalidKek);
248 }
249 let unbound =
250 UnboundKey::new(&CHACHA20_POLY1305, &raw).map_err(|_| CryptoError::InvalidKek)?;
251 map.insert(id, LessSafeKey::new(unbound));
252 }
253 let current = current.into();
254 if !map.contains_key(¤t) {
255 return Err(CryptoError::InvalidKek);
256 }
257 Ok(Self {
258 keys: map,
259 current,
260 rng: SystemRandom::new(),
261 })
262 }
263
264 #[must_use]
266 pub fn current_kek_id(&self) -> &str {
267 &self.current
268 }
269
270 pub fn seal(
275 &self,
276 plaintext: &[u8],
277 binding: SecretBinding<'_>,
278 ) -> Result<SealedSecret, CryptoError> {
279 let key = self
280 .keys
281 .get(&self.current)
282 .ok_or(CryptoError::UnknownKekId)?;
283
284 let mut nonce_bytes = [0u8; NONCE_LEN];
285 self.rng
286 .fill(&mut nonce_bytes)
287 .map_err(|_| CryptoError::SealFailed)?;
288
289 let aad = binding.to_aad_bytes();
290 let mut buf = plaintext.to_vec();
291 key.seal_in_place_append_tag(
292 Nonce::assume_unique_for_key(nonce_bytes),
293 Aad::from(&aad),
294 &mut buf,
295 )
296 .map_err(|_| CryptoError::SealFailed)?;
297
298 Ok(SealedSecret {
299 kek_id: self.current.clone(),
300 nonce: nonce_bytes,
301 ciphertext: buf,
302 })
303 }
304
305 pub fn open(
316 &self,
317 sealed: &SealedSecret,
318 binding: SecretBinding<'_>,
319 ) -> Result<Zeroizing<Vec<u8>>, CryptoError> {
320 let key = self
321 .keys
322 .get(&sealed.kek_id)
323 .ok_or(CryptoError::UnknownKekId)?;
324
325 let aad = binding.to_aad_bytes();
326 let mut buf = sealed.ciphertext.clone();
327 let plaintext = key
328 .open_in_place(
329 Nonce::assume_unique_for_key(sealed.nonce),
330 Aad::from(&aad),
331 &mut buf,
332 )
333 .map_err(|_| CryptoError::OpenFailed)?;
334
335 Ok(Zeroizing::new(plaintext.to_vec()))
336 }
337
338 pub fn reseal(
344 &self,
345 sealed: &SealedSecret,
346 binding: SecretBinding<'_>,
347 ) -> Result<SealedSecret, CryptoError> {
348 let plaintext = self.open(sealed, binding)?;
349 self.seal(&plaintext, binding)
350 }
351}
352
353impl KeyRing {
354 pub fn from_spec(spec: &str, current: &str) -> Result<Self, CryptoError> {
365 let mut keys: Vec<(String, Zeroizing<Vec<u8>>)> = Vec::new();
366 for entry in spec.split(',') {
367 let entry = entry.trim();
368 if entry.is_empty() {
369 continue;
370 }
371 let (id, hex) = entry.split_once(':').ok_or(CryptoError::InvalidKek)?;
375 let id = id.trim();
376 if id.is_empty() || hex.contains(':') {
377 return Err(CryptoError::InvalidKek);
378 }
379 keys.push((id.to_string(), decode_hex(hex.trim())?));
380 }
381 if keys.is_empty() {
382 return Err(CryptoError::InvalidKek);
383 }
384 Self::new(keys, current)
385 }
386}
387
388pub const KEK_SPEC_ENV: &str = "BRAIN_EDGE_LLM_KEKS";
390pub const KEK_CURRENT_ENV: &str = "BRAIN_EDGE_LLM_KEK_CURRENT";
392
393pub fn keyring_from_env() -> Result<Option<KeyRing>, String> {
408 let Ok(spec) = std::env::var(KEK_SPEC_ENV) else {
409 return Ok(None);
410 };
411 if spec.trim().is_empty() {
412 return Ok(None);
413 }
414
415 let ids: Vec<String> = spec
416 .split(',')
417 .filter_map(|e| e.split_once(':').map(|(id, _)| id.trim().to_string()))
418 .filter(|id| !id.is_empty())
419 .collect();
420
421 let current = match std::env::var(KEK_CURRENT_ENV) {
422 Ok(c) if !c.trim().is_empty() => c.trim().to_string(),
423 _ => {
424 if ids.len() == 1 {
425 ids[0].clone()
426 } else {
427 return Err(format!(
428 "{KEK_CURRENT_ENV} must name which key id to seal under when \
429 {KEK_SPEC_ENV} carries more than one key"
430 ));
431 }
432 }
433 };
434
435 KeyRing::from_spec(&spec, ¤t)
438 .map(Some)
439 .map_err(|e| format!("{KEK_SPEC_ENV}: {e}"))
440}
441
442fn decode_hex(hex: &str) -> Result<Zeroizing<Vec<u8>>, CryptoError> {
447 if !hex.len().is_multiple_of(2) {
448 return Err(CryptoError::InvalidKek);
449 }
450 let (pairs, rest) = hex.as_bytes().as_chunks::<2>();
451 debug_assert!(rest.is_empty(), "length checked above");
452 let mut out = Zeroizing::new(Vec::with_capacity(hex.len() / 2));
453 for pair in pairs {
454 let hi = hex_nibble(pair[0]).ok_or(CryptoError::InvalidKek)?;
455 let lo = hex_nibble(pair[1]).ok_or(CryptoError::InvalidKek)?;
456 out.push((hi << 4) | lo);
457 }
458 Ok(out)
459}
460
461const fn hex_nibble(c: u8) -> Option<u8> {
462 match c {
463 b'0'..=b'9' => Some(c - b'0'),
464 b'a'..=b'f' => Some(c - b'a' + 10),
465 b'A'..=b'F' => Some(c - b'A' + 10),
466 _ => None,
467 }
468}
469
470#[cfg(test)]
471mod tests {
472 use super::*;
473
474 fn kek(byte: u8) -> Zeroizing<Vec<u8>> {
475 Zeroizing::new(vec![byte; KEK_LEN])
476 }
477
478 fn ring_with(id: &str, byte: u8) -> KeyRing {
479 KeyRing::new([(id.to_string(), kek(byte))], id).expect("valid ring")
480 }
481
482 fn binding<'a>(ns: &'a str, provider: &'a str) -> SecretBinding<'a> {
483 SecretBinding {
484 namespace: ns,
485 provider,
486 }
487 }
488
489 #[test]
490 fn seals_and_opens_round_trip() {
491 let r = ring_with("k1", 7);
492 let b = binding("acme", "openai");
493 let sealed = r.seal(b"sk-secret-value", b).unwrap();
494 assert_eq!(&*r.open(&sealed, b).unwrap(), b"sk-secret-value");
495 }
496
497 #[test]
498 fn ciphertext_does_not_contain_the_plaintext() {
499 let r = ring_with("k1", 7);
502 let secret = b"sk-live-0123456789";
503 let sealed = r.seal(secret, binding("acme", "openai")).unwrap();
504 assert!(
505 !sealed.ciphertext.windows(secret.len()).any(|w| w == secret),
506 "plaintext leaked into ciphertext"
507 );
508 }
509
510 #[test]
511 fn a_ciphertext_cannot_be_replayed_under_another_namespace() {
512 let r = ring_with("k1", 7);
515 let sealed = r.seal(b"acme-key", binding("acme", "openai")).unwrap();
516 assert_eq!(
517 r.open(&sealed, binding("globex", "openai")).unwrap_err(),
518 CryptoError::OpenFailed
519 );
520 }
521
522 #[test]
523 fn a_ciphertext_cannot_be_replayed_under_another_provider() {
524 let r = ring_with("k1", 7);
525 let sealed = r.seal(b"oa-key", binding("acme", "openai")).unwrap();
526 assert_eq!(
527 r.open(&sealed, binding("acme", "anthropic")).unwrap_err(),
528 CryptoError::OpenFailed
529 );
530 }
531
532 #[test]
533 fn binding_is_unambiguous_across_separator_lookalikes() {
534 let r = ring_with("k1", 7);
538 let sealed = r.seal(b"v", binding("a:b", "c")).unwrap();
539 assert_eq!(
540 r.open(&sealed, binding("a", "b:c")).unwrap_err(),
541 CryptoError::OpenFailed
542 );
543 }
544
545 #[test]
546 fn a_tampered_ciphertext_does_not_open() {
547 let r = ring_with("k1", 7);
548 let b = binding("acme", "openai");
549 let mut sealed = r.seal(b"sk-secret", b).unwrap();
550 sealed.ciphertext[0] ^= 0xFF;
551 assert_eq!(r.open(&sealed, b).unwrap_err(), CryptoError::OpenFailed);
552 }
553
554 #[test]
555 fn a_wrong_key_does_not_open() {
556 let a = ring_with("k1", 1);
557 let bogus = ring_with("k1", 2); let b = binding("acme", "openai");
559 let sealed = a.seal(b"sk-secret", b).unwrap();
560 assert_eq!(bogus.open(&sealed, b).unwrap_err(), CryptoError::OpenFailed);
561 }
562
563 #[test]
564 fn nonces_differ_across_seals_of_identical_input() {
565 let r = ring_with("k1", 7);
568 let b = binding("acme", "openai");
569 let one = r.seal(b"same", b).unwrap();
570 let two = r.seal(b"same", b).unwrap();
571 assert_ne!(one.nonce, two.nonce);
572 assert_ne!(one.ciphertext, two.ciphertext);
573 }
574
575 #[test]
576 fn rotation_keeps_old_ciphertext_readable() {
577 let r = KeyRing::new(
579 [("old".to_string(), kek(1)), ("new".to_string(), kek(2))],
580 "new",
581 )
582 .unwrap();
583 let b = binding("acme", "openai");
584
585 let old_ring = ring_with("old", 1);
587 let old_sealed = old_ring.seal(b"legacy", b).unwrap();
588 assert_eq!(old_sealed.kek_id, "old");
589
590 assert_eq!(&*r.open(&old_sealed, b).unwrap(), b"legacy");
591
592 let resealed = r.reseal(&old_sealed, b).unwrap();
593 assert_eq!(resealed.kek_id, "new");
594 assert_eq!(&*r.open(&resealed, b).unwrap(), b"legacy");
595 }
596
597 #[test]
598 fn an_unknown_kek_id_is_distinguishable_from_a_failed_open() {
599 let r = ring_with("k1", 7);
602 let b = binding("acme", "openai");
603 let mut sealed = r.seal(b"v", b).unwrap();
604 sealed.kek_id = "retired".to_string();
605 assert_eq!(r.open(&sealed, b).unwrap_err(), CryptoError::UnknownKekId);
606 }
607
608 #[test]
609 fn a_short_key_is_rejected_at_construction() {
610 let short = Zeroizing::new(vec![0u8; KEK_LEN - 1]);
611 assert_eq!(
612 KeyRing::new([("k1".to_string(), short)], "k1").unwrap_err(),
613 CryptoError::InvalidKek
614 );
615 }
616
617 #[test]
618 fn a_current_id_not_in_the_ring_is_rejected() {
619 assert_eq!(
621 KeyRing::new([("k1".to_string(), kek(1))], "k2").unwrap_err(),
622 CryptoError::InvalidKek
623 );
624 }
625
626 #[test]
627 fn debug_impls_do_not_leak_secrets() {
628 let r = ring_with("k1", 7);
629 let sealed = r
630 .seal(b"sk-super-secret", binding("acme", "openai"))
631 .unwrap();
632 let d = format!("{sealed:?}");
633 assert!(!d.contains("sk-super-secret"));
634 assert!(!d.contains("nonce"), "nonce must not be printed: {d}");
635 assert!(d.contains("k1"), "kek_id is safe and useful: {d}");
636 let rd = format!("{r:?}");
638 assert!(rd.contains("k1"));
639 assert!(!rd.contains('\u{7}'));
640 }
641}
642
643#[cfg(test)]
644mod spec_tests {
645 use super::*;
646
647 const K1: &str = "1111111111111111111111111111111111111111111111111111111111111111";
648 const K2: &str = "2222222222222222222222222222222222222222222222222222222222222222";
649
650 #[test]
651 fn parses_a_single_key_spec() {
652 let r = KeyRing::from_spec(&format!("k1:{K1}"), "k1").unwrap();
653 assert_eq!(r.current_kek_id(), "k1");
654 }
655
656 #[test]
657 fn parses_a_rotation_spec_and_seals_under_current() {
658 let r = KeyRing::from_spec(&format!("old:{K1},new:{K2}"), "new").unwrap();
659 assert_eq!(r.current_kek_id(), "new");
660 let sealed = r
661 .seal(
662 b"v",
663 SecretBinding {
664 namespace: "acme",
665 provider: "openai",
666 },
667 )
668 .unwrap();
669 assert_eq!(sealed.kek_id, "new");
670 }
671
672 #[test]
673 fn rejects_malformed_specs_rather_than_guessing() {
674 for bad in [
675 "", "k1", &format!("k1:{K1}:extra"), &format!(":{K1}"), "k1:zz", "k1:111", &format!("k1:{}", &K1[..62]), ] {
683 assert_eq!(
684 KeyRing::from_spec(bad, "k1").unwrap_err(),
685 CryptoError::InvalidKek,
686 "should reject {bad:?}"
687 );
688 }
689 }
690
691 #[test]
692 fn rejects_a_current_id_absent_from_the_spec() {
693 assert_eq!(
694 KeyRing::from_spec(&format!("k1:{K1}"), "k2").unwrap_err(),
695 CryptoError::InvalidKek
696 );
697 }
698
699 #[test]
700 fn accepts_uppercase_hex_and_surrounding_whitespace() {
701 let spec = format!(" k1 : {} ", K1.to_uppercase());
702 assert!(KeyRing::from_spec(&spec, "k1").is_ok());
703 }
704}