Lumotia/crates/cloud-providers/src/keystore.rs

use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

/// Store an API key in Lumotia'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 `LUMOTIA_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 Lumotia's process-local keystore.
///
/// Returns a previously stored in-memory key when present, otherwise falls
/// back to the read-only `LUMOTIA_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!("LUMOTIA_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()));
    }
}