1use std::collections::{HashMap, HashSet};
6use std::time::Duration;
7
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone)]
12pub struct AuthConfig {
13 pub enabled: bool,
15
16 pub jwt: Option<JwtConfig>,
18
19 pub oauth: Option<OAuthConfig>,
21
22 pub ldap: Option<LdapConfig>,
24
25 pub api_keys: Option<ApiKeyConfig>,
27
28 pub role_mapping: Vec<RoleMappingRule>,
30
31 pub default_role: Option<String>,
33
34 pub credentials: CredentialConfig,
36
37 pub session: SessionConfig,
39
40 pub rate_limit: AuthRateLimitConfig,
42
43 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 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 pub fn api_keys() -> Self {
77 Self {
78 enabled: true,
79 api_keys: Some(ApiKeyConfig::default()),
80 ..Default::default()
81 }
82 }
83
84 pub fn builder() -> AuthConfigBuilder {
86 AuthConfigBuilder::new()
87 }
88}
89
90#[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#[derive(Debug, Clone)]
153pub struct JwtConfig {
154 pub jwks_url: String,
156
157 pub jwks_refresh_interval: Duration,
159
160 pub allowed_issuers: HashSet<String>,
162
163 pub required_audience: Option<String>,
165
166 pub clock_skew: Duration,
168
169 pub user_id_claim: String,
171
172 pub roles_claim: Option<String>,
174
175 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#[derive(Debug, Clone)]
215pub struct OAuthConfig {
216 pub introspection_url: String,
218
219 pub client_id: String,
221
222 pub client_secret: String,
224
225 pub token_url: Option<String>,
227
228 pub scopes: Vec<String>,
230
231 pub cache_ttl: Duration,
233
234 pub required_scopes: Vec<String>,
236
237 pub issuer: String,
239
240 pub authorization_url: Option<String>,
242
243 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#[derive(Debug, Clone)]
281pub struct LdapConfig {
282 pub server_url: String,
284
285 pub bind_dn: String,
287
288 pub bind_password: String,
290
291 pub user_search_base: String,
293
294 pub user_filter: String,
296
297 pub group_search_base: Option<String>,
299
300 pub group_attribute: String,
302
303 pub timeout: Duration,
305
306 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#[derive(Debug, Clone)]
328pub struct ApiKeyConfig {
329 pub header_name: String,
331
332 pub query_param: Option<String>,
334
335 pub prefix: Option<String>,
337
338 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#[derive(Debug, Clone)]
355pub struct RoleMappingRule {
356 pub name: String,
358
359 pub condition: RoleCondition,
361
362 pub db_role: String,
364
365 pub priority: i32,
367
368 pub assign_roles: Vec<String>,
370
371 pub permissions: Vec<String>,
373
374 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#[derive(Debug, Clone)]
399pub enum RoleCondition {
400 JwtClaim { name: String, value: String },
402
403 JwtClaimAny { name: String, values: Vec<String> },
405
406 OAuthScope(String),
408
409 Group(String),
411
412 EmailDomain(String),
414
415 TenantId(String),
417
418 And(Vec<RoleCondition>),
420
421 Or(Vec<RoleCondition>),
423
424 Always,
426}
427
428#[derive(Debug, Clone)]
431pub enum RoleMappingCondition {
432 HasClaim { claim: String, value: Option<String> },
434
435 InGroup { group: String },
437
438 HasRole { role: String },
440
441 FromTenant { tenant_id: String },
443
444 AuthMethod { method: String },
446
447 EmailDomain { domain: String },
449
450 UsernamePattern { pattern: String },
452
453 And { conditions: Vec<RoleMappingCondition> },
455
456 Or { conditions: Vec<RoleMappingCondition> },
458
459 Not { condition: Box<RoleMappingCondition> },
461}
462
463impl RoleMappingCondition {
464 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 pub fn in_group(group: impl Into<String>) -> Self {
474 Self::InGroup {
475 group: group.into(),
476 }
477 }
478
479 pub fn has_role(role: impl Into<String>) -> Self {
481 Self::HasRole {
482 role: role.into(),
483 }
484 }
485
486 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#[derive(Debug, Clone)]
513pub struct CredentialConfig {
514 pub default_provider: CredentialProvider,
516
517 pub static_credentials: HashMap<String, Credentials>,
519
520 pub vault: Option<VaultConfig>,
522
523 pub aws_secrets: Option<AwsSecretsConfig>,
525
526 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
544pub enum CredentialProvider {
545 Static,
546 Vault,
547 AwsSecrets,
548}
549
550#[derive(Debug, Clone)]
552pub struct Credentials {
553 pub username: String,
555
556 pub password: String,
558
559 pub ttl: Option<Duration>,
561
562 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#[derive(Debug, Clone)]
584pub struct VaultConfig {
585 pub address: String,
587
588 pub auth_method: VaultAuthMethod,
590
591 pub role: String,
593
594 pub secret_path: String,
596
597 pub tls_verify: bool,
599}
600
601#[derive(Debug, Clone)]
603pub enum VaultAuthMethod {
604 Token(String),
605 Kubernetes { role: String },
606 AppRole { role_id: String, secret_id: String },
607}
608
609#[derive(Debug, Clone)]
611pub struct AwsSecretsConfig {
612 pub region: String,
614
615 pub secret_prefix: String,
617
618 pub use_iam_role: bool,
620}
621
622#[derive(Debug, Clone)]
624pub struct SessionConfig {
625 pub timeout: Duration,
627
628 pub max_sessions_per_identity: usize,
630
631 pub max_sessions_per_user: usize,
633
634 pub idle_timeout: Duration,
636
637 pub absolute_timeout: Duration,
639
640 pub secure_cookies: bool,
642
643 pub session_vars: HashMap<String, String>,
645
646 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#[derive(Debug, Clone)]
667pub struct AuthRateLimitConfig {
668 pub enabled: bool,
670
671 pub max_attempts_per_ip: u32,
673
674 pub max_failures_per_ip: u32,
676
677 pub lockout_duration: Duration,
679
680 pub window_seconds: u64,
682
683 pub max_requests_per_user: u32,
685
686 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#[derive(Debug, Clone)]
706pub enum AuthType {
707 Jwt(String),
709
710 OAuth(String),
712
713 Basic { username: String, password: String },
715
716 ApiKey(String),
718
719 None,
721}
722
723#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
726pub enum AuthMethod {
727 Jwt,
729
730 OAuth,
732
733 Ldap,
735
736 ApiKey,
738
739 Basic,
741
742 Trust,
744
745 AgentToken,
747
748 Session,
750
751 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 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#[derive(Debug, Clone, Serialize, Deserialize)]
791pub struct Identity {
792 pub user_id: String,
794
795 pub name: Option<String>,
797
798 pub email: Option<String>,
800
801 pub roles: Vec<String>,
803
804 pub groups: Vec<String>,
806
807 pub tenant_id: Option<String>,
809
810 pub claims: HashMap<String, serde_json::Value>,
812
813 pub auth_method: String,
815
816 pub authenticated_at: chrono::DateTime<chrono::Utc>,
818}
819
820impl Identity {
821 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 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 pub fn has_role(&self, role: &str) -> bool {
849 self.roles.iter().any(|r| r == role)
850 }
851
852 pub fn in_group(&self, group: &str) -> bool {
854 self.groups.iter().any(|g| g == group)
855 }
856
857 pub fn is_admin(&self) -> bool {
859 self.has_role("admin") || self.has_role("db_admin")
860 }
861
862 pub fn get_claim(&self, name: &str) -> Option<&serde_json::Value> {
864 self.claims.get(name)
865 }
866
867 pub fn email_domain(&self) -> Option<&str> {
869 self.email.as_ref().and_then(|e| e.split('@').nth(1))
870 }
871
872 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#[derive(Debug, Clone, Serialize, Deserialize)]
890pub struct JwtClaims {
891 pub sub: String,
893
894 pub iss: String,
896
897 pub aud: Option<Vec<String>>,
899
900 pub exp: i64,
902
903 pub iat: i64,
905
906 pub nbf: Option<i64>,
908
909 pub jti: Option<String>,
911
912 pub name: Option<String>,
914
915 pub email: Option<String>,
917
918 #[serde(default)]
920 pub roles: Vec<String>,
921
922 pub tenant_id: Option<String>,
924
925 #[serde(flatten)]
927 pub custom: HashMap<String, serde_json::Value>,
928}
929
930#[derive(Debug, Clone, Serialize, Deserialize)]
932pub struct AgentIdentity {
933 pub id: String,
935
936 pub agent_type: String,
938
939 pub allowed_tools: Vec<String>,
941
942 pub quota: AgentQuota,
944
945 pub conversation_id: Option<String>,
947
948 pub parent_identity: Option<String>,
950}
951
952#[derive(Debug, Clone, Serialize, Deserialize)]
954pub struct AgentQuota {
955 pub max_queries_per_conversation: u32,
957
958 pub max_rows_per_query: u32,
960
961 pub token_budget: u64,
963
964 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#[derive(Debug, Clone)]
981pub struct ToolPermission {
982 pub db_role: String,
984
985 pub allowed_tables: Vec<String>,
987
988 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}