harmony/harmony-k8s/src/client.rs

use std::sync::Arc;

use kube::config::{KubeConfigOptions, Kubeconfig};
use kube::{Client, Config, Discovery, Error};
use log::{error, info};
use serde::Serialize;
use tokio::sync::{OnceCell, RwLock};

use crate::types::KubernetesDistribution;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ClusterConnection {
    pub name: String,
    pub server: String,
    pub tls_server_name: Option<String>,
    pub proxy_url: Option<String>,
    pub tls_verified: bool,
}

impl ClusterConnection {
    fn from_kubeconfig(
        kubeconfig: &Kubeconfig,
        options: &KubeConfigOptions,
    ) -> Result<Self, String> {
        let context_name = options
            .context
            .as_ref()
            .or(kubeconfig.current_context.as_ref())
            .ok_or("kubeconfig has no selected context")?;
        let context = kubeconfig
            .contexts
            .iter()
            .find(|context| &context.name == context_name)
            .and_then(|context| context.context.as_ref())
            .ok_or_else(|| format!("kubeconfig context '{context_name}' not found"))?;
        let cluster_name = options.cluster.as_ref().unwrap_or(&context.cluster);
        let cluster = kubeconfig
            .clusters
            .iter()
            .find(|cluster| &cluster.name == cluster_name)
            .and_then(|cluster| cluster.cluster.as_ref())
            .ok_or_else(|| format!("kubeconfig cluster '{cluster_name}' not found"))?;
        let server = safe_endpoint(
            cluster
                .server
                .as_deref()
                .ok_or_else(|| format!("kubeconfig cluster '{cluster_name}' has no server"))?,
            "cluster server",
            true,
        )?;
        let proxy_url = cluster
            .proxy_url
            .as_deref()
            .map(|url| safe_endpoint(url, "cluster proxy", false))
            .transpose()?;
        Ok(Self {
            name: cluster_name.clone(),
            server,
            tls_server_name: cluster.tls_server_name.clone(),
            proxy_url,
            tls_verified: cluster.insecure_skip_tls_verify != Some(true),
        })
    }

    fn from_config(name: String, config: &Config) -> Result<Self, String> {
        let server = safe_endpoint(&config.cluster_url.to_string(), "cluster server", true)?;
        let proxy_url = config
            .proxy_url
            .as_ref()
            .map(|url| safe_endpoint(&url.to_string(), "cluster proxy", false))
            .transpose()?;
        Ok(Self {
            name,
            server,
            tls_server_name: config.tls_server_name.clone(),
            proxy_url,
            tls_verified: !config.accept_invalid_certs,
        })
    }
}

fn safe_endpoint(endpoint: &str, name: &str, require_https: bool) -> Result<String, String> {
    let url = url::Url::parse(endpoint).map_err(|error| format!("invalid {name}: {error}"))?;
    if require_https && url.scheme() != "https" {
        return Err(format!("{name} must use HTTPS"));
    }
    if !url.username().is_empty()
        || url.password().is_some()
        || url.query().is_some()
        || url.fragment().is_some()
    {
        return Err(format!(
            "{name} must not contain credentials, a query, or a fragment"
        ));
    }
    Ok(endpoint.to_string())
}

// TODO not cool, should use a proper configuration mechanism
// cli arg, env var, config file
fn read_dry_run_from_env() -> bool {
    std::env::var("DRY_RUN")
        .map(|v| v == "true" || v == "1")
        .unwrap_or(false)
}

#[derive(Clone)]
pub struct K8sClient {
    pub(crate) client: Client,
    /// When `true` no mutation is sent to the API server; diffs are printed
    /// to stdout instead. Initialised from the `DRY_RUN` environment variable.
    pub(crate) dry_run: bool,
    pub(crate) k8s_distribution: Arc<OnceCell<KubernetesDistribution>>,
    /// API discovery cache. Wrapped in `RwLock` so it can be invalidated
    /// after installing CRDs or operators that register new API groups.
    pub(crate) discovery: Arc<RwLock<Option<Arc<Discovery>>>>,
    cluster_connection: Option<ClusterConnection>,
}

