Skip to main content

moq_secure/
key_store.rs

1use base64::Engine;
2#[allow(unused_imports)]
3use zeroize::{Zeroize, Zeroizing};
4
5pub trait KeyStore: Send + Sync {
6    fn aead_key(&self, key_id: u8) -> Option<&[u8; 32]>;
7}
8
9#[derive(Debug, thiserror::Error)]
10pub enum KeyStoreError {
11    #[error("key_id slot {0} is not loaded")]
12    KeyNotLoaded(u8),
13
14    #[error("expected 32 bytes (decoded {0} bytes)")]
15    KeyWrongLength(usize),
16
17    #[error("failed to decode key as hex/base64")]
18    DecodeFailed(#[from] DecodeFailed),
19}
20
21#[derive(Debug, thiserror::Error)]
22pub enum DecodeFailed {
23    #[error("hex decode failed")]
24    Hex(#[from] hex::FromHexError),
25
26    #[error("base64 decode failed")]
27    Base64(#[from] base64::DecodeError),
28
29    #[error("key string looked like hex but had the wrong length")]
30    HexWrongLength,
31
32    #[error("unknown decode error")]
33    Other,
34}
35
36#[derive(Debug)]
37pub struct InMemoryKeyStore {
38    // Unloaded slots contain no key material.
39    // Loaded keys are zeroized when removed or when the store is dropped.
40    keys: [Option<Zeroizing<[u8; 32]>>; 256],
41}
42
43impl InMemoryKeyStore {
44    pub fn empty() -> Self {
45        Self {
46            // `from_fn` avoids initializing the backing storage with
47            // hard-coded cryptographic-looking byte arrays.
48            keys: std::array::from_fn(|_| None),
49        }
50    }
51
52    pub fn set_key(&mut self, key_id: u8, key: [u8; 32]) {
53        self.keys[key_id as usize] = Some(Zeroizing::new(key));
54    }
55
56    /// Accepts either:
57    /// - hex: 64 hex characters representing 32 bytes
58    /// - base64: decodes to exactly 32 bytes, with or without padding
59    pub fn set_key_encoded(
60        &mut self,
61        key_id: u8,
62        key_encoded: &str,
63    ) -> Result<(), KeyStoreError> {
64        let key_encoded = key_encoded.trim();
65
66        // Try hex only when the input has exactly the expected hex length.
67        // This avoids accidentally treating arbitrary base64 as hex.
68        if key_encoded.len() == 64
69            && key_encoded
70                .chars()
71                .all(|character| character.is_ascii_hexdigit())
72        {
73            let decoded = hex::decode(key_encoded).map_err(DecodeFailed::from)?;
74            let decoded = Zeroizing::new(decoded);
75
76            let key: [u8; 32] = decoded
77                .as_slice()
78                .try_into()
79                .map_err(|_| KeyStoreError::KeyWrongLength(decoded.len()))?;
80
81            self.set_key(key_id, key);
82            return Ok(());
83        }
84
85        // Keep decoded key bytes in zeroizing storage while converting them.
86        let decoded = base64::engine::general_purpose::STANDARD
87            .decode(key_encoded)
88            .or_else(|_| {
89                base64::engine::general_purpose::STANDARD_NO_PAD.decode(key_encoded)
90            })
91            .map_err(DecodeFailed::from)?;
92
93        let decoded = Zeroizing::new(decoded);
94
95        if decoded.len() != 32 {
96            return Err(KeyStoreError::KeyWrongLength(decoded.len()));
97        }
98
99        let key: [u8; 32] = decoded
100            .as_slice()
101            .try_into()
102            .map_err(|_| KeyStoreError::KeyWrongLength(decoded.len()))?;
103
104        self.set_key(key_id, key);
105        Ok(())
106    }
107
108    pub fn remove_key(&mut self, key_id: u8) -> Result<(), KeyStoreError> {
109        self.keys[key_id as usize]
110            .take()
111            .map(|_| ())
112            .ok_or(KeyStoreError::KeyNotLoaded(key_id))
113    }
114
115    pub fn contains_key(&self, key_id: u8) -> bool {
116        self.keys[key_id as usize].is_some()
117    }
118}
119
120impl KeyStore for InMemoryKeyStore {
121    fn aead_key(&self, key_id: u8) -> Option<&[u8; 32]> {
122        self.keys[key_id as usize]
123            .as_ref()
124            .map(|key| &**key)
125    }
126}
127