use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; use std::time::{Duration, Instant}; use serde::Serialize; use tokio::sync::Mutex; use super::auth::{ensure_readonly_credential, ReadonlyCredential}; use super::discovery::{parse_workspace_list, workspace_url, Discovery, KcpInfo, WorkspaceInfo}; use super::runner::Runner; pub const PROTECTED_HOSTS: &[&str] = &["socialweb.computer", "hat-7", "hat-9"]; const DISCOVERY_TTL: Duration = Duration::from_secs(30); const TOKEN_REFRESH_MARGIN: Duration = Duration::from_secs(180); const DEFAULT_TOKEN_TTL_SECONDS: u64 = 900; type DiscoveryCache = Option<(Instant, Vec, Vec)>; #[derive(Clone, Debug)] struct CachedCredential { credential: ReadonlyCredential, expires_at: Instant, } pub struct KcpService { discovery: Discovery, runner: Runner, allow_prod: bool, token_ttl_seconds: u64, kcps: Mutex, credentials: Mutex>, } #[derive(Clone)] pub struct KcpServiceHandle(pub Arc); impl KcpServiceHandle { pub fn new(service: KcpService) -> Self { Self(Arc::new(service)) } } impl std::ops::Deref for KcpServiceHandle { type Target = KcpService; fn deref(&self) -> &Self::Target { &self.0 } } impl KcpService { pub fn new( runner: Runner, roots: Vec, kubeconfigs: Vec, allow_prod: bool, ) -> Self { let discovery = Discovery::new(runner.clone(), roots, kubeconfigs); Self { discovery, runner, allow_prod, token_ttl_seconds: DEFAULT_TOKEN_TTL_SECONDS, kcps: Mutex::new(None), credentials: Mutex::new(HashMap::new()), } } pub async fn discover(&self, refresh: bool) -> (Vec, Vec) { let mut cache = self.kcps.lock().await; if !refresh { if let Some((at, kcps, errors)) = cache.as_ref() { if at.elapsed() < DISCOVERY_TTL { return (kcps.clone(), errors.clone()); } } } let (mut kcps, mut errors) = self.discovery.discover().await; for kcp in &mut kcps { let workspace = kcp.workspace.clone(); match self .read( kcp, &workspace, &["get".to_string(), "--raw=/livez".to_string()], ) .await { Ok(output) => kcp.running = output.to_ascii_lowercase().contains("ok"), Err(error) => { kcp.running = false; kcp.error = Some(error.clone()); errors.push(format!("{}: {error}", kcp.id)); } } } *cache = Some((Instant::now(), kcps.clone(), errors.clone())); (kcps, errors) } pub async fn kcps(&self, refresh: bool) -> Result, String> { let (kcps, errors) = self.discover(refresh).await; if kcps.is_empty() { let detail = if errors.is_empty() { "no kcp admin.kubeconfig was found under the configured roots; pass --kubeconfig or --root".to_string() } else { errors.join("; ") }; return Err(format!("no running kcp was discovered: {detail}")); } Ok(kcps) } /// Resolves a kcp id or kubeconfig path. `None` picks the only discovered /// kcp, and fails with the list of ids when there are several. pub async fn resolve_kcp(&self, needle: Option<&str>) -> Result { let Some(needle) = needle.filter(|needle| !needle.is_empty()) else { let mut kcps = self.kcps(false).await?; if kcps.is_empty() { kcps = self.discover(true).await.0; } return match kcps.as_slice() { [only] => Ok(only.clone()), [] => Err("no kcp discovered; pass kcp=".to_string()), many => Err(format!( "several kcps running; pass kcp=, one of: {}", many.iter() .map(|kcp| kcp.id.as_str()) .collect::>() .join(", ") )), }; }; let mut kcps = self.kcps(false).await?; if let Some(found) = find_kcp(&kcps, needle) { return Ok(found.clone()); } let (fresh, _) = self.discover(true).await; kcps = fresh; find_kcp(&kcps, needle).cloned().ok_or_else(|| { format!( "unknown kcp {needle:?}; known kcps: {}", kcps.iter() .map(|kcp| kcp.id.as_str()) .collect::>() .join(", ") ) }) } pub async fn list_workspaces( &self, kcp: &KcpInfo, workspace: Option<&str>, recursive: bool, ) -> Result, String> { self.ensure_allowed(kcp, "list workspaces")?; let start = workspace.unwrap_or(&kcp.workspace).to_string(); let mut seen = BTreeSet::new(); let mut result = Vec::new(); let mut pending = vec![start]; while let Some(parent) = pending.pop() { if !seen.insert(parent.clone()) { continue; } let output = self .read( kcp, &parent, &[ "get".to_string(), "workspaces".to_string(), "-o".to_string(), "json".to_string(), ], ) .await?; for workspace in parse_workspace_list(&output, &parent, &kcp.server)? { if recursive { pending.push(workspace.path.clone()); } result.push(workspace); } } result.sort_by(|left, right| left.path.cmp(&right.path)); Ok(result) } pub fn workspace_or_default(&self, kcp: &KcpInfo, workspace: Option<&str>) -> String { workspace .map(str::to_string) .filter(|workspace| !workspace.is_empty()) .unwrap_or_else(|| kcp.workspace.clone()) } pub fn workspace_url(&self, kcp: &KcpInfo, workspace: &str) -> String { workspace_url(&kcp.server, workspace) } pub async fn read( &self, kcp: &KcpInfo, workspace: &str, args: &[String], ) -> Result { self.ensure_allowed(kcp, "read")?; let credential = self.credential(kcp, workspace).await?; let mut command = vec![ "--kubeconfig".to_string(), credential.kubeconfig.clone(), "--server".to_string(), self.workspace_url(kcp, workspace), "--request-timeout=30s".to_string(), ]; command.extend(args.iter().cloned()); self.runner.kubectl(&command, None).await } /// Starts a streaming read (`kubectl get -w`) with the read credential and /// no request timeout; the caller owns the child. pub async fn read_stream( &self, kcp: &KcpInfo, workspace: &str, args: &[String], ) -> Result { self.ensure_allowed(kcp, "watch")?; let credential = self.credential(kcp, workspace).await?; let mut command = vec![ "--kubeconfig".to_string(), credential.kubeconfig.clone(), "--server".to_string(), self.workspace_url(kcp, workspace), "--request-timeout=0".to_string(), ]; command.extend(args.iter().cloned()); self.runner.spawn_kubectl_stream(&command) } pub async fn admin( &self, kcp: &KcpInfo, workspace: &str, args: &[String], stdin: Option<&str>, ) -> Result { self.ensure_allowed(kcp, "write")?; let mut command = vec![ "--kubeconfig".to_string(), kcp.kubeconfig.clone(), "--server".to_string(), self.workspace_url(kcp, workspace), "--request-timeout=30s".to_string(), ]; command.extend(args.iter().cloned()); self.runner.kubectl(&command, stdin).await } pub async fn readonly_credential_status( &self, kcp: &KcpInfo, workspace: &str, refresh: bool, ) -> Result { self.ensure_allowed(kcp, "mint a read-only token")?; if refresh { self.credentials .lock() .await .remove(&cache_key(kcp, workspace)); } self.credential(kcp, workspace).await } async fn credential( &self, kcp: &KcpInfo, workspace: &str, ) -> Result { let key = cache_key(kcp, workspace); { let cache = self.credentials.lock().await; if let Some(cached) = cache.get(&key) { if cached.expires_at.saturating_duration_since(Instant::now()) > TOKEN_REFRESH_MARGIN { return Ok(cached.credential.clone()); } } } let admin = self.discovery.admin_config(kcp).await?; let credential = ensure_readonly_credential( &self.runner, kcp, workspace, &admin, self.token_ttl_seconds, ) .await?; let expires_at = Instant::now() + Duration::from_secs(self.token_ttl_seconds); self.credentials.lock().await.insert( key, CachedCredential { credential: credential.clone(), expires_at, }, ); Ok(credential) } fn ensure_allowed(&self, kcp: &KcpInfo, operation: &str) -> Result<(), String> { if self.allow_prod { return Ok(()); } if PROTECTED_HOSTS .iter() .any(|protected| kcp.host.contains(protected)) { return Err(format!( "refusing to {operation} on protected host {:?}; re-run with --allow-prod after confirming the target", kcp.host )); } Ok(()) } } pub fn find_kcp<'a>(kcps: &'a [KcpInfo], needle: &str) -> Option<&'a KcpInfo> { if let Some(exact) = kcps .iter() .find(|kcp| kcp.id == needle || kcp.kubeconfig == needle) { return Some(exact); } let matches = kcps .iter() .filter(|kcp| kcp.id.contains(needle) || kcp.kubeconfig.contains(needle)) .collect::>(); if matches.len() == 1 { matches.first().copied() } else { None } } fn cache_key(kcp: &KcpInfo, workspace: &str) -> String { format!("{}|{workspace}", kcp.id) } #[derive(Serialize)] pub struct DiscoveryReport<'a> { pub kcps: &'a [KcpInfo], pub errors: &'a [String], } #[derive(Serialize)] pub struct CredentialReport { pub kcp: String, pub host: String, pub workspace: String, pub token_kubeconfig: String, pub token_expires_in_seconds: u64, pub note: &'static str, }