impl Serialize for K8sClient {
    fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        todo!("K8sClient serialization is not meaningful; remove this impl if unused")
    }
}

impl std::fmt::Debug for K8sClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "K8sClient {{ namespace: {}, dry_run: {} }}",
            self.client.default_namespace(),
            self.dry_run,
        ))
    }
}

impl K8sClient {
    pub fn validate_kubeconfig_context(
        path: &str,
        context: String,
    ) -> Result<ClusterConnection, String> {
        let kubeconfig = Kubeconfig::read_from(path)
            .map_err(|error| format!("failed to load kubeconfig from {path}: {error}"))?;
        ClusterConnection::from_kubeconfig(
            &kubeconfig,
            &KubeConfigOptions {
                context: Some(context),
                ..Default::default()
            },
        )
    }

    /// Create a client, reading `DRY_RUN` from the environment.
    pub fn new(client: Client) -> Self {
        Self {
            dry_run: read_dry_run_from_env(),
            client,
            k8s_distribution: Arc::new(OnceCell::new()),
            discovery: Arc::new(RwLock::new(None)),
            cluster_connection: None,
        }
    }

    /// Create a client that always operates in dry-run mode, regardless of
    /// the environment variable.
    pub fn new_dry_run(client: Client) -> Self {
        Self {
            dry_run: true,
            ..Self::new(client)
        }
    }

    /// Returns `true` if this client is operating in dry-run mode.
    pub fn is_dry_run(&self) -> bool {
        self.dry_run
    }

    pub fn cluster_connection(&self) -> Option<&ClusterConnection> {
        self.cluster_connection.as_ref()
    }

    pub async fn try_default() -> Result<Self, Error> {
        Ok(Self::new(Client::try_default().await?))
    }

    pub async fn from_kubeconfig(path: &str) -> Option<Self> {
        Self::from_kubeconfig_with_opts(path, &KubeConfigOptions::default()).await
    }

    pub async fn from_kubeconfig_with_context(path: &str, context: Option<String>) -> Option<Self> {
        let opts = KubeConfigOptions {
            context,
            ..Default::default()
        };
        Self::from_kubeconfig_with_opts(path, &opts).await
    }

