1use std::collections::HashMap;
6use std::sync::Arc;
7
8use parking_lot::RwLock;
9
10use super::config::{Identity, RoleMappingCondition, RoleMappingRule};
11
12pub struct RoleMapper {
14 rules: Vec<RoleMappingRule>,
16
17 static_roles: Arc<RwLock<HashMap<String, Vec<String>>>>,
19
20 group_roles: HashMap<String, Vec<String>>,
22
23 default_roles: Vec<String>,
25
26 anonymous_role: Option<String>,
28}
29
30impl RoleMapper {
31 pub fn new() -> Self {
33 Self {
34 rules: Vec::new(),
35 static_roles: Arc::new(RwLock::new(HashMap::new())),
36 group_roles: HashMap::new(),
37 default_roles: Vec::new(),
38 anonymous_role: None,
39 }
40 }
41
42 pub fn builder() -> RoleMapperBuilder {
44 RoleMapperBuilder::new()
45 }
46
47 pub fn map_roles(&self, identity: &Identity) -> Vec<String> {
49 let mut roles = Vec::new();
50
51 roles.extend(identity.roles.clone());
53
54 if let Some(static_roles) = self.static_roles.read().get(&identity.user_id) {
56 roles.extend(static_roles.clone());
57 }
58
59 for group in &identity.groups {
61 if let Some(group_roles) = self.group_roles.get(group) {
62 roles.extend(group_roles.clone());
63 }
64 }
65
66 for rule in &self.rules {
68 if self.evaluate_rule(rule, identity) {
69 roles.extend(rule.assign_roles.clone());
70 }
71 }
72
73 if roles.is_empty() {
75 roles.extend(self.default_roles.clone());
76 }
77
78 roles.sort();
80 roles.dedup();
81
82 roles
83 }
84
85 pub fn map_primary_role(&self, identity: &Identity) -> Option<String> {
87 let roles = self.map_roles(identity);
88 roles.into_iter().next()
89 }
90
91 pub fn has_permission(&self, identity: &Identity, permission: &str) -> bool {
93 let roles = self.map_roles(identity);
94
95 for rule in &self.rules {
96 if roles.iter().any(|r| rule.assign_roles.contains(r))
97 && rule.permissions.contains(&permission.to_string())
98 {
99 return true;
100 }
101 }
102
103 false
104 }
105
106 pub fn get_permissions(&self, identity: &Identity) -> Vec<String> {
108 let roles = self.map_roles(identity);
109 let mut permissions = Vec::new();
110
111 for rule in &self.rules {
112 if roles.iter().any(|r| rule.assign_roles.contains(r)) {
113 permissions.extend(rule.permissions.clone());
114 }
115 }
116
117 permissions.sort();
118 permissions.dedup();
119 permissions
120 }
121
122 pub fn assign_static_role(&self, user_id: impl Into<String>, role: impl Into<String>) {
124 let user_id = user_id.into();
125 let role = role.into();
126
127 let mut static_roles = self.static_roles.write();
128 static_roles.entry(user_id).or_default().push(role);
129 }
130
131 pub fn remove_static_role(&self, user_id: &str, role: &str) {
133 let mut static_roles = self.static_roles.write();
134 if let Some(roles) = static_roles.get_mut(user_id) {
135 roles.retain(|r| r != role);
136 }
137 }
138
139 pub fn anonymous_role(&self) -> Option<&String> {
141 self.anonymous_role.as_ref()
142 }
143
144 fn evaluate_rule(&self, rule: &RoleMappingRule, identity: &Identity) -> bool {
146 for condition in &rule.conditions {
148 if !self.evaluate_condition(condition, identity) {
149 return false;
150 }
151 }
152 true
153 }
154
155 fn evaluate_condition(&self, condition: &RoleMappingCondition, identity: &Identity) -> bool {
157 match condition {
158 RoleMappingCondition::HasClaim { claim, value } => {
159 match identity.claims.get(claim) {
160 Some(claim_value) => {
161 if let Some(expected) = value {
162 claim_value.as_str() == Some(expected.as_str())
163 } else {
164 true }
166 }
167 None => false,
168 }
169 }
170
171 RoleMappingCondition::InGroup { group } => identity.groups.contains(group),
172
173 RoleMappingCondition::HasRole { role } => identity.roles.contains(role),
174
175 RoleMappingCondition::FromTenant { tenant_id } => {
176 identity.tenant_id.as_ref() == Some(tenant_id)
177 }
178
179 RoleMappingCondition::AuthMethod { method } => &identity.auth_method == method,
180
181 RoleMappingCondition::EmailDomain { domain } => identity
182 .email
183 .as_ref()
184 .map(|e| e.ends_with(&format!("@{}", domain)))
185 .unwrap_or(false),
186
187 RoleMappingCondition::UsernamePattern { pattern } => {
188 self.match_pattern(&identity.user_id, pattern)
189 }
190
191 RoleMappingCondition::And { conditions } => conditions
192 .iter()
193 .all(|c| self.evaluate_condition(c, identity)),
194
195 RoleMappingCondition::Or { conditions } => conditions
196 .iter()
197 .any(|c| self.evaluate_condition(c, identity)),
198
199 RoleMappingCondition::Not { condition } => {
200 !self.evaluate_condition(condition, identity)
201 }
202 }
203 }
204
205 fn match_pattern(&self, value: &str, pattern: &str) -> bool {
207 if pattern == "*" {
208 return true;
209 }
210
211 if let Some(prefix) = pattern.strip_suffix('*') {
212 return value.starts_with(prefix);
213 }
214
215 if let Some(suffix) = pattern.strip_prefix('*') {
216 return value.ends_with(suffix);
217 }
218
219 value == pattern
220 }
221}
222
223impl Default for RoleMapper {
224 fn default() -> Self {
225 Self::new()
226 }
227}
228
229pub struct RoleMapperBuilder {
231 rules: Vec<RoleMappingRule>,
232 group_roles: HashMap<String, Vec<String>>,
233 default_roles: Vec<String>,
234 anonymous_role: Option<String>,
235}
236
237impl RoleMapperBuilder {
238 pub fn new() -> Self {
240 Self {
241 rules: Vec::new(),
242 group_roles: HashMap::new(),
243 default_roles: Vec::new(),
244 anonymous_role: None,
245 }
246 }
247
248 pub fn rule(mut self, rule: RoleMappingRule) -> Self {
250 self.rules.push(rule);
251 self
252 }
253
254 pub fn group_role(mut self, group: impl Into<String>, role: impl Into<String>) -> Self {
256 let group = group.into();
257 let role = role.into();
258
259 self.group_roles.entry(group).or_default().push(role);
260 self
261 }
262
263 pub fn default_role(mut self, role: impl Into<String>) -> Self {
265 self.default_roles.push(role.into());
266 self
267 }
268
269 pub fn anonymous_role(mut self, role: impl Into<String>) -> Self {
271 self.anonymous_role = Some(role.into());
272 self
273 }
274
275 pub fn build(self) -> RoleMapper {
277 RoleMapper {
278 rules: self.rules,
279 static_roles: Arc::new(RwLock::new(HashMap::new())),
280 group_roles: self.group_roles,
281 default_roles: self.default_roles,
282 anonymous_role: self.anonymous_role,
283 }
284 }
285}
286
287impl Default for RoleMapperBuilder {
288 fn default() -> Self {
289 Self::new()
290 }
291}
292
293#[derive(Debug, Clone)]
295pub struct PermissionSet {
296 pub databases: Vec<String>,
298
299 pub schemas: Vec<String>,
301
302 pub tables: Vec<String>,
304
305 pub operations: Vec<Operation>,
307
308 pub row_predicates: HashMap<String, String>,
310
311 pub column_restrictions: HashMap<String, Vec<String>>,
313}
314
315#[derive(Debug, Clone, PartialEq, Eq)]
317pub enum Operation {
318 Select,
319 Insert,
320 Update,
321 Delete,
322 Create,
323 Drop,
324 Alter,
325 Grant,
326 Execute,
327 All,
328}
329
330impl Operation {
331 #[allow(clippy::should_implement_trait)]
333 pub fn from_str(s: &str) -> Option<Self> {
334 match s.to_uppercase().as_str() {
335 "SELECT" => Some(Self::Select),
336 "INSERT" => Some(Self::Insert),
337 "UPDATE" => Some(Self::Update),
338 "DELETE" => Some(Self::Delete),
339 "CREATE" => Some(Self::Create),
340 "DROP" => Some(Self::Drop),
341 "ALTER" => Some(Self::Alter),
342 "GRANT" => Some(Self::Grant),
343 "EXECUTE" => Some(Self::Execute),
344 "ALL" => Some(Self::All),
345 _ => None,
346 }
347 }
348}
349
350impl PermissionSet {
351 pub fn empty() -> Self {
353 Self {
354 databases: Vec::new(),
355 schemas: Vec::new(),
356 tables: Vec::new(),
357 operations: Vec::new(),
358 row_predicates: HashMap::new(),
359 column_restrictions: HashMap::new(),
360 }
361 }
362
363 pub fn full_access() -> Self {
365 Self {
366 databases: vec!["*".to_string()],
367 schemas: vec!["*".to_string()],
368 tables: vec!["*".to_string()],
369 operations: vec![Operation::All],
370 row_predicates: HashMap::new(),
371 column_restrictions: HashMap::new(),
372 }
373 }
374
375 pub fn is_operation_allowed(&self, operation: &Operation, table: &str) -> bool {
377 if !self.operations.contains(&Operation::All) && !self.operations.contains(operation) {
379 return false;
380 }
381
382 if self.tables.is_empty() {
384 return true;
385 }
386
387 for pattern in &self.tables {
388 if pattern == "*" || pattern == table {
389 return true;
390 }
391
392 if pattern.ends_with('*') {
394 let prefix = &pattern[..pattern.len() - 1];
395 if table.starts_with(prefix) {
396 return true;
397 }
398 }
399 }
400
401 false
402 }
403
404 pub fn row_predicate(&self, table: &str) -> Option<&String> {
406 self.row_predicates.get(table)
407 }
408
409 pub fn allowed_columns(&self, table: &str) -> Option<&Vec<String>> {
411 self.column_restrictions.get(table)
412 }
413}
414
415#[derive(Debug, Clone)]
417pub struct AuthorizationContext {
418 pub identity: Identity,
420
421 pub roles: Vec<String>,
423
424 pub permissions: PermissionSet,
426
427 pub session_start: chrono::DateTime<chrono::Utc>,
429
430 pub context: HashMap<String, String>,
432}
433
434impl AuthorizationContext {
435 pub fn new(identity: Identity, roles: Vec<String>, permissions: PermissionSet) -> Self {
437 Self {
438 identity,
439 roles,
440 permissions,
441 session_start: chrono::Utc::now(),
442 context: HashMap::new(),
443 }
444 }
445
446 pub fn is_allowed(&self, operation: &Operation, table: &str) -> bool {
448 self.permissions.is_operation_allowed(operation, table)
449 }
450
451 pub fn with_context(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
453 self.context.insert(key.into(), value.into());
454 self
455 }
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461 use crate::auth::config::RoleCondition;
462
463 fn test_identity() -> Identity {
464 Identity {
465 user_id: "user123".to_string(),
466 name: Some("Test User".to_string()),
467 email: Some("test@example.com".to_string()),
468 roles: vec!["user".to_string()],
469 groups: vec!["developers".to_string()],
470 tenant_id: Some("tenant1".to_string()),
471 claims: {
472 let mut claims = HashMap::new();
473 claims.insert("department".to_string(), serde_json::json!("engineering"));
474 claims
475 },
476 auth_method: "jwt".to_string(),
477 authenticated_at: chrono::Utc::now(),
478 }
479 }
480
481 #[test]
482 fn test_basic_role_mapping() {
483 let mapper = RoleMapper::builder()
484 .group_role("developers", "db_developer")
485 .default_role("db_user")
486 .build();
487
488 let identity = test_identity();
489 let roles = mapper.map_roles(&identity);
490
491 assert!(roles.contains(&"user".to_string())); assert!(roles.contains(&"db_developer".to_string())); }
494
495 #[test]
496 fn test_rule_based_mapping() {
497 let mapper = RoleMapper::builder()
498 .rule(RoleMappingRule {
499 name: "admin_from_claim".to_string(),
500 condition: RoleCondition::Always,
501 db_role: String::new(),
502 conditions: vec![RoleMappingCondition::HasClaim {
503 claim: "department".to_string(),
504 value: Some("engineering".to_string()),
505 }],
506 assign_roles: vec!["db_admin".to_string()],
507 permissions: vec!["read".to_string(), "write".to_string()],
508 priority: 1,
509 })
510 .build();
511
512 let identity = test_identity();
513 let roles = mapper.map_roles(&identity);
514
515 assert!(roles.contains(&"db_admin".to_string()));
516 }
517
518 #[test]
519 fn test_tenant_condition() {
520 let mapper = RoleMapper::builder()
521 .rule(RoleMappingRule {
522 name: "tenant1_role".to_string(),
523 condition: RoleCondition::Always,
524 db_role: String::new(),
525 conditions: vec![RoleMappingCondition::FromTenant {
526 tenant_id: "tenant1".to_string(),
527 }],
528 assign_roles: vec!["tenant1_user".to_string()],
529 permissions: Vec::new(),
530 priority: 1,
531 })
532 .build();
533
534 let identity = test_identity();
535 let roles = mapper.map_roles(&identity);
536
537 assert!(roles.contains(&"tenant1_user".to_string()));
538 }
539
540 #[test]
541 fn test_email_domain_condition() {
542 let mapper = RoleMapper::builder()
543 .rule(RoleMappingRule {
544 name: "example_domain".to_string(),
545 condition: RoleCondition::Always,
546 db_role: String::new(),
547 conditions: vec![RoleMappingCondition::EmailDomain {
548 domain: "example.com".to_string(),
549 }],
550 assign_roles: vec!["internal_user".to_string()],
551 permissions: Vec::new(),
552 priority: 1,
553 })
554 .build();
555
556 let identity = test_identity();
557 let roles = mapper.map_roles(&identity);
558
559 assert!(roles.contains(&"internal_user".to_string()));
560 }
561
562 #[test]
563 fn test_and_condition() {
564 let mapper = RoleMapper::builder()
565 .rule(RoleMappingRule {
566 name: "combined".to_string(),
567 condition: RoleCondition::Always,
568 db_role: String::new(),
569 conditions: vec![RoleMappingCondition::And {
570 conditions: vec![
571 RoleMappingCondition::HasRole {
572 role: "user".to_string(),
573 },
574 RoleMappingCondition::InGroup {
575 group: "developers".to_string(),
576 },
577 ],
578 }],
579 assign_roles: vec!["power_user".to_string()],
580 permissions: Vec::new(),
581 priority: 1,
582 })
583 .build();
584
585 let identity = test_identity();
586 let roles = mapper.map_roles(&identity);
587
588 assert!(roles.contains(&"power_user".to_string()));
589 }
590
591 #[test]
592 fn test_or_condition() {
593 let mapper = RoleMapper::builder()
594 .rule(RoleMappingRule {
595 name: "either".to_string(),
596 condition: RoleCondition::Always,
597 db_role: String::new(),
598 conditions: vec![RoleMappingCondition::Or {
599 conditions: vec![
600 RoleMappingCondition::HasRole {
601 role: "admin".to_string(),
602 },
603 RoleMappingCondition::HasRole {
604 role: "user".to_string(),
605 },
606 ],
607 }],
608 assign_roles: vec!["authenticated".to_string()],
609 permissions: Vec::new(),
610 priority: 1,
611 })
612 .build();
613
614 let identity = test_identity();
615 let roles = mapper.map_roles(&identity);
616
617 assert!(roles.contains(&"authenticated".to_string()));
618 }
619
620 #[test]
621 fn test_not_condition() {
622 let mapper = RoleMapper::builder()
623 .rule(RoleMappingRule {
624 name: "not_admin".to_string(),
625 condition: RoleCondition::Always,
626 db_role: String::new(),
627 conditions: vec![RoleMappingCondition::Not {
628 condition: Box::new(RoleMappingCondition::HasRole {
629 role: "admin".to_string(),
630 }),
631 }],
632 assign_roles: vec!["regular_user".to_string()],
633 permissions: Vec::new(),
634 priority: 1,
635 })
636 .build();
637
638 let identity = test_identity();
639 let roles = mapper.map_roles(&identity);
640
641 assert!(roles.contains(&"regular_user".to_string()));
642 }
643
644 #[test]
645 fn test_static_role_assignment() {
646 let mapper = RoleMapper::new();
647 mapper.assign_static_role("user123", "special_role");
648
649 let identity = test_identity();
650 let roles = mapper.map_roles(&identity);
651
652 assert!(roles.contains(&"special_role".to_string()));
653 }
654
655 #[test]
656 fn test_default_roles() {
657 let mapper = RoleMapper::builder().default_role("guest").build();
658
659 let identity = Identity {
661 user_id: "empty".to_string(),
662 name: None,
663 email: None,
664 roles: Vec::new(),
665 groups: Vec::new(),
666 tenant_id: None,
667 claims: HashMap::new(),
668 auth_method: "none".to_string(),
669 authenticated_at: chrono::Utc::now(),
670 };
671
672 let roles = mapper.map_roles(&identity);
673 assert!(roles.contains(&"guest".to_string()));
674 }
675
676 #[test]
677 fn test_permission_set() {
678 let permissions = PermissionSet {
679 databases: vec!["mydb".to_string()],
680 schemas: vec!["public".to_string()],
681 tables: vec!["users".to_string(), "orders*".to_string()],
682 operations: vec![Operation::Select, Operation::Insert],
683 row_predicates: {
684 let mut p = HashMap::new();
685 p.insert("users".to_string(), "id = current_user_id()".to_string());
686 p
687 },
688 column_restrictions: HashMap::new(),
689 };
690
691 assert!(permissions.is_operation_allowed(&Operation::Select, "users"));
692 assert!(permissions.is_operation_allowed(&Operation::Insert, "orders_2024"));
693 assert!(!permissions.is_operation_allowed(&Operation::Delete, "users"));
694 assert!(!permissions.is_operation_allowed(&Operation::Select, "secrets"));
695
696 assert_eq!(
697 permissions.row_predicate("users"),
698 Some(&"id = current_user_id()".to_string())
699 );
700 }
701
702 #[test]
703 fn test_pattern_matching() {
704 let mapper = RoleMapper::new();
705
706 assert!(mapper.match_pattern("admin_user", "admin*"));
707 assert!(mapper.match_pattern("user_admin", "*admin"));
708 assert!(mapper.match_pattern("anything", "*"));
709 assert!(mapper.match_pattern("exact", "exact"));
710 assert!(!mapper.match_pattern("mismatch", "exact"));
711 }
712
713 #[test]
714 fn test_authorization_context() {
715 let identity = test_identity();
716 let permissions = PermissionSet::full_access();
717 let roles = vec!["admin".to_string()];
718
719 let ctx = AuthorizationContext::new(identity, roles, permissions)
720 .with_context("client_ip", "192.168.1.1");
721
722 assert!(ctx.is_allowed(&Operation::Select, "any_table"));
723 assert_eq!(
724 ctx.context.get("client_ip"),
725 Some(&"192.168.1.1".to_string())
726 );
727 }
728}