78 lines
2.5 KiB
Rust
78 lines
2.5 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::{Mutex, OnceLock};
|
|
|
|
/// Store an API key in Kon's process-local keystore.
|
|
///
|
|
/// Keys are held in memory for the lifetime of the process and are lost on
|
|
/// exit. This avoids the undefined behaviour of mutating process environment
|
|
/// variables from arbitrary threads while keeping the public API safe.
|
|
///
|
|
/// `retrieve_api_key` still falls back to `KON_API_KEY_<PROVIDER>` environment
|
|
/// variables so externally injected secrets continue to work.
|
|
///
|
|
/// TODO: Replace with the `keyring` crate (or platform-native credential
|
|
/// storage) so keys persist across sessions and are accessed safely.
|
|
pub fn store_api_key(provider: &str, key: &str) {
|
|
api_key_store()
|
|
.lock()
|
|
.unwrap()
|
|
.insert(provider_env_key(provider), key.to_string());
|
|
}
|
|
|
|
/// Retrieve an API key from Kon's process-local keystore.
|
|
///
|
|
/// Returns a previously stored in-memory key when present, otherwise falls
|
|
/// back to the read-only `KON_API_KEY_<PROVIDER>` environment variable so
|
|
/// operator-supplied secrets still work.
|
|
pub fn retrieve_api_key(provider: &str) -> Option<String> {
|
|
let env_key = provider_env_key(provider);
|
|
api_key_store()
|
|
.lock()
|
|
.unwrap()
|
|
.get(&env_key)
|
|
.cloned()
|
|
.or_else(|| std::env::var(env_key).ok())
|
|
}
|
|
|
|
fn api_key_store() -> &'static Mutex<HashMap<String, String>> {
|
|
static STORE: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
|
|
STORE.get_or_init(|| Mutex::new(HashMap::new()))
|
|
}
|
|
|
|
fn provider_env_key(provider: &str) -> String {
|
|
format!("KON_API_KEY_{}", provider.to_uppercase())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
fn unique_provider(prefix: &str) -> String {
|
|
static NEXT_ID: AtomicUsize = AtomicUsize::new(1);
|
|
format!("{prefix}_{}", NEXT_ID.fetch_add(1, Ordering::Relaxed))
|
|
}
|
|
|
|
#[test]
|
|
fn stored_key_is_retrievable_without_env_mutation() {
|
|
let provider = unique_provider("provider");
|
|
store_api_key(&provider, "secret-token");
|
|
assert_eq!(
|
|
retrieve_api_key(&provider),
|
|
Some("secret-token".to_string())
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn providers_do_not_overlap() {
|
|
let first = unique_provider("first");
|
|
let second = unique_provider("second");
|
|
|
|
store_api_key(&first, "alpha");
|
|
store_api_key(&second, "beta");
|
|
|
|
assert_eq!(retrieve_api_key(&first), Some("alpha".to_string()));
|
|
assert_eq!(retrieve_api_key(&second), Some("beta".to_string()));
|
|
}
|
|
}
|