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;
const DEFAULT_TOKEN_TTL_MS: i64 = 18000000;
pub struct AuthManager {
credentials: Credentials,
token: Arc<RwLock<AccessToken>>,
server_addresses: Vec<ServerAddress>,
}
impl AuthManager {
pub fn new(credentials: Credentials, server_addresses: Vec<ServerAddress>) -> Self {
Self {
credentials,
token: Arc::new(RwLock::new(AccessToken::default())),
server_addresses,
}
}
pub fn is_auth_enabled(&self) -> bool {
self.credentials.is_configured()
}
pub async fn get_token(&self) -> Result<Option<String>> {
if !self.is_auth_enabled() {
return Ok(None);
}
{
let token = self.token.read();
if token.is_valid() {
return Ok(Some(token.token.clone()));
}
}
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(),
})
}
}
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(),
})
}
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() {
if let Some(sig) = self.credentials.generate_signature(&server.address()) {
params.insert("accessKey".to_string(), sig.access_key.clone());
}
}
let response = client
.post(&url)
.form(¶ms)
.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,
})
}
pub async fn get_auth_headers(&self) -> Result<HashMap<String, String>> {
let mut headers = HashMap::new();
if !self.is_auth_enabled() {
return Ok(headers);
}
if let Some(token) = self.get_token().await? {
headers.insert("accessToken".to_string(), token);
}
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());
}
}