Skip to main content

heliosdb_proxy/auth/
role_mapper.rs

1//! Role Mapping
2//!
3//! Maps authenticated identities to database roles and permissions.
4
5use std::collections::HashMap;
6use std::sync::Arc;
7
8use parking_lot::RwLock;
9
10use super::config::{Identity, RoleCondition, RoleMappingRule, RoleMappingCondition};
11
12/// Role mapper
13pub struct RoleMapper {
14    /// Mapping rules
15    rules: Vec<RoleMappingRule>,
16
17    /// Static role assignments (user_id -> roles)
18    static_roles: Arc<RwLock<HashMap<String, Vec<String>>>>,
19
20    /// Group to role mappings
21    group_roles: HashMap<String, Vec<String>>,
22
23    /// Default roles for authenticated users
24    default_roles: Vec<String>,
25
26    /// Default role for anonymous users
27    anonymous_role: Option<String>,
28}
29
30impl RoleMapper {
31    /// Create a new role mapper
32    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    /// Create a builder
43    pub fn builder() -> RoleMapperBuilder {
44        RoleMapperBuilder::new()
45    }
46
47    /// Map an identity to database roles
48    pub fn map_roles(&self, identity: &Identity) -> Vec<String> {
49        let mut roles = Vec::new();
50
51        // Add roles from identity
52        roles.extend(identity.roles.clone());
53
54        // Add static roles
55        if let Some(static_roles) = self.static_roles.read().get(&identity.user_id) {
56            roles.extend(static_roles.clone());
57        }
58
59        // Add group-based roles
60        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        // Apply mapping rules
67        for rule in &self.rules {
68            if self.evaluate_rule(rule, identity) {
69                roles.extend(rule.assign_roles.clone());
70            }
71        }
72
73        // Add default roles if no roles assigned
74        if roles.is_empty() {
75            roles.extend(self.default_roles.clone());
76        }
77
78        // Deduplicate
79        roles.sort();
80        roles.dedup();
81
82        roles
83    }
84
85    /// Map an identity to a single database role (primary)
86    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    /// Check if identity has a specific permission
92    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    /// Get all permissions for an identity
107    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    /// Add a static role assignment
123    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    /// Remove a static role assignment
135    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    /// Get anonymous role
143    pub fn anonymous_role(&self) -> Option<&String> {
144        self.anonymous_role.as_ref()
145    }
146
147    /// Evaluate a mapping rule against an identity
148    fn evaluate_rule(&self, rule: &RoleMappingRule, identity: &Identity) -> bool {
149        // All conditions must match
150        for condition in &rule.conditions {
151            if !self.evaluate_condition(condition, identity) {
152                return false;
153            }
154        }
155        true
156    }
157
158    /// Evaluate a single condition
159    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 // Just check claim exists
168                        }
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    /// Simple pattern matching (supports * wildcard)
216    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
239/// Role mapper builder
240pub 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    /// Create a new builder
249    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    /// Add a mapping rule
259    pub fn rule(mut self, rule: RoleMappingRule) -> Self {
260        self.rules.push(rule);
261        self
262    }
263
264    /// Add a group to role mapping
265    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    /// Add a default role
277    pub fn default_role(mut self, role: impl Into<String>) -> Self {
278        self.default_roles.push(role.into());
279        self
280    }
281
282    /// Set anonymous role
283    pub fn anonymous_role(mut self, role: impl Into<String>) -> Self {
284        self.anonymous_role = Some(role.into());
285        self
286    }
287
288    /// Build the role mapper
289    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/// Permission set for authorization
307#[derive(Debug, Clone)]
308pub struct PermissionSet {
309    /// Allowed databases
310    pub databases: Vec<String>,
311
312    /// Allowed schemas
313    pub schemas: Vec<String>,
314
315    /// Allowed tables (database.schema.table patterns)
316    pub tables: Vec<String>,
317
318    /// Allowed operations
319    pub operations: Vec<Operation>,
320
321    /// Row-level security predicates
322    pub row_predicates: HashMap<String, String>,
323
324    /// Column restrictions
325    pub column_restrictions: HashMap<String, Vec<String>>,
326}
327
328/// Database operation types
329#[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    /// Parse operation from string
345    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    /// Create an empty permission set
364    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    /// Create a full access permission set
376    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    /// Check if operation is allowed on table
388    pub fn is_operation_allowed(&self, operation: &Operation, table: &str) -> bool {
389        // Check operation
390        if !self.operations.contains(&Operation::All) && !self.operations.contains(operation) {
391            return false;
392        }
393
394        // Check table access
395        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            // Simple wildcard matching
405            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    /// Get row predicate for a table
417    pub fn row_predicate(&self, table: &str) -> Option<&String> {
418        self.row_predicates.get(table)
419    }
420
421    /// Get allowed columns for a table
422    pub fn allowed_columns(&self, table: &str) -> Option<&Vec<String>> {
423        self.column_restrictions.get(table)
424    }
425}
426
427/// Authorization context
428#[derive(Debug, Clone)]
429pub struct AuthorizationContext {
430    /// User identity
431    pub identity: Identity,
432
433    /// Mapped roles
434    pub roles: Vec<String>,
435
436    /// Permission set
437    pub permissions: PermissionSet,
438
439    /// Session start time
440    pub session_start: chrono::DateTime<chrono::Utc>,
441
442    /// Additional context
443    pub context: HashMap<String, String>,
444}
445
446impl AuthorizationContext {
447    /// Create a new authorization context
448    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    /// Check if operation is allowed
459    pub fn is_allowed(&self, operation: &Operation, table: &str) -> bool {
460        self.permissions.is_operation_allowed(operation, table)
461    }
462
463    /// Add context value
464    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())); // From identity
503        assert!(roles.contains(&"db_developer".to_string())); // From group
504    }
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        // Empty identity with no roles
677        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}