Skip to main content

heliosdb_proxy/auth/
config.rs

1//! Authentication Proxy Configuration
2//!
3//! Configuration types for authentication, authorization, and credential management.
4
5use std::collections::{HashMap, HashSet};
6use std::time::Duration;
7
8use serde::{Deserialize, Serialize};
9
10/// Main authentication configuration
11#[derive(Debug, Clone)]
12pub struct AuthConfig {
13    /// Whether authentication is enabled
14    pub enabled: bool,
15
16    /// JWT configuration
17    pub jwt: Option<JwtConfig>,
18
19    /// OAuth configuration
20    pub oauth: Option<OAuthConfig>,
21
22    /// LDAP configuration
23    pub ldap: Option<LdapConfig>,
24
25    /// API key configuration
26    pub api_keys: Option<ApiKeyConfig>,
27
28    /// Role mapping rules
29    pub role_mapping: Vec<RoleMappingRule>,
30
31    /// Default role if no mapping matches
32    pub default_role: Option<String>,
33
34    /// Credential providers configuration
35    pub credentials: CredentialConfig,
36
37    /// Session configuration
38    pub session: SessionConfig,
39
40    /// Rate limiting for auth endpoints
41    pub rate_limit: AuthRateLimitConfig,
42
43    /// Ordered list of authentication methods to try
44    pub auth_methods: Vec<AuthMethod>,
45}
46
47impl Default for AuthConfig {
48    fn default() -> Self {
49        Self {
50            enabled: false,
51            jwt: None,
52            oauth: None,
53            ldap: None,
54            api_keys: None,
55            role_mapping: Vec::new(),
56            default_role: Some("db_minimal".to_string()),
57            credentials: CredentialConfig::default(),
58            session: SessionConfig::default(),
59            rate_limit: AuthRateLimitConfig::default(),
60            auth_methods: Vec::new(),
61        }
62    }
63}
64
65impl AuthConfig {
66    /// Create an enabled config with JWT
67    pub fn jwt(jwks_url: impl Into<String>) -> Self {
68        Self {
69            enabled: true,
70            jwt: Some(JwtConfig::new(jwks_url)),
71            ..Default::default()
72        }
73    }
74
75    /// Create an enabled config with API keys
76    pub fn api_keys() -> Self {
77        Self {
78            enabled: true,
79            api_keys: Some(ApiKeyConfig::default()),
80            ..Default::default()
81        }
82    }
83
84    /// Builder for AuthConfig
85    pub fn builder() -> AuthConfigBuilder {
86        AuthConfigBuilder::new()
87    }
88}
89
90/// Builder for AuthConfig
91#[derive(Default)]
92pub struct AuthConfigBuilder {
93    config: AuthConfig,
94}
95
96impl AuthConfigBuilder {
97    pub fn new() -> Self {
98        Self {
99            config: AuthConfig {
100                enabled: true,
101                ..Default::default()
102            },
103        }
104    }
105
106    pub fn jwt(mut self, config: JwtConfig) -> Self {
107        self.config.jwt = Some(config);
108        self
109    }
110
111    pub fn oauth(mut self, config: OAuthConfig) -> Self {
112        self.config.oauth = Some(config);
113        self
114    }
115
116    pub fn ldap(mut self, config: LdapConfig) -> Self {
117        self.config.ldap = Some(config);
118        self
119    }
120
121    pub fn api_keys(mut self, config: ApiKeyConfig) -> Self {
122        self.config.api_keys = Some(config);
123        self
124    }
125
126    pub fn add_role_mapping(mut self, rule: RoleMappingRule) -> Self {
127        self.config.role_mapping.push(rule);
128        self
129    }
130
131    pub fn default_role(mut self, role: impl Into<String>) -> Self {
132        self.config.default_role = Some(role.into());
133        self
134    }
135
136    pub fn credentials(mut self, config: CredentialConfig) -> Self {
137        self.config.credentials = config;
138        self
139    }
140
141    pub fn session(mut self, config: SessionConfig) -> Self {
142        self.config.session = config;
143        self
144    }
145
146    pub fn build(self) -> AuthConfig {
147        self.config
148    }
149}
150
151/// JWT authentication configuration
152#[derive(Debug, Clone)]
153pub struct JwtConfig {
154    /// JWKS URL for key fetching
155    pub jwks_url: String,
156
157    /// How often to refresh JWKS
158    pub jwks_refresh_interval: Duration,
159
160    /// Allowed token issuers
161    pub allowed_issuers: HashSet<String>,
162
163    /// Required audience claim
164    pub required_audience: Option<String>,
165
166    /// Clock skew tolerance
167    pub clock_skew: Duration,
168
169    /// Claim to use as user ID
170    pub user_id_claim: String,
171
172    /// Claim to use as roles
173    pub roles_claim: Option<String>,
174
175    /// Algorithm restrictions
176    pub allowed_algorithms: Vec<String>,
177}
178
179impl Default for JwtConfig {
180    fn default() -> Self {
181        Self {
182            jwks_url: String::new(),
183            jwks_refresh_interval: Duration::from_secs(3600),
184            allowed_issuers: HashSet::new(),
185            required_audience: None,
186            clock_skew: Duration::from_secs(60),
187            user_id_claim: "sub".to_string(),
188            roles_claim: Some("roles".to_string()),
189            allowed_algorithms: vec!["RS256".to_string(), "ES256".to_string()],
190        }
191    }
192}
193
194impl JwtConfig {
195    pub fn new(jwks_url: impl Into<String>) -> Self {
196        Self {
197            jwks_url: jwks_url.into(),
198            ..Default::default()
199        }
200    }
201
202    pub fn with_issuer(mut self, issuer: impl Into<String>) -> Self {
203        self.allowed_issuers.insert(issuer.into());
204        self
205    }
206
207    pub fn with_audience(mut self, audience: impl Into<String>) -> Self {
208        self.required_audience = Some(audience.into());
209        self
210    }
211}
212
213/// OAuth introspection configuration
214#[derive(Debug, Clone)]
215pub struct OAuthConfig {
216    /// Token introspection endpoint
217    pub introspection_url: String,
218
219    /// Client ID for introspection
220    pub client_id: String,
221
222    /// Client secret for introspection
223    pub client_secret: String,
224
225    /// Token endpoint (for client credentials)
226    pub token_url: Option<String>,
227
228    /// Scopes to request
229    pub scopes: Vec<String>,
230
231    /// Cache introspection results
232    pub cache_ttl: Duration,
233
234    /// Required scopes that must be present on a validated token
235    pub required_scopes: Vec<String>,
236
237    /// Token issuer identifier
238    pub issuer: String,
239
240    /// Authorization endpoint URL (for authorization code flow)
241    pub authorization_url: Option<String>,
242
243    /// Expected audience claim
244    pub audience: Option<String>,
245}
246
247impl Default for OAuthConfig {
248    fn default() -> Self {
249        Self {
250            introspection_url: String::new(),
251            client_id: String::new(),
252            client_secret: String::new(),
253            token_url: None,
254            scopes: Vec::new(),
255            cache_ttl: Duration::from_secs(60),
256            required_scopes: Vec::new(),
257            issuer: String::new(),
258            authorization_url: None,
259            audience: None,
260        }
261    }
262}
263
264impl OAuthConfig {
265    pub fn new(
266        introspection_url: impl Into<String>,
267        client_id: impl Into<String>,
268        client_secret: impl Into<String>,
269    ) -> Self {
270        Self {
271            introspection_url: introspection_url.into(),
272            client_id: client_id.into(),
273            client_secret: client_secret.into(),
274            ..Default::default()
275        }
276    }
277}
278
279/// LDAP authentication configuration
280#[derive(Debug, Clone)]
281pub struct LdapConfig {
282    /// LDAP server URL
283    pub server_url: String,
284
285    /// Bind DN for searches
286    pub bind_dn: String,
287
288    /// Bind password
289    pub bind_password: String,
290
291    /// User search base
292    pub user_search_base: String,
293
294    /// User search filter (use {0} for username placeholder)
295    pub user_filter: String,
296
297    /// Group search base
298    pub group_search_base: Option<String>,
299
300    /// Group attribute to read
301    pub group_attribute: String,
302
303    /// Connection timeout
304    pub timeout: Duration,
305
306    /// Use STARTTLS
307    pub starttls: bool,
308}
309
310impl Default for LdapConfig {
311    fn default() -> Self {
312        Self {
313            server_url: "ldap://localhost:389".to_string(),
314            bind_dn: String::new(),
315            bind_password: String::new(),
316            user_search_base: String::new(),
317            user_filter: "(uid={0})".to_string(),
318            group_search_base: None,
319            group_attribute: "memberOf".to_string(),
320            timeout: Duration::from_secs(10),
321            starttls: false,
322        }
323    }
324}
325
326/// API key authentication configuration
327#[derive(Debug, Clone)]
328pub struct ApiKeyConfig {
329    /// Header name to read API key from
330    pub header_name: String,
331
332    /// Query parameter name to read API key from (optional)
333    pub query_param: Option<String>,
334
335    /// Key prefix for generated keys (e.g., "hdb_")
336    pub prefix: Option<String>,
337
338    /// Hash algorithm for storage
339    pub hash_algorithm: String,
340}
341
342impl Default for ApiKeyConfig {
343    fn default() -> Self {
344        Self {
345            header_name: "X-API-Key".to_string(),
346            query_param: None,
347            prefix: Some("hpk_".to_string()),
348            hash_algorithm: "sha256".to_string(),
349        }
350    }
351}
352
353/// Role mapping rule
354#[derive(Debug, Clone)]
355pub struct RoleMappingRule {
356    /// Rule name for identification
357    pub name: String,
358
359    /// Condition to match
360    pub condition: RoleCondition,
361
362    /// Database role to assign
363    pub db_role: String,
364
365    /// Priority (higher = evaluated first)
366    pub priority: i32,
367
368    /// Roles to assign when this rule matches
369    pub assign_roles: Vec<String>,
370
371    /// Permissions granted by this rule
372    pub permissions: Vec<String>,
373
374    /// Conditions that must be met (using RoleMappingCondition)
375    pub conditions: Vec<RoleMappingCondition>,
376}
377
378impl RoleMappingRule {
379    pub fn new(condition: RoleCondition, db_role: impl Into<String>) -> Self {
380        Self {
381            name: String::new(),
382            condition,
383            db_role: db_role.into(),
384            priority: 0,
385            assign_roles: Vec::new(),
386            permissions: Vec::new(),
387            conditions: Vec::new(),
388        }
389    }
390
391    pub fn with_priority(mut self, priority: i32) -> Self {
392        self.priority = priority;
393        self
394    }
395}
396
397/// Conditions for role mapping
398#[derive(Debug, Clone)]
399pub enum RoleCondition {
400    /// Match JWT claim value
401    JwtClaim { name: String, value: String },
402
403    /// Match any JWT claim value from list
404    JwtClaimAny { name: String, values: Vec<String> },
405
406    /// Match OAuth scope
407    OAuthScope(String),
408
409    /// Match group membership
410    Group(String),
411
412    /// Match email domain
413    EmailDomain(String),
414
415    /// Match tenant ID
416    TenantId(String),
417
418    /// Compound AND condition
419    And(Vec<RoleCondition>),
420
421    /// Compound OR condition
422    Or(Vec<RoleCondition>),
423
424    /// Always match (catch-all)
425    Always,
426}
427
428/// Role mapping condition (alias for backward compatibility)
429/// Provides a more expressive condition language for role mapping
430#[derive(Debug, Clone)]
431pub enum RoleMappingCondition {
432    /// Match a specific claim value
433    HasClaim { claim: String, value: Option<String> },
434
435    /// Match group membership
436    InGroup { group: String },
437
438    /// Match existing role
439    HasRole { role: String },
440
441    /// Match tenant ID
442    FromTenant { tenant_id: String },
443
444    /// Match authentication method
445    AuthMethod { method: String },
446
447    /// Match email domain
448    EmailDomain { domain: String },
449
450    /// Match username pattern (supports wildcards)
451    UsernamePattern { pattern: String },
452
453    /// All conditions must match
454    And { conditions: Vec<RoleMappingCondition> },
455
456    /// Any condition must match
457    Or { conditions: Vec<RoleMappingCondition> },
458
459    /// Negate a condition
460    Not { condition: Box<RoleMappingCondition> },
461}
462
463impl RoleMappingCondition {
464    /// Create a HasClaim condition
465    pub fn has_claim(claim: impl Into<String>, value: Option<String>) -> Self {
466        Self::HasClaim {
467            claim: claim.into(),
468            value,
469        }
470    }
471
472    /// Create an InGroup condition
473    pub fn in_group(group: impl Into<String>) -> Self {
474        Self::InGroup {
475            group: group.into(),
476        }
477    }
478
479    /// Create a HasRole condition
480    pub fn has_role(role: impl Into<String>) -> Self {
481        Self::HasRole {
482            role: role.into(),
483        }
484    }
485
486    /// Create an AuthMethod condition
487    pub fn auth_method(method: impl Into<String>) -> Self {
488        Self::AuthMethod {
489            method: method.into(),
490        }
491    }
492}
493
494impl RoleCondition {
495    pub fn jwt_claim(name: impl Into<String>, value: impl Into<String>) -> Self {
496        Self::JwtClaim {
497            name: name.into(),
498            value: value.into(),
499        }
500    }
501
502    pub fn group(name: impl Into<String>) -> Self {
503        Self::Group(name.into())
504    }
505
506    pub fn email_domain(domain: impl Into<String>) -> Self {
507        Self::EmailDomain(domain.into())
508    }
509}
510
511/// Credential provider configuration
512#[derive(Debug, Clone)]
513pub struct CredentialConfig {
514    /// Default credential provider
515    pub default_provider: CredentialProvider,
516
517    /// Static credentials
518    pub static_credentials: HashMap<String, Credentials>,
519
520    /// Vault configuration
521    pub vault: Option<VaultConfig>,
522
523    /// AWS Secrets Manager configuration
524    pub aws_secrets: Option<AwsSecretsConfig>,
525
526    /// Credential cache TTL
527    pub cache_ttl: Duration,
528}
529
530impl Default for CredentialConfig {
531    fn default() -> Self {
532        Self {
533            default_provider: CredentialProvider::Static,
534            static_credentials: HashMap::new(),
535            vault: None,
536            aws_secrets: None,
537            cache_ttl: Duration::from_secs(300),
538        }
539    }
540}
541
542/// Credential provider type
543#[derive(Debug, Clone, Copy, PartialEq, Eq)]
544pub enum CredentialProvider {
545    Static,
546    Vault,
547    AwsSecrets,
548}
549
550/// Database credentials
551#[derive(Debug, Clone)]
552pub struct Credentials {
553    /// Username
554    pub username: String,
555
556    /// Password
557    pub password: String,
558
559    /// Time-to-live (for dynamic credentials)
560    pub ttl: Option<Duration>,
561
562    /// Additional connection options
563    pub options: HashMap<String, String>,
564}
565
566impl Credentials {
567    pub fn new(username: impl Into<String>, password: impl Into<String>) -> Self {
568        Self {
569            username: username.into(),
570            password: password.into(),
571            ttl: None,
572            options: HashMap::new(),
573        }
574    }
575
576    pub fn with_ttl(mut self, ttl: Duration) -> Self {
577        self.ttl = Some(ttl);
578        self
579    }
580}
581
582/// Vault configuration
583#[derive(Debug, Clone)]
584pub struct VaultConfig {
585    /// Vault address
586    pub address: String,
587
588    /// Authentication method
589    pub auth_method: VaultAuthMethod,
590
591    /// Vault role
592    pub role: String,
593
594    /// Secret path prefix
595    pub secret_path: String,
596
597    /// TLS configuration
598    pub tls_verify: bool,
599}
600
601/// Vault authentication method
602#[derive(Debug, Clone)]
603pub enum VaultAuthMethod {
604    Token(String),
605    Kubernetes { role: String },
606    AppRole { role_id: String, secret_id: String },
607}
608
609/// AWS Secrets Manager configuration
610#[derive(Debug, Clone)]
611pub struct AwsSecretsConfig {
612    /// AWS region
613    pub region: String,
614
615    /// Secret name prefix
616    pub secret_prefix: String,
617
618    /// Use IAM role for authentication
619    pub use_iam_role: bool,
620}
621
622/// Session configuration
623#[derive(Debug, Clone)]
624pub struct SessionConfig {
625    /// Session timeout
626    pub timeout: Duration,
627
628    /// Maximum sessions per identity
629    pub max_sessions_per_identity: usize,
630
631    /// Maximum sessions per user (used by session manager)
632    pub max_sessions_per_user: usize,
633
634    /// Idle timeout — session expires after this duration of inactivity
635    pub idle_timeout: Duration,
636
637    /// Absolute timeout — maximum session lifetime regardless of activity
638    pub absolute_timeout: Duration,
639
640    /// Whether to use secure cookies
641    pub secure_cookies: bool,
642
643    /// Session variables to set
644    pub session_vars: HashMap<String, String>,
645
646    /// Extend session on activity
647    pub extend_on_activity: bool,
648}
649
650impl Default for SessionConfig {
651    fn default() -> Self {
652        Self {
653            timeout: Duration::from_secs(3600),
654            max_sessions_per_identity: 10,
655            max_sessions_per_user: 10,
656            idle_timeout: Duration::from_secs(1800),
657            absolute_timeout: Duration::from_secs(86400),
658            secure_cookies: true,
659            session_vars: HashMap::new(),
660            extend_on_activity: true,
661        }
662    }
663}
664
665/// Rate limiting configuration for auth
666#[derive(Debug, Clone)]
667pub struct AuthRateLimitConfig {
668    /// Enable rate limiting
669    pub enabled: bool,
670
671    /// Max auth attempts per minute per IP
672    pub max_attempts_per_ip: u32,
673
674    /// Max auth failures per minute per IP
675    pub max_failures_per_ip: u32,
676
677    /// Lockout duration after too many failures
678    pub lockout_duration: Duration,
679
680    /// Rate limit window in seconds
681    pub window_seconds: u64,
682
683    /// Max requests per user within the window
684    pub max_requests_per_user: u32,
685
686    /// Max requests per IP within the window
687    pub max_requests_per_ip: u32,
688}
689
690impl Default for AuthRateLimitConfig {
691    fn default() -> Self {
692        Self {
693            enabled: true,
694            max_attempts_per_ip: 60,
695            max_failures_per_ip: 10,
696            lockout_duration: Duration::from_secs(300),
697            window_seconds: 60,
698            max_requests_per_user: 120,
699            max_requests_per_ip: 60,
700        }
701    }
702}
703
704/// Authentication type in request
705#[derive(Debug, Clone)]
706pub enum AuthType {
707    /// JWT bearer token
708    Jwt(String),
709
710    /// OAuth access token
711    OAuth(String),
712
713    /// Basic auth (username/password)
714    Basic { username: String, password: String },
715
716    /// API key
717    ApiKey(String),
718
719    /// No authentication
720    None,
721}
722
723/// Authentication method enum
724/// Used for role mapping and audit logging
725#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
726pub enum AuthMethod {
727    /// JWT-based authentication
728    Jwt,
729
730    /// OAuth token introspection
731    OAuth,
732
733    /// LDAP authentication
734    Ldap,
735
736    /// API key authentication
737    ApiKey,
738
739    /// HTTP Basic authentication
740    Basic,
741
742    /// Trust-based (internal services)
743    Trust,
744
745    /// Agent token authentication
746    AgentToken,
747
748    /// Session token authentication
749    Session,
750
751    /// No authentication (anonymous)
752    Anonymous,
753}
754
755impl std::fmt::Display for AuthMethod {
756    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
757        match self {
758            Self::Jwt => write!(f, "jwt"),
759            Self::OAuth => write!(f, "oauth"),
760            Self::Ldap => write!(f, "ldap"),
761            Self::ApiKey => write!(f, "api_key"),
762            Self::Basic => write!(f, "basic"),
763            Self::Trust => write!(f, "trust"),
764            Self::AgentToken => write!(f, "agent_token"),
765            Self::Session => write!(f, "session"),
766            Self::Anonymous => write!(f, "anonymous"),
767        }
768    }
769}
770
771impl AuthMethod {
772    /// Parse from string
773    pub fn from_str(s: &str) -> Option<Self> {
774        match s.to_lowercase().as_str() {
775            "jwt" => Some(Self::Jwt),
776            "oauth" => Some(Self::OAuth),
777            "ldap" => Some(Self::Ldap),
778            "api_key" | "apikey" => Some(Self::ApiKey),
779            "basic" => Some(Self::Basic),
780            "trust" => Some(Self::Trust),
781            "agent_token" | "agent" => Some(Self::AgentToken),
782            "session" => Some(Self::Session),
783            "anonymous" | "none" => Some(Self::Anonymous),
784            _ => None,
785        }
786    }
787}
788
789/// Authenticated identity
790#[derive(Debug, Clone, Serialize, Deserialize)]
791pub struct Identity {
792    /// Unique user identifier
793    pub user_id: String,
794
795    /// Display name
796    pub name: Option<String>,
797
798    /// Email address
799    pub email: Option<String>,
800
801    /// Assigned roles
802    pub roles: Vec<String>,
803
804    /// Group memberships
805    pub groups: Vec<String>,
806
807    /// Tenant ID (for multi-tenancy)
808    pub tenant_id: Option<String>,
809
810    /// Additional claims/attributes
811    pub claims: HashMap<String, serde_json::Value>,
812
813    /// How the identity was authenticated
814    pub auth_method: String,
815
816    /// When authentication occurred
817    pub authenticated_at: chrono::DateTime<chrono::Utc>,
818}
819
820impl Identity {
821    /// Create a new identity
822    pub fn new(user_id: impl Into<String>, auth_method: impl Into<String>) -> Self {
823        Self {
824            user_id: user_id.into(),
825            name: None,
826            email: None,
827            roles: Vec::new(),
828            groups: Vec::new(),
829            tenant_id: None,
830            claims: HashMap::new(),
831            auth_method: auth_method.into(),
832            authenticated_at: chrono::Utc::now(),
833        }
834    }
835
836    /// Create from JWT claims
837    pub fn from_jwt_claims(claims: &JwtClaims) -> Self {
838        let mut identity = Self::new(&claims.sub, "jwt");
839        identity.name = claims.name.clone();
840        identity.email = claims.email.clone();
841        identity.roles = claims.roles.clone();
842        identity.tenant_id = claims.tenant_id.clone();
843        identity.claims = claims.custom.clone();
844        identity
845    }
846
847    /// Check if identity has a role
848    pub fn has_role(&self, role: &str) -> bool {
849        self.roles.iter().any(|r| r == role)
850    }
851
852    /// Check if identity is in a group
853    pub fn in_group(&self, group: &str) -> bool {
854        self.groups.iter().any(|g| g == group)
855    }
856
857    /// Check if identity is admin
858    pub fn is_admin(&self) -> bool {
859        self.has_role("admin") || self.has_role("db_admin")
860    }
861
862    /// Get claim value
863    pub fn get_claim(&self, name: &str) -> Option<&serde_json::Value> {
864        self.claims.get(name)
865    }
866
867    /// Get email domain
868    pub fn email_domain(&self) -> Option<&str> {
869        self.email.as_ref().and_then(|e| e.split('@').nth(1))
870    }
871
872    /// Create an anonymous identity
873    pub fn anonymous() -> Self {
874        Self {
875            user_id: "anonymous".to_string(),
876            name: None,
877            email: None,
878            roles: Vec::new(),
879            groups: Vec::new(),
880            tenant_id: None,
881            claims: HashMap::new(),
882            auth_method: "anonymous".to_string(),
883            authenticated_at: chrono::Utc::now(),
884        }
885    }
886}
887
888/// JWT claims structure
889#[derive(Debug, Clone, Serialize, Deserialize)]
890pub struct JwtClaims {
891    /// Subject (user ID)
892    pub sub: String,
893
894    /// Issuer
895    pub iss: String,
896
897    /// Audience
898    pub aud: Option<Vec<String>>,
899
900    /// Expiration time
901    pub exp: i64,
902
903    /// Issued at
904    pub iat: i64,
905
906    /// Not before
907    pub nbf: Option<i64>,
908
909    /// JWT ID
910    pub jti: Option<String>,
911
912    /// User's name
913    pub name: Option<String>,
914
915    /// User's email
916    pub email: Option<String>,
917
918    /// User's roles
919    #[serde(default)]
920    pub roles: Vec<String>,
921
922    /// Tenant ID
923    pub tenant_id: Option<String>,
924
925    /// Custom claims
926    #[serde(flatten)]
927    pub custom: HashMap<String, serde_json::Value>,
928}
929
930/// Agent-specific identity for AI workloads
931#[derive(Debug, Clone, Serialize, Deserialize)]
932pub struct AgentIdentity {
933    /// Agent ID
934    pub id: String,
935
936    /// Agent type (e.g., "claude", "gpt", "custom")
937    pub agent_type: String,
938
939    /// Allowed tools
940    pub allowed_tools: Vec<String>,
941
942    /// Resource quota
943    pub quota: AgentQuota,
944
945    /// Conversation ID (optional)
946    pub conversation_id: Option<String>,
947
948    /// Parent identity (human that authorized the agent)
949    pub parent_identity: Option<String>,
950}
951
952/// Agent resource quota
953#[derive(Debug, Clone, Serialize, Deserialize)]
954pub struct AgentQuota {
955    /// Maximum queries per conversation
956    pub max_queries_per_conversation: u32,
957
958    /// Maximum rows per query
959    pub max_rows_per_query: u32,
960
961    /// Token budget (for LLM cost tracking)
962    pub token_budget: u64,
963
964    /// Allowed tables
965    pub allowed_tables: Option<Vec<String>>,
966}
967
968impl Default for AgentQuota {
969    fn default() -> Self {
970        Self {
971            max_queries_per_conversation: 100,
972            max_rows_per_query: 1000,
973            token_budget: 100000,
974            allowed_tables: None,
975        }
976    }
977}
978
979/// Tool permission for AI agents
980#[derive(Debug, Clone)]
981pub struct ToolPermission {
982    /// Database role to use
983    pub db_role: String,
984
985    /// Tables allowed for this tool
986    pub allowed_tables: Vec<String>,
987
988    /// Read-only mode
989    pub read_only: bool,
990}
991
992#[cfg(test)]
993mod tests {
994    use super::*;
995
996    #[test]
997    fn test_auth_config_builder() {
998        let config = AuthConfig::builder()
999            .jwt(JwtConfig::new("https://auth.example.com/.well-known/jwks.json"))
1000            .add_role_mapping(RoleMappingRule::new(
1001                RoleCondition::jwt_claim("role", "admin"),
1002                "db_admin",
1003            ).with_priority(100))
1004            .default_role("db_minimal")
1005            .build();
1006
1007        assert!(config.enabled);
1008        assert!(config.jwt.is_some());
1009        assert_eq!(config.role_mapping.len(), 1);
1010    }
1011
1012    #[test]
1013    fn test_identity() {
1014        let identity = Identity::new("user123", "jwt");
1015        assert_eq!(identity.user_id, "user123");
1016        assert_eq!(identity.auth_method, "jwt");
1017        assert!(!identity.is_admin());
1018    }
1019
1020    #[test]
1021    fn test_identity_roles() {
1022        let mut identity = Identity::new("admin123", "jwt");
1023        identity.roles = vec!["admin".to_string(), "db_readwrite".to_string()];
1024
1025        assert!(identity.is_admin());
1026        assert!(identity.has_role("admin"));
1027        assert!(identity.has_role("db_readwrite"));
1028        assert!(!identity.has_role("superuser"));
1029    }
1030
1031    #[test]
1032    fn test_email_domain() {
1033        let mut identity = Identity::new("user", "jwt");
1034        identity.email = Some("alice@example.com".to_string());
1035
1036        assert_eq!(identity.email_domain(), Some("example.com"));
1037    }
1038
1039    #[test]
1040    fn test_credentials() {
1041        let creds = Credentials::new("dbuser", "password123")
1042            .with_ttl(Duration::from_secs(3600));
1043
1044        assert_eq!(creds.username, "dbuser");
1045        assert!(creds.ttl.is_some());
1046    }
1047
1048    #[test]
1049    fn test_role_mapping() {
1050        let rule = RoleMappingRule::new(
1051            RoleCondition::group("developers"),
1052            "db_readwrite",
1053        ).with_priority(50);
1054
1055        assert_eq!(rule.db_role, "db_readwrite");
1056        assert_eq!(rule.priority, 50);
1057    }
1058}