    pub async fn from_kubeconfig_with_opts(path: &str, opts: &KubeConfigOptions) -> Option<Self> {
        let k = match Kubeconfig::read_from(path) {
            Ok(k) => k,
            Err(e) => {
                error!("Failed to load kubeconfig from {path}: {e}");
                return None;
            }
        };
        let context = opts
            .context
            .clone()
            .or_else(|| k.current_context.clone())
            .unwrap_or_default();
        let cluster_name = ClusterConnection::from_kubeconfig(&k, opts)
            .map(|connection| connection.name)
            .map_err(|error| error!("Tenant credentials unavailable for {path}: {error}"))
            .ok();
        let config = match Config::from_custom_kubeconfig(k, opts).await {
            Ok(config) => config,
            Err(error) => {
                error!("Failed to load kubeconfig from {path}: {error}");
                return None;
            }
        };
        let connection = cluster_name.and_then(|name| {
            ClusterConnection::from_config(name, &config)
                .map_err(|error| error!("Tenant credentials unavailable for {path}: {error}"))
                .ok()
        });
        if let Some(connection) = &connection {
            info!(
                "Loaded Kubernetes context '{context}' for cluster '{}' at '{}'",
                connection.name, connection.server
            );
        }
        let client = match Client::try_from(config) {
            Ok(client) => client,
            Err(error) => {
                error!("Failed to create Kubernetes client from {path}: {error}");
                return None;
            }
        };
        let mut client = Self::new(client);
        client.cluster_connection = connection;
        Some(client)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn connection_uses_selected_context_without_authentication() {
        let kubeconfig: Kubeconfig = serde_yaml::from_str(
            r#"
current-context: admin
contexts:
  - name: admin
    context:
      cluster: shared
      user: cluster-admin
clusters:
  - name: shared
    cluster:
      server: https://api.example.com:6443
      certificate-authority-data: Y2E=
users:
  - name: cluster-admin
    user:
      token: super-secret
"#,
        )
        .unwrap();

        assert_eq!(
            ClusterConnection::from_kubeconfig(&kubeconfig, &KubeConfigOptions::default()).unwrap(),
            ClusterConnection {
                name: "shared".to_string(),
                server: "https://api.example.com:6443".to_string(),
                tls_server_name: None,
                proxy_url: None,
                tls_verified: true,
            }
        );
    }

    #[tokio::test]
    async fn resolved_connection_uses_explicit_context_in_multi_cluster_config() {
        let kubeconfig: Kubeconfig = serde_yaml::from_str(
            r#"
current-context: cluster-a-admin
contexts:
  - name: cluster-a-admin
    context: { cluster: cluster-a, user: cluster-a-admin }
  - name: cluster-b-admin
    context: { cluster: cluster-b, user: cluster-b-admin }
clusters:
  - name: cluster-a
    cluster: { server: "https://192.0.2.10:6443" }
  - name: cluster-b
    cluster:
      server: https://api.cluster-b.example.com:6443
      tls-server-name: api.cluster-b.example.com
users:
  - name: cluster-a-admin
    user: { token: cluster-a-secret }
  - name: cluster-b-admin
    user: { token: cluster-b-secret }
"#,
        )
        .unwrap();
        let options = KubeConfigOptions {
            context: Some("cluster-b-admin".into()),
            ..Default::default()
        };
        let name = ClusterConnection::from_kubeconfig(&kubeconfig, &options)
            .unwrap()
            .name;
        let config = Config::from_custom_kubeconfig(kubeconfig, &options)
            .await
            .unwrap();

        assert_eq!(
            ClusterConnection::from_config(name, &config).unwrap(),
            ClusterConnection {
                name: "cluster-b".into(),
                server: "https://api.cluster-b.example.com:6443/".into(),
                tls_server_name: Some("api.cluster-b.example.com".into()),
                proxy_url: None,
                tls_verified: true,
            }
        );
    }

    #[test]
    fn connection_records_insecure_tls() {
        let kubeconfig: Kubeconfig = serde_yaml::from_str(
            r#"
current-context: admin
contexts:
  - name: admin
    context: { cluster: shared }
clusters:
  - name: shared
    cluster:
      server: https://api.example.com
      insecure-skip-tls-verify: true
"#,
        )
        .unwrap();

        assert!(
            !ClusterConnection::from_kubeconfig(&kubeconfig, &KubeConfigOptions::default())
                .unwrap()
                .tls_verified
        );
    }

    #[test]
    fn connection_rejects_credentials_in_server_url() {
        let kubeconfig: Kubeconfig = serde_yaml::from_str(
            r#"
current-context: admin
contexts:
  - name: admin
    context: { cluster: shared }
clusters:
  - name: shared
    cluster: { server: "https://admin:secret@api.example.com" }
"#,
        )
        .unwrap();

        assert!(
            ClusterConnection::from_kubeconfig(&kubeconfig, &KubeConfigOptions::default()).is_err()
        );
    }

    #[test]
    fn connection_rejects_plain_http_server() {
        let kubeconfig: Kubeconfig = serde_yaml::from_str(
            r#"
current-context: admin
contexts:
  - name: admin
    context: { cluster: shared }
clusters:
  - name: shared
    cluster: { server: "http://api.example.com" }
"#,
        )
        .unwrap();

        assert!(
            ClusterConnection::from_kubeconfig(&kubeconfig, &KubeConfigOptions::default()).is_err()
        );
    }
}