heliosdb_proxy/auth/
mod.rs

1//! Authentication Proxy Module
2//!
3//! Provides comprehensive authentication and authorization for HeliosProxy.
4//!
5//! # Features
6//!
7//! - **JWT Validation**: JWKS-based JWT token validation
8//! - **OAuth Introspection**: RFC 7662 token introspection
9//! - **API Key Management**: Generate, validate, and revoke API keys
10//! - **Role Mapping**: Map identities to database roles
11//! - **Credential Providers**: Vault, AWS Secrets Manager, environment
12//! - **Session Management**: Token-based session handling
13//!
14//! # Architecture
15//!
16//! ```text
17//!                    ┌─────────────────────────────────────┐
18//!                    │      AuthenticationHandler          │
19//!                    │  (Main entry point for auth)        │
20//!                    └─────────────┬───────────────────────┘
21//!                                  │
22//!          ┌───────────────────────┼───────────────────────┐
23//!          │                       │                       │
24//!   ┌──────▼──────┐        ┌───────▼───────┐       ┌───────▼───────┐
25//!   │ JwtValidator│        │  OAuthClient  │       │ ApiKeyManager │
26//!   │  (JWKS)     │        │(Introspection)│       │ (Key mgmt)    │
27//!   └─────────────┘        └───────────────┘       └───────────────┘
28//!          │                       │                       │
29//!          └───────────────────────┼───────────────────────┘
30//!                                  │
31//!                    ┌─────────────▼───────────────────────┐
32//!                    │          Identity                    │
33//!                    │  (Unified user representation)       │
34//!                    └─────────────┬───────────────────────┘
35//!                                  │
36//!                    ┌─────────────▼───────────────────────┐
37//!                    │          RoleMapper                  │
38//!                    │  (Identity → Database Roles)         │
39//!                    └─────────────────────────────────────┘
40//! ```
41//!
42//! # Example
43//!
44//! ```rust,ignore
45//! use heliosdb::proxy::auth::{
46//!     AuthenticationHandler, AuthRequest, JwtConfig,
47//!     RoleMapper, SessionManager,
48//! };
49//!
50//! // Create authentication handler
51//! let handler = AuthenticationHandler::builder()
52//!     .enabled(true)
53//!     .with_jwt(JwtConfig::new("https://auth.example.com/.well-known/jwks.json"))
54//!     .with_api_keys(ApiKeyConfig::default())
55//!     .default_role("db_user")
56//!     .build();
57//!
58//! // Authenticate a request
59//! let request = AuthRequest::new()
60//!     .with_header("Authorization", "Bearer eyJ...");
61//!
62//! let result = handler.authenticate(&request).await?;
63//! println!("Authenticated: {}", result.identity.user_id);
64//!
65//! // Map to database roles
66//! let mapper = RoleMapper::builder()
67//!     .group_role("admins", "db_admin")
68//!     .group_role("developers", "db_developer")
69//!     .default_role("db_readonly")
70//!     .build();
71//!
72//! let roles = mapper.map_roles(&result.identity);
73//! ```
74//!
75//! # AI/Agent Authentication
76//!
77//! HeliosProxy supports special authentication patterns for AI agents:
78//!
79//! - **Agent Tokens**: Short-lived tokens with conversation scope
80//! - **Tool Authorization**: Role-based tool access control
81//! - **Quota Management**: Per-agent resource quotas
82//!
83//! ```rust,ignore
84//! use heliosdb::proxy::auth::{AgentIdentity, AgentQuota};
85//!
86//! let agent_identity = AgentIdentity {
87//!     agent_id: "claude-code".to_string(),
88//!     parent_user_id: "user123".to_string(),
89//!     conversation_id: Some("conv_abc".to_string()),
90//!     allowed_tools: vec!["query", "insert".to_string()],
91//!     quota: AgentQuota::default(),
92//! };
93//! ```
94
95pub mod api_keys;
96pub mod config;
97pub mod credentials;
98pub mod handler;
99pub mod jwt;
100pub mod oauth;
101pub mod role_mapper;
102pub mod session;
103
104// Re-export main types
105pub use config::{
106    AgentIdentity, AgentQuota, ApiKeyConfig, AuthConfig, AuthMethod, AuthRateLimitConfig,
107    CredentialConfig, Identity, JwtClaims, JwtConfig, LdapConfig, OAuthConfig,
108    RoleMappingCondition, RoleMappingRule, SessionConfig,
109};
110
111pub use jwt::{Jwk, Jwks, JwtError, JwtHeader, JwtValidator, TokenCache};
112
113pub use handler::{
114    AuthError, AuthRequest, AuthResult, AuthenticationHandler, AuthenticationHandlerBuilder,
115    CacheStats,
116};
117
118pub use oauth::{IntrospectionResponse, OAuthClient, OAuthError, TokenExchange, TokenResponse};
119
120pub use api_keys::{ApiKey, ApiKeyBuilder, ApiKeyError, ApiKeyManager, ApiKeyStats};
121
122pub use role_mapper::{
123    AuthorizationContext, Operation, PermissionSet, RoleMapper, RoleMapperBuilder,
124};
125
126pub use credentials::{
127    AwsSecretsManagerProvider, CredentialError, CredentialManager, CredentialManagerBuilder,
128    CredentialProvider, CredentialSource, DatabaseCredential, EnvironmentCredentialProvider,
129    StaticCredentialProvider, VaultCredentialProvider,
130};
131
132pub use session::{
133    CookieOptions, SameSite, Session, SessionError, SessionManager, SessionManagerBuilder,
134    SessionStats,
135};
136
137/// Authentication proxy facade
138///
139/// High-level facade that combines authentication, authorization, and
140/// session management into a single interface.
141pub struct AuthProxy {
142    /// Authentication handler
143    handler: AuthenticationHandler,
144
145    /// Role mapper
146    role_mapper: RoleMapper,
147
148    /// Session manager
149    session_manager: SessionManager,
150
151    /// Credential manager
152    credential_manager: Option<CredentialManager>,
153}
154
155impl AuthProxy {
156    /// Create a new auth proxy
157    pub fn new(
158        handler: AuthenticationHandler,
159        role_mapper: RoleMapper,
160        session_manager: SessionManager,
161    ) -> Self {
162        Self {
163            handler,
164            role_mapper,
165            session_manager,
166            credential_manager: None,
167        }
168    }
169
170    /// Create a builder
171    pub fn builder() -> AuthProxyBuilder {
172        AuthProxyBuilder::new()
173    }
174
175    /// Authenticate a request
176    pub async fn authenticate(&self, request: &AuthRequest) -> Result<AuthResult, AuthError> {
177        self.handler.authenticate(request).await
178    }
179
180    /// Authenticate and create session
181    pub async fn authenticate_and_create_session(
182        &self,
183        request: &AuthRequest,
184    ) -> Result<(AuthResult, Session), AuthError> {
185        let result = self.handler.authenticate(request).await?;
186
187        let session = self
188            .session_manager
189            .create_session(
190                result.identity.clone(),
191                request.client_ip,
192                request.headers.get("user-agent").cloned(),
193            )
194            .map_err(|e| AuthError::Session(e.to_string()))?;
195
196        Ok((result, session))
197    }
198
199    /// Validate session token
200    pub fn validate_session(&self, token: &str) -> Result<Identity, AuthError> {
201        self.session_manager
202            .validate_token(token)
203            .map_err(|e| AuthError::Session(e.to_string()))
204    }
205
206    /// Map identity to database roles
207    pub fn map_roles(&self, identity: &Identity) -> Vec<String> {
208        self.role_mapper.map_roles(identity)
209    }
210
211    /// Get database credentials for an identity
212    pub fn get_credentials(&self, key: &str) -> Result<DatabaseCredential, AuthError> {
213        self.credential_manager
214            .as_ref()
215            .ok_or_else(|| {
216                AuthError::Configuration("Credential manager not configured".to_string())
217            })?
218            .get_credential(key)
219            .map_err(|e| AuthError::Configuration(e.to_string()))
220    }
221
222    /// Invalidate session
223    pub fn invalidate_session(&self, token: &str) -> Result<(), AuthError> {
224        self.session_manager
225            .invalidate_session(token)
226            .map_err(|e| AuthError::Session(e.to_string()))
227    }
228
229    /// Get authentication handler
230    pub fn handler(&self) -> &AuthenticationHandler {
231        &self.handler
232    }
233
234    /// Get role mapper
235    pub fn role_mapper(&self) -> &RoleMapper {
236        &self.role_mapper
237    }
238
239    /// Get session manager
240    pub fn session_manager(&self) -> &SessionManager {
241        &self.session_manager
242    }
243}
244
245/// Auth proxy builder
246pub struct AuthProxyBuilder {
247    handler: Option<AuthenticationHandler>,
248    role_mapper: Option<RoleMapper>,
249    session_manager: Option<SessionManager>,
250    credential_manager: Option<CredentialManager>,
251}
252
253impl AuthProxyBuilder {
254    /// Create a new builder
255    pub fn new() -> Self {
256        Self {
257            handler: None,
258            role_mapper: None,
259            session_manager: None,
260            credential_manager: None,
261        }
262    }
263
264    /// Set authentication handler
265    pub fn handler(mut self, handler: AuthenticationHandler) -> Self {
266        self.handler = Some(handler);
267        self
268    }
269
270    /// Set role mapper
271    pub fn role_mapper(mut self, mapper: RoleMapper) -> Self {
272        self.role_mapper = Some(mapper);
273        self
274    }
275
276    /// Set session manager
277    pub fn session_manager(mut self, manager: SessionManager) -> Self {
278        self.session_manager = Some(manager);
279        self
280    }
281
282    /// Set credential manager
283    pub fn credential_manager(mut self, manager: CredentialManager) -> Self {
284        self.credential_manager = Some(manager);
285        self
286    }
287
288    /// Build the auth proxy
289    pub fn build(self) -> AuthProxy {
290        let handler = self
291            .handler
292            .unwrap_or_else(|| AuthenticationHandler::builder().enabled(false).build());
293
294        let role_mapper = self.role_mapper.unwrap_or_default();
295
296        let session_manager = self
297            .session_manager
298            .unwrap_or_else(|| SessionManager::new(SessionConfig::default()));
299
300        let mut proxy = AuthProxy::new(handler, role_mapper, session_manager);
301        proxy.credential_manager = self.credential_manager;
302        proxy
303    }
304}
305
306impl Default for AuthProxyBuilder {
307    fn default() -> Self {
308        Self::new()
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use std::collections::HashMap;
316
317    fn test_identity() -> Identity {
318        Identity {
319            user_id: "testuser".to_string(),
320            name: Some("Test User".to_string()),
321            email: Some("test@example.com".to_string()),
322            roles: vec!["user".to_string()],
323            groups: vec!["developers".to_string()],
324            tenant_id: None,
325            claims: HashMap::new(),
326            auth_method: "test".to_string(),
327            authenticated_at: chrono::Utc::now(),
328        }
329    }
330
331    #[test]
332    fn test_auth_proxy_builder() {
333        let proxy = AuthProxy::builder()
334            .handler(AuthenticationHandler::builder().enabled(false).build())
335            .role_mapper(RoleMapper::builder().default_role("user").build())
336            .session_manager(SessionManager::new(SessionConfig::default()))
337            .build();
338
339        assert!(!proxy.handler().is_enabled());
340    }
341
342    #[test]
343    fn test_auth_proxy_map_roles() {
344        let proxy = AuthProxy::builder()
345            .role_mapper(
346                RoleMapper::builder()
347                    .group_role("developers", "db_developer")
348                    .build(),
349            )
350            .build();
351
352        let roles = proxy.map_roles(&test_identity());
353        assert!(roles.contains(&"db_developer".to_string()));
354    }
355
356    #[tokio::test]
357    async fn test_auth_proxy_disabled() {
358        let proxy = AuthProxy::builder().build();
359
360        let request = AuthRequest::new();
361        let result = proxy.authenticate(&request).await.unwrap();
362
363        assert_eq!(result.identity.auth_method, "anonymous");
364    }
365
366    #[test]
367    fn test_session_integration() {
368        let proxy = AuthProxy::builder()
369            .session_manager(SessionManager::builder().max_sessions_per_user(5).build())
370            .build();
371
372        // Create session directly via session manager
373        let session = proxy
374            .session_manager()
375            .create_session(test_identity(), None, None)
376            .unwrap();
377
378        // Validate via proxy
379        let identity = proxy.validate_session(&session.token).unwrap();
380        assert_eq!(identity.user_id, "testuser");
381
382        // Invalidate via proxy
383        proxy.invalidate_session(&session.token).unwrap();
384        assert!(proxy.validate_session(&session.token).is_err());
385    }
386}
heliosdb_proxy/auth/mod.rs

heliosdb_proxy/auth/
mod.rs
