johnandersen777.bsky.social/asyncai

Switch repository

Loading…

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<KcpInfo>, Vec<String>)>;

#[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<DiscoveryCache>,
    credentials: Mutex<HashMap<String, CachedCredential>>,
}

#[derive(Clone)]
pub struct KcpServiceHandle(pub Arc<KcpService>);

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<String>,
        kubeconfigs: Vec<String>,
        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<KcpInfo>, Vec<String>) {
        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<Vec<KcpInfo>, 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<KcpInfo, String> {
        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=<id or kubeconfig path>".to_string()),
                many => Err(format!(
                    "several kcps running; pass kcp=<id>, one of: {}",
                    many.iter()
                        .map(|kcp| kcp.id.as_str())
                        .collect::<Vec<_>>()
                        .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::<Vec<_>>()
                    .join(", ")
            )
        })
    }

    pub async fn list_workspaces(
        &self,
        kcp: &KcpInfo,
        workspace: Option<&str>,
        recursive: bool,
    ) -> Result<Vec<WorkspaceInfo>, 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<String, String> {
        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<tokio::process::Child, String> {
        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<String, String> {
        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<ReadonlyCredential, String> {
        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<ReadonlyCredential, String> {
        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::<Vec<_>>();
    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,
}