ubiquity_core/
auth.rs

1//! Authentication and authorization
2
3use crate::{AuthToken, AuthType, AgentCapability, UbiquityError};
4use chrono::{Utc, Duration};
5use std::collections::HashMap;
6use std::sync::Arc;
7use tokio::sync::RwLock;
8
9/// Authentication service
10pub struct AuthService {
11    auth_type: AuthType,
12    tokens: Arc<RwLock<HashMap<String, AuthToken>>>,
13    api_keys: Arc<RwLock<HashMap<String, String>>>, // api_key -> agent_id
14    secret_key: String,
15}
16
17impl AuthService {
18    pub fn new(auth_type: AuthType) -> Self {
19        Self {
20            auth_type,
21            tokens: Arc::new(RwLock::new(HashMap::new())),
22            api_keys: Arc::new(RwLock::new(HashMap::new())),
23            secret_key: generate_secret_key(),
24        }
25    }
26    
27    /// Generate a new authentication token
28    pub async fn generate_token(
29        &self,
30        agent_id: String,
31        capabilities: AgentCapability,
32        expires_in: Option<Duration>,
33    ) -> Result<AuthToken, UbiquityError> {
34        let token_string = generate_token_string();
35        let expires_at = expires_in.map(|d| Utc::now() + d);
36        
37        let token = AuthToken {
38            token: token_string.clone(),
39            agent_id: agent_id.clone(),
40            capabilities,
41            expires_at,
42        };
43        
44        self.tokens.write().await.insert(token_string, token.clone());
45        
46        Ok(token)
47    }
48    
49    /// Validate a token
50    pub async fn validate_token(&self, token_str: &str) -> Result<AuthToken, UbiquityError> {
51        let tokens = self.tokens.read().await;
52        
53        match tokens.get(token_str) {
54            Some(token) => {
55                // Check expiration
56                if let Some(expires_at) = token.expires_at {
57                    if expires_at < Utc::now() {
58                        return Err(UbiquityError::ConfigError("Token expired".into()));
59                    }
60                }
61                Ok(token.clone())
62            }
63            None => Err(UbiquityError::ConfigError("Invalid token".into())),
64        }
65    }
66    
67    /// Register an API key
68    pub async fn register_api_key(&self, api_key: String, agent_id: String) -> Result<(), UbiquityError> {
69        self.api_keys.write().await.insert(api_key, agent_id);
70        Ok(())
71    }
72    
73    /// Validate an API key
74    pub async fn validate_api_key(&self, api_key: &str) -> Result<String, UbiquityError> {
75        self.api_keys
76            .read()
77            .await
78            .get(api_key)
79            .cloned()
80            .ok_or_else(|| UbiquityError::ConfigError("Invalid API key".into()))
81    }
82    
83    /// Authenticate based on configured method
84    pub async fn authenticate(&self, credentials: &str) -> Result<String, UbiquityError> {
85        match self.auth_type {
86            AuthType::None => Ok("anonymous".to_string()),
87            AuthType::Token => {
88                let token = self.validate_token(credentials).await?;
89                Ok(token.agent_id)
90            }
91            AuthType::ApiKey => {
92                self.validate_api_key(credentials).await
93            }
94            AuthType::Certificate => {
95                // TODO: Implement certificate validation
96                Err(UbiquityError::ConfigError("Certificate auth not implemented".into()))
97            }
98        }
99    }
100    
101    /// Check if an agent has a specific capability
102    pub async fn check_capability(
103        &self,
104        token_str: &str,
105        check: impl FnOnce(&AgentCapability) -> bool,
106    ) -> Result<bool, UbiquityError> {
107        let token = self.validate_token(token_str).await?;
108        Ok(check(&token.capabilities))
109    }
110    
111    /// Revoke a token
112    pub async fn revoke_token(&self, token_str: &str) -> Result<(), UbiquityError> {
113        self.tokens.write().await.remove(token_str);
114        Ok(())
115    }
116    
117    /// Clean up expired tokens
118    pub async fn cleanup_expired_tokens(&self) {
119        let now = Utc::now();
120        let mut tokens = self.tokens.write().await;
121        
122        tokens.retain(|_, token| {
123            token.expires_at.map(|exp| exp > now).unwrap_or(true)
124        });
125    }
126}
127
128/// Generate a secure random token string
129fn generate_token_string() -> String {
130    use rand::Rng;
131    const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
132    const TOKEN_LEN: usize = 32;
133    
134    let mut rng = rand::thread_rng();
135    (0..TOKEN_LEN)
136        .map(|_| {
137            let idx = rng.gen_range(0..CHARSET.len());
138            CHARSET[idx] as char
139        })
140        .collect()
141}
142
143/// Generate a secret key for signing
144fn generate_secret_key() -> String {
145    generate_token_string()
146}
147
148/// Authorization guard
149pub struct AuthGuard {
150    required_capabilities: Vec<Box<dyn Fn(&AgentCapability) -> bool + Send + Sync>>,
151}
152
153impl AuthGuard {
154    pub fn new() -> Self {
155        Self {
156            required_capabilities: Vec::new(),
157        }
158    }
159    
160    /// Require write code capability
161    pub fn require_write_code(mut self) -> Self {
162        self.required_capabilities.push(Box::new(|cap| cap.can_write_code));
163        self
164    }
165    
166    /// Require execute commands capability
167    pub fn require_execute_commands(mut self) -> Self {
168        self.required_capabilities.push(Box::new(|cap| cap.can_execute_commands));
169        self
170    }
171    
172    /// Require coordination capability
173    pub fn require_coordination(mut self) -> Self {
174        self.required_capabilities.push(Box::new(|cap| cap.can_coordinate));
175        self
176    }
177    
178    /// Check if token meets all requirements
179    pub async fn check(&self, auth_service: &AuthService, token: &str) -> Result<(), UbiquityError> {
180        let token_data = auth_service.validate_token(token).await?;
181        
182        for check in &self.required_capabilities {
183            if !check(&token_data.capabilities) {
184                return Err(UbiquityError::ConfigError("Insufficient capabilities".into()));
185            }
186        }
187        
188        Ok(())
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    
196    #[tokio::test]
197    async fn test_token_generation_and_validation() {
198        let auth = AuthService::new(AuthType::Token);
199        let capabilities = AgentCapability::default();
200        
201        let token = auth.generate_token(
202            "test-agent".to_string(),
203            capabilities,
204            Some(Duration::hours(1)),
205        ).await.unwrap();
206        
207        let validated = auth.validate_token(&token.token).await.unwrap();
208        assert_eq!(validated.agent_id, "test-agent");
209    }
210    
211    #[tokio::test]
212    async fn test_expired_token() {
213        let auth = AuthService::new(AuthType::Token);
214        let capabilities = AgentCapability::default();
215        
216        let token = auth.generate_token(
217            "test-agent".to_string(),
218            capabilities,
219            Some(Duration::seconds(-1)), // Already expired
220        ).await.unwrap();
221        
222        let result = auth.validate_token(&token.token).await;
223        assert!(result.is_err());
224    }
225    
226    #[tokio::test]
227    async fn test_api_key() {
228        let auth = AuthService::new(AuthType::ApiKey);
229        
230        auth.register_api_key("test-key".to_string(), "test-agent".to_string())
231            .await
232            .unwrap();
233        
234        let agent_id = auth.validate_api_key("test-key").await.unwrap();
235        assert_eq!(agent_id, "test-agent");
236    }
237    
238    #[tokio::test]
239    async fn test_auth_guard() {
240        let auth = AuthService::new(AuthType::Token);
241        let mut capabilities = AgentCapability::default();
242        capabilities.can_coordinate = true;
243        
244        let token = auth.generate_token(
245            "test-agent".to_string(),
246            capabilities,
247            None,
248        ).await.unwrap();
249        
250        let guard = AuthGuard::new()
251            .require_write_code()
252            .require_coordination();
253        
254        let result = guard.check(&auth, &token.token).await;
255        assert!(result.is_ok());
256    }
257}
ubiquity_core/auth.rs

ubiquity_core/
auth.rs
