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, RoleMappingCondition, RoleMappingRule};
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                && rule.permissions.contains(&permission.to_string())
98            {
99                return true;
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.entry(user_id).or_default().push(role);
129    }
130
131    /// Remove a static role assignment
132    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    /// Get anonymous role
140    pub fn anonymous_role(&self) -> Option<&String> {
141        self.anonymous_role.as_ref()
142    }
143
144    /// Evaluate a mapping rule against an identity
145    fn evaluate_rule(&self, rule: &RoleMappingRule, identity: &Identity) -> bool {
146        // All conditions must match
147        for condition in &rule.conditions {
148            if !self.evaluate_condition(condition, identity) {
149                return false;
150            }
151        }
152        true
153    }
154
155    /// Evaluate a single condition
156    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 // Just check claim exists
165                        }
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    /// Simple pattern matching (supports * wildcard)
206    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
229/// Role mapper builder
230pub 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    /// Create a new builder
239    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    /// Add a mapping rule
249    pub fn rule(mut self, rule: RoleMappingRule) -> Self {
250        self.rules.push(rule);
251        self
252    }
253
254    /// Add a group to role mapping
255    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    /// Add a default role
264    pub fn default_role(mut self, role: impl Into<String>) -> Self {
265        self.default_roles.push(role.into());
266        self
267    }
268
269    /// Set anonymous role
270    pub fn anonymous_role(mut self, role: impl Into<String>) -> Self {
271        self.anonymous_role = Some(role.into());
272        self
273    }
274
275    /// Build the role mapper
276    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/// Permission set for authorization
294#[derive(Debug, Clone)]
295pub struct PermissionSet {
296    /// Allowed databases
297    pub databases: Vec<String>,
298
299    /// Allowed schemas
300    pub schemas: Vec<String>,
301
302    /// Allowed tables (database.schema.table patterns)
303    pub tables: Vec<String>,
304
305    /// Allowed operations
306    pub operations: Vec<Operation>,
307
308    /// Row-level security predicates
309    pub row_predicates: HashMap<String, String>,
310
311    /// Column restrictions
312    pub column_restrictions: HashMap<String, Vec<String>>,
313}
314
315/// Database operation types
316#[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    /// Parse operation from string
332    #[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    /// Create an empty permission set
352    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    /// Create a full access permission set
364    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    /// Check if operation is allowed on table
376    pub fn is_operation_allowed(&self, operation: &Operation, table: &str) -> bool {
377        // Check operation
378        if !self.operations.contains(&Operation::All) && !self.operations.contains(operation) {
379            return false;
380        }
381
382        // Check table access
383        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            // Simple wildcard matching
393            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    /// Get row predicate for a table
405    pub fn row_predicate(&self, table: &str) -> Option<&String> {
406        self.row_predicates.get(table)
407    }
408
409    /// Get allowed columns for a table
410    pub fn allowed_columns(&self, table: &str) -> Option<&Vec<String>> {
411        self.column_restrictions.get(table)
412    }
413}
414
415/// Authorization context
416#[derive(Debug, Clone)]
417pub struct AuthorizationContext {
418    /// User identity
419    pub identity: Identity,
420
421    /// Mapped roles
422    pub roles: Vec<String>,
423
424    /// Permission set
425    pub permissions: PermissionSet,
426
427    /// Session start time
428    pub session_start: chrono::DateTime<chrono::Utc>,
429
430    /// Additional context
431    pub context: HashMap<String, String>,
432}
433
434impl AuthorizationContext {
435    /// Create a new authorization context
436    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    /// Check if operation is allowed
447    pub fn is_allowed(&self, operation: &Operation, table: &str) -> bool {
448        self.permissions.is_operation_allowed(operation, table)
449    }
450
451    /// Add context value
452    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())); // From identity
492        assert!(roles.contains(&"db_developer".to_string())); // From group
493    }
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        // Empty identity with no roles
660        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}