use anyhow::{Context, Result};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use vaultrs::client::{VaultClient, VaultClientSettingsBuilder};
use vaultrs::kv2;
#[allow(dead_code)]
const DEFAULT_SECRETS_PATH: &str = "secret/data/codetether/providers";
#[derive(Clone)]
pub struct SecretsManager {
client: Option<Arc<VaultClient>>,
pub cache: Arc<RwLock<HashMap<String, String>>>,
mount: String,
path: String,
}
impl Default for SecretsManager {
fn default() -> Self {
Self {
client: None,
cache: Arc::new(RwLock::new(HashMap::new())),
mount: "secret".to_string(),
path: "codetether/providers".to_string(),
}
}
}
impl SecretsManager {
pub async fn new(config: &VaultConfig) -> Result<Self> {
let settings = VaultClientSettingsBuilder::default()
.address(&config.address)
.token(&config.token)
.build()
.context("Failed to build Vault client settings")?;
let client = VaultClient::new(settings).context("Failed to create Vault client")?;
Ok(Self {
client: Some(Arc::new(client)),
cache: Arc::new(RwLock::new(HashMap::new())),
mount: config.mount.clone().unwrap_or_else(|| "secret".to_string()),
path: config
.path
.clone()
.unwrap_or_else(|| "codetether/providers".to_string()),
})
}
pub async fn from_env() -> Result<Self> {
let address = std::env::var("VAULT_ADDR").context("VAULT_ADDR not set")?;
let token = std::env::var("VAULT_TOKEN").context("VAULT_TOKEN not set")?;
let mount = std::env::var("VAULT_MOUNT").ok();
let path = std::env::var("VAULT_SECRETS_PATH").ok();
let config = VaultConfig {
address,
token,
mount,
path,
};
Self::new(&config).await
}
pub fn is_connected(&self) -> bool {
self.client.is_some()
}
pub async fn get_api_key(&self, provider_id: &str) -> Result<Option<String>> {
{
let cache = self.cache.read().await;
if let Some(key) = cache.get(provider_id) {
return Ok(Some(key.clone()));
}
}
let client = match &self.client {
Some(c) => c,
None => return Ok(None),
};
let secret_path = format!("{}/{}", self.path, provider_id);
match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
Ok(secret) => {
if let Some(ref api_key) = secret.api_key {
let mut cache = self.cache.write().await;
cache.insert(provider_id.to_string(), api_key.clone());
}
Ok(secret.api_key)
}
Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
Err(e) => {
tracing::warn!("Failed to fetch secret for {}: {}", provider_id, e);
Ok(None)
}
}
}
pub async fn get_provider_secrets(&self, provider_id: &str) -> Result<Option<ProviderSecrets>> {
let client = match &self.client {
Some(c) => c,
None => return Ok(None),
};
let secret_path = format!("{}/{}", self.path, provider_id);
match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
Ok(secret) => Ok(Some(secret)),
Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
Err(e) => {
tracing::warn!("Failed to fetch secrets for {}: {}", provider_id, e);
Ok(None)
}
}
}
pub async fn set_provider_secrets(
&self,
provider_id: &str,
secrets: &ProviderSecrets,
) -> Result<()> {
let client = match &self.client {
Some(c) => c,
None => anyhow::bail!("Vault client not configured"),
};
let secret_path = format!("{}/{}", self.path, provider_id);
kv2::set(client.as_ref(), &self.mount, &secret_path, secrets)
.await
.with_context(|| format!("Failed to write provider secrets for {}", provider_id))?;
let mut cache = self.cache.write().await;
if let Some(api_key) = secrets.api_key.clone() {
cache.insert(provider_id.to_string(), api_key);
} else {
cache.remove(provider_id);
}
Ok(())
}
pub async fn has_api_key(&self, provider_id: &str) -> bool {
matches!(self.get_api_key(provider_id).await, Ok(Some(_)))
}
pub async fn list_configured_providers(&self) -> Result<Vec<String>> {
let client = match &self.client {
Some(c) => c,
None => return Ok(Vec::new()),
};
match kv2::list(client.as_ref(), &self.mount, &self.path).await {
Ok(keys) => Ok(keys),
Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(Vec::new()),
Err(e) => {
tracing::warn!("Failed to list providers: {}", e);
Ok(Vec::new())
}
}
}
pub async fn clear_cache(&self) {
let mut cache = self.cache.write().await;
cache.clear();
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct VaultConfig {
pub address: String,
pub token: String,
#[serde(default)]
pub mount: Option<String>,
#[serde(default)]
pub path: Option<String>,
}
impl Default for VaultConfig {
fn default() -> Self {
Self {
address: String::new(),
token: String::new(),
mount: Some("secret".to_string()),
path: Some("codetether/providers".to_string()),
}
}
}
#[derive(Clone, serde::Serialize, serde::Deserialize)]
pub struct ProviderSecrets {
#[serde(default)]
pub api_key: Option<String>,
#[serde(default)]
pub base_url: Option<String>,
#[serde(default)]
pub organization: Option<String>,
#[serde(default)]
pub headers: Option<HashMap<String, String>>,
#[serde(flatten)]
pub extra: HashMap<String, serde_json::Value>,
}
impl std::fmt::Debug for ProviderSecrets {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ProviderSecrets")
.field("api_key", &self.api_key.as_ref().map(|_| "<REDACTED>"))
.field("api_key_len", &self.api_key.as_ref().map(|k| k.len()))
.field("base_url", &self.base_url)
.field("organization", &self.organization)
.field("headers_present", &self.headers.is_some())
.field("extra_fields", &self.extra.len())
.finish()
}
}
impl ProviderSecrets {
pub fn has_valid_api_key(&self) -> bool {
self.api_key
.as_ref()
.map(|k| !k.is_empty())
.unwrap_or(false)
}
pub fn api_key_len(&self) -> Option<usize> {
self.api_key.as_ref().map(|k| k.len())
}
}
static SECRETS_MANAGER: tokio::sync::OnceCell<SecretsManager> = tokio::sync::OnceCell::const_new();
pub async fn init_secrets_manager(config: &VaultConfig) -> Result<()> {
let manager = SecretsManager::new(config).await?;
SECRETS_MANAGER
.set(manager)
.map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
Ok(())
}
pub fn init_from_manager(manager: SecretsManager) -> Result<()> {
SECRETS_MANAGER
.set(manager)
.map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
Ok(())
}
pub fn secrets_manager() -> Option<&'static SecretsManager> {
SECRETS_MANAGER.get()
}
pub async fn get_api_key(provider_id: &str) -> Option<String> {
match SECRETS_MANAGER.get() {
Some(manager) => manager.get_api_key(provider_id).await.ok().flatten(),
None => None,
}
}
pub async fn has_api_key(provider_id: &str) -> bool {
match SECRETS_MANAGER.get() {
Some(manager) => manager.has_api_key(provider_id).await,
None => false,
}
}
pub async fn get_provider_secrets(provider_id: &str) -> Option<ProviderSecrets> {
match SECRETS_MANAGER.get() {
Some(manager) => manager
.get_provider_secrets(provider_id)
.await
.ok()
.flatten(),
None => None,
}
}
pub async fn set_provider_secrets(provider_id: &str, secrets: &ProviderSecrets) -> Result<()> {
match SECRETS_MANAGER.get() {
Some(manager) => manager.set_provider_secrets(provider_id, secrets).await,
None => anyhow::bail!("Secrets manager not initialized"),
}
}