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 config;
96pub mod jwt;
97pub mod handler;
98pub mod oauth;
99pub mod api_keys;
100pub mod role_mapper;
101pub mod credentials;
102pub mod session;
103
104// Re-export main types
105pub use config::{
106    AuthConfig, AuthMethod, Identity, AgentIdentity, AgentQuota,
107    JwtConfig, JwtClaims, OAuthConfig, LdapConfig, ApiKeyConfig,
108    RoleMappingRule, RoleMappingCondition, CredentialConfig, SessionConfig,
109    AuthRateLimitConfig,
110};
111
112pub use jwt::{JwtValidator, JwtError, JwtHeader, Jwks, Jwk, TokenCache};
113
114pub use handler::{
115    AuthenticationHandler, AuthenticationHandlerBuilder,
116    AuthRequest, AuthResult, AuthError, CacheStats,
117};
118
119pub use oauth::{
120    OAuthClient, OAuthError, IntrospectionResponse,
121    TokenExchange, TokenResponse,
122};
123
124pub use api_keys::{
125    ApiKeyManager, ApiKey, ApiKeyError, ApiKeyStats, ApiKeyBuilder,
126};
127
128pub use role_mapper::{
129    RoleMapper, RoleMapperBuilder, PermissionSet, Operation,
130    AuthorizationContext,
131};
132
133pub use credentials::{
134    CredentialManager, CredentialManagerBuilder, CredentialProvider,
135    DatabaseCredential, CredentialSource, CredentialError,
136    StaticCredentialProvider, EnvironmentCredentialProvider,
137    VaultCredentialProvider, AwsSecretsManagerProvider,
138};
139
140pub use session::{
141    SessionManager, SessionManagerBuilder, Session, SessionError,
142    SessionStats, CookieOptions, SameSite,
143};
144
145/// Authentication proxy facade
146///
147/// High-level facade that combines authentication, authorization, and
148/// session management into a single interface.
149pub struct AuthProxy {
150    /// Authentication handler
151    handler: AuthenticationHandler,
152
153    /// Role mapper
154    role_mapper: RoleMapper,
155
156    /// Session manager
157    session_manager: SessionManager,
158
159    /// Credential manager
160    credential_manager: Option<CredentialManager>,
161}
162
163impl AuthProxy {
164    /// Create a new auth proxy
165    pub fn new(
166        handler: AuthenticationHandler,
167        role_mapper: RoleMapper,
168        session_manager: SessionManager,
169    ) -> Self {
170        Self {
171            handler,
172            role_mapper,
173            session_manager,
174            credential_manager: None,
175        }
176    }
177
178    /// Create a builder
179    pub fn builder() -> AuthProxyBuilder {
180        AuthProxyBuilder::new()
181    }
182
183    /// Authenticate a request
184    pub async fn authenticate(&self, request: &AuthRequest) -> Result<AuthResult, AuthError> {
185        self.handler.authenticate(request).await
186    }
187
188    /// Authenticate and create session
189    pub async fn authenticate_and_create_session(
190        &self,
191        request: &AuthRequest,
192    ) -> Result<(AuthResult, Session), AuthError> {
193        let result = self.handler.authenticate(request).await?;
194
195        let session = self.session_manager.create_session(
196            result.identity.clone(),
197            request.client_ip,
198            request.headers.get("user-agent").cloned(),
199        ).map_err(|e| AuthError::Session(e.to_string()))?;
200
201        Ok((result, session))
202    }
203
204    /// Validate session token
205    pub fn validate_session(&self, token: &str) -> Result<Identity, AuthError> {
206        self.session_manager.validate_token(token)
207            .map_err(|e| AuthError::Session(e.to_string()))
208    }
209
210    /// Map identity to database roles
211    pub fn map_roles(&self, identity: &Identity) -> Vec<String> {
212        self.role_mapper.map_roles(identity)
213    }
214
215    /// Get database credentials for an identity
216    pub fn get_credentials(&self, key: &str) -> Result<DatabaseCredential, AuthError> {
217        self.credential_manager
218            .as_ref()
219            .ok_or_else(|| AuthError::Configuration("Credential manager not configured".to_string()))?
220            .get_credential(key)
221            .map_err(|e| AuthError::Configuration(e.to_string()))
222    }
223
224    /// Invalidate session
225    pub fn invalidate_session(&self, token: &str) -> Result<(), AuthError> {
226        self.session_manager.invalidate_session(token)
227            .map_err(|e| AuthError::Session(e.to_string()))
228    }
229
230    /// Get authentication handler
231    pub fn handler(&self) -> &AuthenticationHandler {
232        &self.handler
233    }
234
235    /// Get role mapper
236    pub fn role_mapper(&self) -> &RoleMapper {
237        &self.role_mapper
238    }
239
240    /// Get session manager
241    pub fn session_manager(&self) -> &SessionManager {
242        &self.session_manager
243    }
244}
245
246/// Auth proxy builder
247pub struct AuthProxyBuilder {
248    handler: Option<AuthenticationHandler>,
249    role_mapper: Option<RoleMapper>,
250    session_manager: Option<SessionManager>,
251    credential_manager: Option<CredentialManager>,
252}
253
254impl AuthProxyBuilder {
255    /// Create a new builder
256    pub fn new() -> Self {
257        Self {
258            handler: None,
259            role_mapper: None,
260            session_manager: None,
261            credential_manager: None,
262        }
263    }
264
265    /// Set authentication handler
266    pub fn handler(mut self, handler: AuthenticationHandler) -> Self {
267        self.handler = Some(handler);
268        self
269    }
270
271    /// Set role mapper
272    pub fn role_mapper(mut self, mapper: RoleMapper) -> Self {
273        self.role_mapper = Some(mapper);
274        self
275    }
276
277    /// Set session manager
278    pub fn session_manager(mut self, manager: SessionManager) -> Self {
279        self.session_manager = Some(manager);
280        self
281    }
282
283    /// Set credential manager
284    pub fn credential_manager(mut self, manager: CredentialManager) -> Self {
285        self.credential_manager = Some(manager);
286        self
287    }
288
289    /// Build the auth proxy
290    pub fn build(self) -> AuthProxy {
291        let handler = self.handler.unwrap_or_else(|| {
292            AuthenticationHandler::builder()
293                .enabled(false)
294                .build()
295        });
296
297        let role_mapper = self.role_mapper.unwrap_or_else(|| {
298            RoleMapper::new()
299        });
300
301        let session_manager = self.session_manager.unwrap_or_else(|| {
302            SessionManager::new(SessionConfig::default())
303        });
304
305        let mut proxy = AuthProxy::new(handler, role_mapper, session_manager);
306        proxy.credential_manager = self.credential_manager;
307        proxy
308    }
309}
310
311impl Default for AuthProxyBuilder {
312    fn default() -> Self {
313        Self::new()
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320    use std::collections::HashMap;
321
322    fn test_identity() -> Identity {
323        Identity {
324            user_id: "testuser".to_string(),
325            name: Some("Test User".to_string()),
326            email: Some("test@example.com".to_string()),
327            roles: vec!["user".to_string()],
328            groups: vec!["developers".to_string()],
329            tenant_id: None,
330            claims: HashMap::new(),
331            auth_method: "test".to_string(),
332            authenticated_at: chrono::Utc::now(),
333        }
334    }
335
336    #[test]
337    fn test_auth_proxy_builder() {
338        let proxy = AuthProxy::builder()
339            .handler(AuthenticationHandler::builder().enabled(false).build())
340            .role_mapper(RoleMapper::builder().default_role("user").build())
341            .session_manager(SessionManager::new(SessionConfig::default()))
342            .build();
343
344        assert!(!proxy.handler().is_enabled());
345    }
346
347    #[test]
348    fn test_auth_proxy_map_roles() {
349        let proxy = AuthProxy::builder()
350            .role_mapper(
351                RoleMapper::builder()
352                    .group_role("developers", "db_developer")
353                    .build()
354            )
355            .build();
356
357        let roles = proxy.map_roles(&test_identity());
358        assert!(roles.contains(&"db_developer".to_string()));
359    }
360
361    #[tokio::test]
362    async fn test_auth_proxy_disabled() {
363        let proxy = AuthProxy::builder().build();
364
365        let request = AuthRequest::new();
366        let result = proxy.authenticate(&request).await.unwrap();
367
368        assert_eq!(result.identity.auth_method, "anonymous");
369    }
370
371    #[test]
372    fn test_session_integration() {
373        let proxy = AuthProxy::builder()
374            .session_manager(
375                SessionManager::builder()
376                    .max_sessions_per_user(5)
377                    .build()
378            )
379            .build();
380
381        // Create session directly via session manager
382        let session = proxy.session_manager()
383            .create_session(test_identity(), None, None)
384            .unwrap();
385
386        // Validate via proxy
387        let identity = proxy.validate_session(&session.token).unwrap();
388        assert_eq!(identity.user_id, "testuser");
389
390        // Invalidate via proxy
391        proxy.invalidate_session(&session.token).unwrap();
392        assert!(proxy.validate_session(&session.token).is_err());
393    }
394}
heliosdb_proxy/auth/mod.rs

heliosdb_proxy/auth/
mod.rs
