Skip to main content

moq_secure/
key_store.rs

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