Skip to main content

codetether_agent/secrets/
mod.rs

1//! Secrets management via HashiCorp Vault
2//!
3//! All API keys and secrets are loaded exclusively from HashiCorp Vault.
4//! Environment variables are NOT used for secrets.
5
6use anyhow::{Context, Result};
7use std::collections::HashMap;
8use std::sync::Arc;
9use tokio::sync::RwLock;
10use vaultrs::client::{VaultClient, VaultClientSettingsBuilder};
11use vaultrs::kv2;
12
13/// Path in Vault where provider secrets are stored
14#[allow(dead_code)]
15const DEFAULT_SECRETS_PATH: &str = "secret/data/codetether/providers";
16
17/// Vault-based secrets manager
18#[derive(Clone)]
19pub struct SecretsManager {
20    client: Option<Arc<VaultClient>>,
21    /// Cache of loaded API keys (provider_id -> api_key)
22    pub cache: Arc<RwLock<HashMap<String, String>>>,
23    mount: String,
24    path: String,
25}
26
27impl Default for SecretsManager {
28    fn default() -> Self {
29        Self {
30            client: None,
31            cache: Arc::new(RwLock::new(HashMap::new())),
32            mount: "secret".to_string(),
33            path: "codetether/providers".to_string(),
34        }
35    }
36}
37
38impl SecretsManager {
39    /// Create a new secrets manager with Vault configuration
40    pub async fn new(config: &VaultConfig) -> Result<Self> {
41        let settings = VaultClientSettingsBuilder::default()
42            .address(&config.address)
43            .token(&config.token)
44            .build()
45            .context("Failed to build Vault client settings")?;
46
47        let client = VaultClient::new(settings).context("Failed to create Vault client")?;
48
49        Ok(Self {
50            client: Some(Arc::new(client)),
51            cache: Arc::new(RwLock::new(HashMap::new())),
52            mount: config.mount.clone().unwrap_or_else(|| "secret".to_string()),
53            path: config
54                .path
55                .clone()
56                .unwrap_or_else(|| "codetether/providers".to_string()),
57        })
58    }
59
60    /// Authenticate to Vault using the pod's Kubernetes service account JWT.
61    ///
62    /// Reads the SA JWT from the standard Kubernetes mount path (overridable via
63    /// `VAULT_K8S_SA_JWT_PATH`) then calls the Vault `auth/kubernetes/login`
64    /// endpoint.  The returned manager holds the short-lived token that Vault
65    /// issued — no `VAULT_TOKEN` environment variable is required.
66    pub async fn from_k8s_auth(
67        address: &str,
68        role: &str,
69        mount: &str,
70        kv_mount: Option<&str>,
71        kv_path: Option<&str>,
72    ) -> Result<Self> {
73        let jwt_path = std::env::var("VAULT_K8S_SA_JWT_PATH")
74            .unwrap_or_else(|_| "/var/run/secrets/kubernetes.io/serviceaccount/token".to_string());
75
76        let jwt = tokio::fs::read_to_string(&jwt_path)
77            .await
78            .with_context(|| {
79                format!(
80                    "Failed to read Kubernetes service account token from {}",
81                    jwt_path
82                )
83            })?;
84        let jwt = jwt.trim().to_string();
85
86        // Bootstrap client with an empty token — only used for the one-shot
87        // auth call; the real authenticated client is built from the result.
88        let bootstrap_settings = VaultClientSettingsBuilder::default()
89            .address(address)
90            .token("")
91            .build()
92            .context("Failed to build bootstrap Vault client settings")?;
93        let bootstrap_client = VaultClient::new(bootstrap_settings)
94            .context("Failed to create bootstrap Vault client")?;
95
96        let auth_info = vaultrs::auth::kubernetes::login(&bootstrap_client, mount, role, &jwt)
97            .await
98            .context("Vault Kubernetes auth login failed")?;
99
100        let settings = VaultClientSettingsBuilder::default()
101            .address(address)
102            .token(&auth_info.client_token)
103            .build()
104            .context("Failed to build authenticated Vault client settings")?;
105        let client =
106            VaultClient::new(settings).context("Failed to create authenticated Vault client")?;
107
108        Ok(Self {
109            client: Some(Arc::new(client)),
110            cache: Arc::new(RwLock::new(HashMap::new())),
111            mount: kv_mount.unwrap_or("secret").to_string(),
112            path: kv_path.unwrap_or("codetether/providers").to_string(),
113        })
114    }
115
116    /// Try to create from environment (for initial bootstrap only).
117    ///
118    /// When `VAULT_ROLE` is set the worker authenticates via Kubernetes service
119    /// account — no static token is needed and the resulting Vault token is
120    /// short-lived and automatically rotated by Vault itself.  Falls back to
121    /// `VAULT_TOKEN` when `VAULT_ROLE` is absent or K8s auth fails.
122    pub async fn from_env() -> Result<Self> {
123        let address = std::env::var("VAULT_ADDR").context("VAULT_ADDR not set")?;
124        let kv_mount = std::env::var("VAULT_MOUNT").ok();
125        let kv_path = std::env::var("VAULT_SECRETS_PATH").ok();
126
127        // Prefer Kubernetes service-account auth when VAULT_ROLE is set.
128        // This eliminates the dependency on a static VAULT_TOKEN; the pod's own
129        // SA JWT (mounted by k8s at the standard path) is the only credential
130        // the container needs to carry.
131        if let Ok(role) = std::env::var("VAULT_ROLE") {
132            let role = role.trim().to_string();
133            if !role.is_empty() {
134                let k8s_mount =
135                    std::env::var("VAULT_AUTH_MOUNT").unwrap_or_else(|_| "kubernetes".to_string());
136
137                match Self::from_k8s_auth(
138                    &address,
139                    &role,
140                    &k8s_mount,
141                    kv_mount.as_deref(),
142                    kv_path.as_deref(),
143                )
144                .await
145                {
146                    Ok(manager) => {
147                        tracing::info!(
148                            role = %role,
149                            mount = %k8s_mount,
150                            "Authenticated to Vault via Kubernetes service account"
151                        );
152                        return Ok(manager);
153                    }
154                    Err(e) => {
155                        tracing::warn!(
156                            error = %e,
157                            "Vault Kubernetes auth failed; falling back to VAULT_TOKEN"
158                        );
159                    }
160                }
161            }
162        }
163
164        let token = std::env::var("VAULT_TOKEN").context("VAULT_TOKEN not set")?;
165        let config = VaultConfig {
166            address,
167            token,
168            mount: kv_mount,
169            path: kv_path,
170        };
171
172        Self::new(&config).await
173    }
174
175    /// Check if Vault is configured and connected
176    pub fn is_connected(&self) -> bool {
177        self.client.is_some()
178    }
179
180    /// Get an API key for a provider from Vault
181    pub async fn get_api_key(&self, provider_id: &str) -> Result<Option<String>> {
182        // Check cache first
183        {
184            let cache = self.cache.read().await;
185            if let Some(key) = cache.get(provider_id) {
186                return Ok(Some(key.clone()));
187            }
188        }
189
190        // Fetch from Vault
191        let client = match &self.client {
192            Some(c) => c,
193            None => return Ok(None),
194        };
195
196        let secret_path = format!("{}/{}", self.path, provider_id);
197
198        match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
199            Ok(secret) => {
200                // Cache the result
201                if let Some(ref api_key) = secret.api_key {
202                    let mut cache = self.cache.write().await;
203                    cache.insert(provider_id.to_string(), api_key.clone());
204                }
205                Ok(secret.api_key)
206            }
207            Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
208            Err(e) => {
209                tracing::warn!("Failed to fetch secret for {}: {}", provider_id, e);
210                Ok(None)
211            }
212        }
213    }
214
215    /// Get all secrets for a provider
216    pub async fn get_provider_secrets(&self, provider_id: &str) -> Result<Option<ProviderSecrets>> {
217        let client = match &self.client {
218            Some(c) => c,
219            None => return Ok(None),
220        };
221
222        let secret_path = format!("{}/{}", self.path, provider_id);
223
224        match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
225            Ok(secret) => Ok(Some(secret)),
226            Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
227            Err(e) => {
228                tracing::warn!("Failed to fetch secrets for {}: {}", provider_id, e);
229                Ok(None)
230            }
231        }
232    }
233
234    /// Set/update secrets for a provider in Vault
235    pub async fn set_provider_secrets(
236        &self,
237        provider_id: &str,
238        secrets: &ProviderSecrets,
239    ) -> Result<()> {
240        let client = match &self.client {
241            Some(c) => c,
242            None => anyhow::bail!("Vault client not configured"),
243        };
244
245        let secret_path = format!("{}/{}", self.path, provider_id);
246        kv2::set(client.as_ref(), &self.mount, &secret_path, secrets)
247            .await
248            .with_context(|| format!("Failed to write provider secrets for {}", provider_id))?;
249
250        // Update cache with latest API key value
251        let mut cache = self.cache.write().await;
252        if let Some(api_key) = secrets.api_key.clone() {
253            cache.insert(provider_id.to_string(), api_key);
254        } else {
255            cache.remove(provider_id);
256        }
257
258        Ok(())
259    }
260
261    /// Check if a provider has an API key in Vault
262    pub async fn has_api_key(&self, provider_id: &str) -> bool {
263        matches!(self.get_api_key(provider_id).await, Ok(Some(_)))
264    }
265
266    /// List all providers that have secrets configured
267    pub async fn list_configured_providers(&self) -> Result<Vec<String>> {
268        let client = match &self.client {
269            Some(c) => c,
270            None => return Ok(Vec::new()),
271        };
272
273        match kv2::list(client.as_ref(), &self.mount, &self.path).await {
274            Ok(keys) => Ok(keys),
275            Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(Vec::new()),
276            Err(e) => {
277                tracing::warn!("Failed to list providers: {}", e);
278                Ok(Vec::new())
279            }
280        }
281    }
282
283    /// Clear the cache (useful when secrets are rotated)
284    pub async fn clear_cache(&self) {
285        let mut cache = self.cache.write().await;
286        cache.clear();
287    }
288}
289
290/// Vault configuration
291#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
292pub struct VaultConfig {
293    /// Vault server address (e.g., "https://vault.example.com:8200")
294    pub address: String,
295
296    /// Vault token for authentication
297    pub token: String,
298
299    /// KV secrets engine mount path (default: "secret")
300    #[serde(default)]
301    pub mount: Option<String>,
302
303    /// Path prefix for provider secrets (default: "codetether/providers")
304    #[serde(default)]
305    pub path: Option<String>,
306}
307
308impl Default for VaultConfig {
309    fn default() -> Self {
310        Self {
311            address: String::new(),
312            token: String::new(),
313            mount: Some("secret".to_string()),
314            path: Some("codetether/providers".to_string()),
315        }
316    }
317}
318
319/// Provider secrets stored in Vault
320#[derive(Clone, serde::Serialize, serde::Deserialize)]
321pub struct ProviderSecrets {
322    /// API key for the provider
323    #[serde(default)]
324    pub api_key: Option<String>,
325
326    /// Base URL override
327    #[serde(default)]
328    pub base_url: Option<String>,
329
330    /// Organization ID (for OpenAI)
331    #[serde(default)]
332    pub organization: Option<String>,
333
334    /// Additional headers as JSON
335    #[serde(default)]
336    pub headers: Option<HashMap<String, String>>,
337
338    /// Any provider-specific extra fields
339    #[serde(flatten)]
340    pub extra: HashMap<String, serde_json::Value>,
341}
342
343impl std::fmt::Debug for ProviderSecrets {
344    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345        f.debug_struct("ProviderSecrets")
346            .field("api_key", &self.api_key.as_ref().map(|_| "<REDACTED>"))
347            .field("api_key_len", &self.api_key.as_ref().map(|k| k.len()))
348            .field("base_url", &self.base_url)
349            .field("organization", &self.organization)
350            .field("headers_present", &self.headers.is_some())
351            .field("extra_fields", &self.extra.len())
352            .finish()
353    }
354}
355
356impl ProviderSecrets {
357    /// Check if API key is present and valid (non-empty)
358    pub fn has_valid_api_key(&self) -> bool {
359        self.api_key
360            .as_ref()
361            .map(|k| !k.is_empty())
362            .unwrap_or(false)
363    }
364
365    /// Get API key length without exposing the key
366    pub fn api_key_len(&self) -> Option<usize> {
367        self.api_key.as_ref().map(|k| k.len())
368    }
369}
370
371/// Global secrets manager instance
372static SECRETS_MANAGER: tokio::sync::OnceCell<SecretsManager> = tokio::sync::OnceCell::const_new();
373
374/// Initialize the global secrets manager
375pub async fn init_secrets_manager(config: &VaultConfig) -> Result<()> {
376    let manager = SecretsManager::new(config).await?;
377    SECRETS_MANAGER
378        .set(manager)
379        .map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
380    Ok(())
381}
382
383/// Initialize the global secrets manager from an existing manager instance
384pub fn init_from_manager(manager: SecretsManager) -> Result<()> {
385    SECRETS_MANAGER
386        .set(manager)
387        .map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
388    Ok(())
389}
390
391/// Get the global secrets manager
392pub fn secrets_manager() -> Option<&'static SecretsManager> {
393    SECRETS_MANAGER.get()
394}
395
396/// Get API key for a provider (convenience function)
397pub async fn get_api_key(provider_id: &str) -> Option<String> {
398    match SECRETS_MANAGER.get() {
399        Some(manager) => manager.get_api_key(provider_id).await.ok().flatten(),
400        None => None,
401    }
402}
403
404/// Check if a provider has an API key (convenience function)
405pub async fn has_api_key(provider_id: &str) -> bool {
406    match SECRETS_MANAGER.get() {
407        Some(manager) => manager.has_api_key(provider_id).await,
408        None => false,
409    }
410}
411
412/// Get full provider secrets (convenience function)
413pub async fn get_provider_secrets(provider_id: &str) -> Option<ProviderSecrets> {
414    match SECRETS_MANAGER.get() {
415        Some(manager) => manager
416            .get_provider_secrets(provider_id)
417            .await
418            .ok()
419            .flatten(),
420        None => None,
421    }
422}
423
424/// Set full provider secrets (convenience function)
425pub async fn set_provider_secrets(provider_id: &str, secrets: &ProviderSecrets) -> Result<()> {
426    match SECRETS_MANAGER.get() {
427        Some(manager) => manager.set_provider_secrets(provider_id, secrets).await,
428        None => anyhow::bail!("Secrets manager not initialized"),
429    }
430}