use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use parking_lot::RwLock;
use crate::errors::FnError;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct SecretHandle(u64);
impl SecretHandle {
#[must_use]
pub fn opaque_id(&self) -> u64 {
self.0
}
}
#[derive(Debug)]
pub struct SecretStore {
by_name: RwLock<HashMap<String, Vec<u8>>>,
by_handle: RwLock<HashMap<u64, String>>,
next: AtomicU64,
}
impl Default for SecretStore {
fn default() -> Self {
Self {
by_name: RwLock::default(),
by_handle: RwLock::default(),
next: AtomicU64::new(1),
}
}
}
impl SecretStore {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn seal(&self, name: impl Into<String>, bytes: Vec<u8>) {
let name = name.into();
self.by_name.write().insert(name, bytes);
}
pub fn acquire(&self, name: &str) -> Result<SecretHandle, FnError> {
let exists = self.by_name.read().contains_key(name);
if !exists {
return Err(FnError::new(
0xA00,
format!("secret `{name}` not found in store"),
));
}
let id = self.next.fetch_add(1, Ordering::SeqCst);
let id = if id == 0 {
self.next.fetch_add(1, Ordering::SeqCst)
} else {
id
};
self.by_handle.write().insert(id, name.to_owned());
tracing::debug!(secret_id = name, handle_opaque = id, "secret.acquire");
Ok(SecretHandle(id))
}
pub fn unseal_for_host_use(&self, h: SecretHandle) -> Result<Vec<u8>, FnError> {
let by_handle = self.by_handle.read();
let name = by_handle.get(&h.0).ok_or_else(|| {
FnError::new(
0xA01,
format!("secret handle {} is invalid or revoked", h.0),
)
})?;
let by_name = self.by_name.read();
by_name.get(name).cloned().ok_or_else(|| {
FnError::new(0xA02, format!("secret `{name}` was sealed but is now gone"))
})
}
pub fn revoke(&self, h: SecretHandle) {
self.by_handle.write().remove(&h.0);
}
pub fn clear(&self) {
self.by_name.write().clear();
self.by_handle.write().clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn acquire_returns_handle_for_sealed_secret() {
let s = SecretStore::new();
s.seal("api_key", b"sk-test-abc".to_vec());
let h = s.acquire("api_key").unwrap();
assert_ne!(h.opaque_id(), 0);
}
#[test]
fn acquire_missing_secret_errors() {
let s = SecretStore::new();
assert!(s.acquire("nope").is_err());
}
#[test]
fn host_unseal_returns_bytes() {
let s = SecretStore::new();
s.seal("api_key", b"sk-test-abc".to_vec());
let h = s.acquire("api_key").unwrap();
let bytes = s.unseal_for_host_use(h).unwrap();
assert_eq!(bytes, b"sk-test-abc");
}
#[test]
fn revoked_handle_cannot_unseal() {
let s = SecretStore::new();
s.seal("api_key", b"x".to_vec());
let h = s.acquire("api_key").unwrap();
s.revoke(h);
assert!(s.unseal_for_host_use(h).is_err());
}
#[test]
fn separate_stores_handles_dont_cross() {
let a = SecretStore::new();
let b = SecretStore::new();
a.seal("k", b"av".to_vec());
b.seal("k", b"bv".to_vec());
let ha = a.acquire("k").unwrap();
assert!(b.unseal_for_host_use(ha).is_err());
}
#[test]
fn acquire_never_returns_zero_handle() {
let s = SecretStore::new();
s.seal("k", b"v".to_vec());
s.next.store(0, Ordering::SeqCst);
let h = s.acquire("k").unwrap();
assert_ne!(h.opaque_id(), 0);
}
}