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