sol/src/sdk/vault.rs

use std::sync::Arc;

use reqwest::Client as HttpClient;
use serde::Deserialize;
use tokio::sync::Mutex;
use tracing::{debug, info, warn};

const SA_TOKEN_PATH: &str = "/var/run/secrets/kubernetes.io/serviceaccount/token";

/// Client for OpenBao/Vault KV v2 operations, authenticated via Kubernetes auth.
pub struct VaultClient {
    url: String,
    role: String,
    kv_mount: String,
    http: HttpClient,
    token: Mutex<Option<String>>,
}

#[derive(Debug, Deserialize)]
struct AuthResponse {
    auth: AuthData,
}

#[derive(Debug, Deserialize)]
struct AuthData {
    client_token: String,
}

#[derive(Debug, Deserialize)]
struct KvReadResponse {
    data: KvData,
}

#[derive(Debug, Deserialize)]
struct KvData {
    data: serde_json::Value,
}

impl VaultClient {
    pub fn new(url: &str, role: &str, kv_mount: &str) -> Self {
        Self {
            url: url.trim_end_matches('/').to_string(),
            role: role.to_string(),
            kv_mount: kv_mount.to_string(),
            http: HttpClient::new(),
            token: Mutex::new(None),
        }
    }

    /// Authenticate with OpenBao via Kubernetes auth method.
    /// Reads the service account JWT from the mounted token file.
    async fn authenticate(&self) -> Result<String, String> {
        let jwt = tokio::fs::read_to_string(SA_TOKEN_PATH)
            .await
            .map_err(|e| format!("failed to read SA token at {SA_TOKEN_PATH}: {e}"))?;

        let resp = self
            .http
            .post(format!("{}/v1/auth/kubernetes/login", self.url))
            .json(&serde_json::json!({
                "role": self.role,
                "jwt": jwt.trim(),
            }))
            .send()
            .await
            .map_err(|e| format!("vault auth failed: {e}"))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            return Err(format!("vault auth failed (HTTP {status}): {text}"));
        }

        let auth: AuthResponse = resp
            .json()
            .await
            .map_err(|e| format!("failed to parse auth response: {e}"))?;

        info!("Authenticated with OpenBao via Kubernetes auth");
        Ok(auth.auth.client_token)
    }

    /// Get a valid Vault token, authenticating if needed.
    async fn get_token(&self) -> Result<String, String> {
        let mut cached = self.token.lock().await;
        if let Some(ref token) = *cached {
            return Ok(token.clone());
        }
        let token = self.authenticate().await?;
        *cached = Some(token.clone());
        Ok(token)
    }

    /// Clear the cached token (on 403, triggers re-auth on next call).
    async fn clear_token(&self) {
        let mut cached = self.token.lock().await;
        *cached = None;
    }

    /// Read a KV v2 secret. Returns None if not found.
    pub async fn kv_get(&self, path: &str) -> Result<Option<serde_json::Value>, String> {
        let token = self.get_token().await?;
        let url = format!("{}/v1/{}/data/{}", self.url, self.kv_mount, path);

        let resp = self
            .http
            .get(&url)
            .header("X-Vault-Token", &token)
            .send()
            .await
            .map_err(|e| format!("vault read failed: {e}"))?;

        if resp.status().as_u16() == 404 {
            return Ok(None);
        }

        if resp.status().as_u16() == 403 {
            // Token expired — re-auth and retry once
            warn!("Vault token rejected, re-authenticating");
            self.clear_token().await;
            let token = self.get_token().await?;
            let resp = self
                .http
                .get(&url)
                .header("X-Vault-Token", &token)
                .send()
                .await
                .map_err(|e| format!("vault read failed (retry): {e}"))?;

            if resp.status().as_u16() == 404 {
                return Ok(None);
            }
            if !resp.status().is_success() {
                let text = resp.text().await.unwrap_or_default();
                return Err(format!("vault read failed: {text}"));
            }
            let kv: KvReadResponse = resp
                .json()
                .await
                .map_err(|e| format!("vault parse failed: {e}"))?;
            return Ok(Some(kv.data.data));
        }

        if !resp.status().is_success() {
            let text = resp.text().await.unwrap_or_default();
            return Err(format!("vault read failed: {text}"));
        }

        let kv: KvReadResponse = resp
            .json()
            .await
            .map_err(|e| format!("vault parse failed: {e}"))?;

        Ok(Some(kv.data.data))
    }

    /// Write a KV v2 secret (create or update).
    pub async fn kv_put(
        &self,
        path: &str,
        data: serde_json::Value,
    ) -> Result<(), String> {
        let token = self.get_token().await?;
        let url = format!("{}/v1/{}/data/{}", self.url, self.kv_mount, path);

        let body = serde_json::json!({ "data": data });

        let resp = self
            .http
            .post(&url)
            .header("X-Vault-Token", &token)
            .json(&body)
            .send()
            .await
            .map_err(|e| format!("vault write failed: {e}"))?;

        if resp.status().as_u16() == 403 {
            warn!("Vault token rejected on write, re-authenticating");
            self.clear_token().await;
            let token = self.get_token().await?;
            let resp = self
                .http
                .post(&url)
                .header("X-Vault-Token", &token)
                .json(&body)
                .send()
                .await
                .map_err(|e| format!("vault write failed (retry): {e}"))?;

            if !resp.status().is_success() {
                let text = resp.text().await.unwrap_or_default();
                return Err(format!("vault write failed: {text}"));
            }
            return Ok(());
        }

        if !resp.status().is_success() {
            let text = resp.text().await.unwrap_or_default();
            return Err(format!("vault write failed: {text}"));
        }

        debug!(path, "Wrote secret to Vault");
        Ok(())
    }

    /// Delete a KV v2 secret.
    pub async fn kv_delete(&self, path: &str) -> Result<(), String> {
        let token = self.get_token().await?;
        let url = format!(
            "{}/v1/{}/metadata/{}",
            self.url, self.kv_mount, path
        );

        let resp = self
            .http
            .delete(&url)
            .header("X-Vault-Token", &token)
            .send()
            .await
            .map_err(|e| format!("vault delete failed: {e}"))?;

        // 404 = already gone, that's fine
        if resp.status().as_u16() == 404 {
            return Ok(());
        }

        if resp.status().as_u16() == 403 {
            self.clear_token().await;
            let token = self.get_token().await?;
            let resp = self
                .http
                .delete(&url)
                .header("X-Vault-Token", &token)
                .send()
                .await
                .map_err(|e| format!("vault delete failed (retry): {e}"))?;

            if !resp.status().is_success() && resp.status().as_u16() != 404 {
                let text = resp.text().await.unwrap_or_default();
                return Err(format!("vault delete failed: {text}"));
            }
            return Ok(());
        }

        if !resp.status().is_success() {
            let text = resp.text().await.unwrap_or_default();
            return Err(format!("vault delete failed: {text}"));
        }

        Ok(())
    }
}