batata-client 0.0.2

Rust client for Batata/Nacos service discovery and configuration management
Documentation
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use parking_lot::RwLock;
use tracing::{debug, warn};

use crate::auth::{AccessToken, Credentials};
use crate::error::{BatataError, Result};
use crate::remote::ServerAddress;

/// Default token TTL in milliseconds (5 hours, matching Nacos default)
const DEFAULT_TOKEN_TTL_MS: i64 = 18000000;

/// Auth manager for handling authentication
pub struct AuthManager {
    credentials: Credentials,
    token: Arc<RwLock<AccessToken>>,
    server_addresses: Vec<ServerAddress>,
}

impl AuthManager {
    /// Create a new auth manager
    pub fn new(credentials: Credentials, server_addresses: Vec<ServerAddress>) -> Self {
        Self {
            credentials,
            token: Arc::new(RwLock::new(AccessToken::default())),
            server_addresses,
        }
    }

    /// Check if authentication is required
    pub fn is_auth_enabled(&self) -> bool {
        self.credentials.is_configured()
    }

    /// Get current valid token, refreshing if necessary
    pub async fn get_token(&self) -> Result<Option<String>> {
        if !self.is_auth_enabled() {
            return Ok(None);
        }

        // Check if token is still valid
        {
            let token = self.token.read();
            if token.is_valid() {
                return Ok(Some(token.token.clone()));
            }
        }

        // Refresh token
        self.refresh_token().await?;

        let token = self.token.read();
        if token.is_valid() {
            Ok(Some(token.token.clone()))
        } else {
            Err(BatataError::AuthError {
                message: "Failed to obtain valid token".to_string(),
            })
        }
    }

    /// Refresh the access token
    pub async fn refresh_token(&self) -> Result<()> {
        if !self.credentials.is_configured() {
            return Ok(());
        }

        for server in &self.server_addresses {
            match self.login_to_server(server).await {
                Ok(token) => {
                    *self.token.write() = token;
                    debug!("Token refreshed successfully from {}", server.address());
                    return Ok(());
                }
                Err(e) => {
                    warn!("Failed to login to {}: {}", server.address(), e);
                    continue;
                }
            }
        }

        Err(BatataError::AuthError {
            message: "Failed to login to any server".to_string(),
        })
    }

    /// Login to a specific server
    async fn login_to_server(&self, server: &ServerAddress) -> Result<AccessToken> {
        let url = format!(
            "http://{}:{}/nacos/v1/auth/login",
            server.host(),
            server.port()
        );

        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(5))
            .build()
            .map_err(|e| BatataError::connection_error(format!("Failed to create HTTP client: {}", e)))?;

        let mut params = HashMap::new();

        if let (Some(username), Some(password)) = (&self.credentials.username, &self.credentials.password) {
            params.insert("username".to_string(), username.clone());
            params.insert("password".to_string(), password.clone());
        } else if self.credentials.has_ak_sk_auth() {
            // For AK/SK auth, we might need different endpoint or headers
            // This is a simplified version
            if let Some(sig) = self.credentials.generate_signature(&server.address()) {
                params.insert("accessKey".to_string(), sig.access_key.clone());
                // Note: actual implementation may vary based on Nacos version
            }
        }

        let response = client
            .post(&url)
            .form(&params)
            .send()
            .await
            .map_err(|e| BatataError::connection_error(format!("Login request failed: {}", e)))?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(BatataError::AuthError {
                message: format!("Login failed with status {}: {}", status, body),
            });
        }

        let body: serde_json::Value = response
            .json()
            .await
            .map_err(|e| BatataError::AuthError {
                message: format!("Failed to parse login response: {}", e),
            })?;

        let token = body["accessToken"]
            .as_str()
            .ok_or_else(|| BatataError::AuthError {
                message: "No accessToken in response".to_string(),
            })?
            .to_string();

        let ttl = body["tokenTtl"]
            .as_i64()
            .unwrap_or(DEFAULT_TOKEN_TTL_MS / 1000) * 1000;

        let global_admin = body["globalAdmin"].as_bool().unwrap_or(false);

        Ok(AccessToken {
            token,
            expire_time: crate::common::current_time_millis() + ttl,
            global_admin,
        })
    }

    /// Get auth headers to include in requests
    pub async fn get_auth_headers(&self) -> Result<HashMap<String, String>> {
        let mut headers = HashMap::new();

        if !self.is_auth_enabled() {
            return Ok(headers);
        }

        // Get token
        if let Some(token) = self.get_token().await? {
            headers.insert("accessToken".to_string(), token);
        }

        // Add AK/SK signature if configured
        if self.credentials.has_ak_sk_auth()
            && let Some(sig) = self.credentials.generate_signature("+")
        {
            headers.insert("ak".to_string(), sig.access_key);
            headers.insert("data".to_string(), sig.timestamp);
            headers.insert("signature".to_string(), sig.signature);
        }

        Ok(headers)
    }
}

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

    #[test]
    fn test_auth_manager_no_auth() {
        let manager = AuthManager::new(Credentials::new(), vec![]);
        assert!(!manager.is_auth_enabled());
    }

    #[test]
    fn test_auth_manager_with_credentials() {
        let creds = Credentials::with_username_password("admin", "password");
        let servers = vec![ServerAddress::new("localhost", 8848)];
        let manager = AuthManager::new(creds, servers);
        assert!(manager.is_auth_enabled());
    }
}