uqa_sql/catalog/roles/
memberships.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use super::{RoleDefinition, RoleMembership, RoleMembershipKey};
8use crate::ast::{GrantRoleStmt, RoleAttribute, RoleMembershipOptions};
9use crate::SQLError;
10use std::collections::{BTreeMap, BTreeSet, VecDeque};
11use uqa_core::Value;
12
13pub fn role_is_superuser(roles: &BTreeMap<String, RoleDefinition>, role: &str) -> bool {
14    roles
15        .get(role)
16        .is_some_and(|definition| definition.has(RoleAttribute::Superuser))
17}
18
19pub fn require_role_attribute_authority(
20    roles: &BTreeMap<String, RoleDefinition>,
21    current: &str,
22    attributes: impl IntoIterator<Item = RoleAttribute>,
23    action: &str,
24) -> Result<(), SQLError> {
25    let current_role = roles.get(current).ok_or_else(|| undefined_role(current))?;
26    if current_role.has(RoleAttribute::Superuser) {
27        return Ok(());
28    }
29    for attribute in attributes {
30        let restricted = matches!(
31            attribute,
32            RoleAttribute::Superuser
33                | RoleAttribute::CreateRole
34                | RoleAttribute::CreateDb
35                | RoleAttribute::Replication
36                | RoleAttribute::BypassRls
37        );
38        if restricted && !current_role.has(attribute) {
39            return Err(insufficient_privilege(&format!(
40                "permission denied to {action}"
41            )));
42        }
43    }
44    Ok(())
45}
46
47pub fn role_has_admin(
48    memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
49    member: &str,
50    role: &str,
51) -> bool {
52    memberships.values().any(|membership| {
53        membership.member == member && membership.role == role && membership.admin_option
54    })
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum RolePrivilegeCheck {
59    Member,
60    Usage,
61    Set,
62    Admin,
63}
64
65pub fn resolve_pg_has_role_identifier(
66    value: &Value,
67    roles: &BTreeMap<String, RoleDefinition>,
68) -> Result<Option<String>, SQLError> {
69    match value {
70        Value::Str(name) | Value::FixedChar(name) => {
71            if roles.contains_key(name) {
72                Ok(Some(name.clone()))
73            } else {
74                Err(undefined_role(name))
75            }
76        }
77        Value::Int(oid) => Ok(roles
78            .values()
79            .find(|role| role.oid == *oid)
80            .map(|role| role.name.clone())),
81        _ => Err(SQLError::TypeMismatch(
82            "pg_has_role role arguments must be name or oid".into(),
83        )),
84    }
85}
86
87pub fn role_privilege_text(value: &Value) -> Result<&str, SQLError> {
88    match value {
89        Value::Str(privilege) | Value::FixedChar(privilege) => Ok(privilege),
90        _ => Err(SQLError::TypeMismatch(
91            "pg_has_role privilege argument must be text".into(),
92        )),
93    }
94}
95
96pub fn parse_pg_has_role_privileges(privileges: &str) -> Result<Vec<RolePrivilegeCheck>, SQLError> {
97    privileges
98        .split(',')
99        .map(|privilege| {
100            let privilege = privilege.trim();
101            if [
102                "MEMBER WITH ADMIN OPTION",
103                "MEMBER WITH GRANT OPTION",
104                "USAGE WITH ADMIN OPTION",
105                "USAGE WITH GRANT OPTION",
106                "SET WITH ADMIN OPTION",
107                "SET WITH GRANT OPTION",
108            ]
109            .iter()
110            .any(|candidate| privilege.eq_ignore_ascii_case(candidate))
111            {
112                return Ok(RolePrivilegeCheck::Admin);
113            }
114            if privilege.eq_ignore_ascii_case("MEMBER") {
115                Ok(RolePrivilegeCheck::Member)
116            } else if privilege.eq_ignore_ascii_case("USAGE") {
117                Ok(RolePrivilegeCheck::Usage)
118            } else if privilege.eq_ignore_ascii_case("SET") {
119                Ok(RolePrivilegeCheck::Set)
120            } else {
121                Err(SQLError::Routine {
122                    sqlstate: "22023".into(),
123                    message: format!("unrecognized privilege type: \"{privilege}\""),
124                })
125            }
126        })
127        .collect()
128}
129
130pub fn pg_has_role_privilege(
131    roles: &BTreeMap<String, RoleDefinition>,
132    memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
133    subject: Option<&str>,
134    target: Option<&str>,
135    privilege: RolePrivilegeCheck,
136) -> bool {
137    let Some(subject) = subject else {
138        return false;
139    };
140    if role_is_superuser(roles, subject) {
141        return true;
142    }
143    let Some(target) = target else {
144        return false;
145    };
146    match privilege {
147        RolePrivilegeCheck::Member => role_reaches(memberships, subject, target, |_| true),
148        RolePrivilegeCheck::Usage => role_inherits(roles, memberships, subject, target),
149        RolePrivilegeCheck::Set => role_can_set(roles, memberships, subject, target),
150        RolePrivilegeCheck::Admin => role_has_transitive_admin(memberships, subject, target),
151    }
152}
153
154pub fn role_has_transitive_admin(
155    memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
156    member: &str,
157    role: &str,
158) -> bool {
159    let mut queue = VecDeque::from([member.to_string()]);
160    let mut visited = BTreeSet::from([member.to_string()]);
161    while let Some(current) = queue.pop_front() {
162        for membership in memberships
163            .values()
164            .filter(|membership| membership.member == current)
165        {
166            if membership.role == role && membership.admin_option {
167                return true;
168            }
169            if visited.insert(membership.role.clone()) {
170                queue.push_back(membership.role.clone());
171            }
172        }
173    }
174    false
175}
176
177pub fn role_reaches(
178    memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
179    member: &str,
180    role: &str,
181    usable: impl Fn(&RoleMembership) -> bool,
182) -> bool {
183    if member == role {
184        return true;
185    }
186    let mut queue = VecDeque::from([member.to_string()]);
187    let mut visited = BTreeSet::from([member.to_string()]);
188    while let Some(current) = queue.pop_front() {
189        for membership in memberships
190            .values()
191            .filter(|membership| membership.member == current && usable(membership))
192        {
193            if membership.role == role {
194                return true;
195            }
196            if visited.insert(membership.role.clone()) {
197                queue.push_back(membership.role.clone());
198            }
199        }
200    }
201    false
202}
203
204pub fn role_can_set(
205    roles: &BTreeMap<String, RoleDefinition>,
206    memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
207    member: &str,
208    role: &str,
209) -> bool {
210    role_is_superuser(roles, member)
211        || role_reaches(memberships, member, role, |membership| {
212            membership.set_option
213        })
214}
215
216pub fn role_inherits(
217    roles: &BTreeMap<String, RoleDefinition>,
218    memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
219    member: &str,
220    role: &str,
221) -> bool {
222    role_is_superuser(roles, member)
223        || role_reaches(memberships, member, role, |membership| {
224            membership.inherit_option
225        })
226}
227
228pub fn membership_error(message: impl Into<String>) -> SQLError {
229    SQLError::Routine {
230        sqlstate: "0LP01".into(),
231        message: message.into(),
232    }
233}
234
235pub fn undefined_role(name: &str) -> SQLError {
236    SQLError::Routine {
237        sqlstate: "42704".into(),
238        message: format!("role \"{name}\" does not exist"),
239    }
240}
241
242pub fn apply_grant_role_statement(
243    roles: &BTreeMap<String, RoleDefinition>,
244    memberships: &mut BTreeMap<RoleMembershipKey, RoleMembership>,
245    current: &str,
246    statement: &GrantRoleStmt,
247) -> Result<(), SQLError> {
248    for role in statement
249        .granted_roles
250        .iter()
251        .chain(statement.grantee_roles.iter())
252    {
253        if !roles.contains_key(role) {
254            return Err(undefined_role(role));
255        }
256    }
257    let grantor = statement.grantor.as_deref().unwrap_or(current);
258    if !roles.contains_key(grantor) {
259        return Err(undefined_role(grantor));
260    }
261    if statement.grantor.is_some() && !role_can_set(roles, memberships, current, grantor) {
262        return Err(insufficient_privilege(&format!(
263            "permission denied to grant privileges as role \"{grantor}\""
264        )));
265    }
266    for role in &statement.granted_roles {
267        let superuser_revoke = !statement.is_grant && role_is_superuser(roles, current);
268        if !superuser_revoke
269            && !role_is_superuser(roles, grantor)
270            && !role_has_admin(memberships, grantor, role)
271        {
272            return Err(insufficient_privilege(&format!(
273                "permission denied to {} role \"{role}\"",
274                if statement.is_grant {
275                    "grant"
276                } else {
277                    "revoke"
278                }
279            )));
280        }
281    }
282    for role in &statement.granted_roles {
283        for member in &statement.grantee_roles {
284            let key = RoleMembershipKey {
285                role: role.clone(),
286                member: member.clone(),
287                grantor: grantor.to_string(),
288            };
289            if statement.is_grant {
290                if role_reaches(memberships, role, member, |_| true) {
291                    return Err(membership_error(format!(
292                        "role \"{role}\" is a member of role \"{member}\""
293                    )));
294                }
295                insert_membership(memberships, role, member, grantor, statement.options, roles);
296            } else if statement.options == RoleMembershipOptions::default() {
297                revoke_membership(memberships, &key, statement.cascade, true)?;
298            } else if let Some(existing) = memberships.get(&key).cloned() {
299                if statement.options.admin == Some(false) && existing.admin_option {
300                    clear_membership_admin(memberships, &key, statement.cascade)?;
301                }
302                if let Some(membership) = memberships.get_mut(&key) {
303                    if statement.options.inherit == Some(false) {
304                        membership.inherit_option = false;
305                    }
306                    if statement.options.set == Some(false) {
307                        membership.set_option = false;
308                    }
309                }
310            }
311        }
312    }
313    Ok(())
314}
315
316pub fn insert_membership(
317    memberships: &mut BTreeMap<RoleMembershipKey, RoleMembership>,
318    role: &str,
319    member: &str,
320    grantor: &str,
321    options: RoleMembershipOptions,
322    roles: &BTreeMap<String, RoleDefinition>,
323) {
324    let key = RoleMembershipKey {
325        role: role.to_string(),
326        member: member.to_string(),
327        grantor: grantor.to_string(),
328    };
329    if let Some(existing) = memberships.get_mut(&key) {
330        if let Some(value) = options.admin {
331            existing.admin_option = value;
332        }
333        if let Some(value) = options.inherit {
334            existing.inherit_option = value;
335        }
336        if let Some(value) = options.set {
337            existing.set_option = value;
338        }
339        return;
340    }
341    let oid = allocate_role_membership_oid(memberships, &key);
342    memberships.insert(
343        key,
344        RoleMembership {
345            oid,
346            role: role.to_string(),
347            member: member.to_string(),
348            grantor: grantor.to_string(),
349            admin_option: options.admin.unwrap_or(false),
350            inherit_option: options.inherit.unwrap_or_else(|| {
351                roles
352                    .get(member)
353                    .is_some_and(|role| role.has(RoleAttribute::Inherit))
354            }),
355            set_option: options.set.unwrap_or(true),
356        },
357    );
358}
359
360pub fn allocate_role_membership_oid(
361    memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
362    key: &RoleMembershipKey,
363) -> i64 {
364    let mut hash = 14_695_981_039_346_656_037_u64;
365    for part in [&key.role, &key.member, &key.grantor] {
366        for byte in part.as_bytes().iter().copied().chain(std::iter::once(0)) {
367            hash ^= u64::from(byte);
368            hash = hash.wrapping_mul(1_099_511_628_211);
369        }
370    }
371    let mut oid = 2_500_000_000_i64 + i64::try_from(hash % 1_500_000_000).unwrap_or(0);
372    while memberships.values().any(|membership| membership.oid == oid) {
373        oid = if oid == 3_999_999_999 {
374            2_500_000_000
375        } else {
376            oid + 1
377        };
378    }
379    oid
380}
381
382pub fn clear_membership_admin(
383    memberships: &mut BTreeMap<RoleMembershipKey, RoleMembership>,
384    key: &RoleMembershipKey,
385    cascade: bool,
386) -> Result<(), SQLError> {
387    let Some(existing) = memberships.get_mut(key) else {
388        return Ok(());
389    };
390    existing.admin_option = false;
391    revoke_dependent_memberships(memberships, &key.role, &key.member, cascade)
392}
393
394pub fn revoke_membership(
395    memberships: &mut BTreeMap<RoleMembershipKey, RoleMembership>,
396    key: &RoleMembershipKey,
397    cascade: bool,
398    check_dependents: bool,
399) -> Result<(), SQLError> {
400    let Some(existing) = memberships.remove(key) else {
401        return Ok(());
402    };
403    if check_dependents && existing.admin_option {
404        revoke_dependent_memberships(memberships, &existing.role, &existing.member, cascade)?;
405    }
406    Ok(())
407}
408
409pub fn revoke_dependent_memberships(
410    memberships: &mut BTreeMap<RoleMembershipKey, RoleMembership>,
411    role: &str,
412    former_admin: &str,
413    cascade: bool,
414) -> Result<(), SQLError> {
415    if role_has_admin(memberships, former_admin, role) {
416        return Ok(());
417    }
418    let dependent = memberships
419        .iter()
420        .filter(|(_, membership)| membership.role == role && membership.grantor == former_admin)
421        .map(|(key, _)| key.clone())
422        .collect::<Vec<_>>();
423    if dependent.is_empty() {
424        return Ok(());
425    }
426    if !cascade {
427        return Err(SQLError::Routine {
428            sqlstate: "2BP01".into(),
429            message: "dependent privileges exist".into(),
430        });
431    }
432    for key in dependent {
433        revoke_membership(memberships, &key, true, true)?;
434    }
435    Ok(())
436}
437
438pub fn insufficient_privilege(message: &str) -> SQLError {
439    SQLError::Routine {
440        sqlstate: "42501".into(),
441        message: message.into(),
442    }
443}
444
445#[cfg(test)]
446mod tests {
447    use super::super::role_oid;
448    use super::*;
449
450    fn role(name: &str) -> RoleDefinition {
451        RoleDefinition {
452            oid: role_oid(name),
453            name: name.into(),
454            attributes: BTreeSet::new(),
455            connection_limit: -1,
456        }
457    }
458
459    fn membership(
460        memberships: &mut BTreeMap<RoleMembershipKey, RoleMembership>,
461        role: &str,
462        member: &str,
463        admin: bool,
464        inherit: bool,
465        set: bool,
466    ) {
467        let key = RoleMembershipKey {
468            role: role.into(),
469            member: member.into(),
470            grantor: "uqa".into(),
471        };
472        memberships.insert(
473            key.clone(),
474            RoleMembership {
475                oid: role_oid(&format!("{role}/{member}")),
476                role: key.role,
477                member: key.member,
478                grantor: key.grantor,
479                admin_option: admin,
480                inherit_option: inherit,
481                set_option: set,
482            },
483        );
484    }
485
486    #[test]
487    fn pg_has_role_privilege_names_include_lists_and_admin_aliases() {
488        assert_eq!(
489            parse_pg_has_role_privileges(" member, USAGE , set ").unwrap(),
490            vec![
491                RolePrivilegeCheck::Member,
492                RolePrivilegeCheck::Usage,
493                RolePrivilegeCheck::Set,
494            ]
495        );
496        for privilege in [
497            "MEMBER WITH ADMIN OPTION",
498            "USAGE WITH GRANT OPTION",
499            "SET WITH ADMIN OPTION",
500        ] {
501            assert_eq!(
502                parse_pg_has_role_privileges(privilege).unwrap(),
503                vec![RolePrivilegeCheck::Admin]
504            );
505        }
506        assert_eq!(
507            parse_pg_has_role_privileges("ADMIN")
508                .unwrap_err()
509                .sqlstate(),
510            Some("22023")
511        );
512    }
513
514    #[test]
515    fn pg_has_role_checks_member_usage_set_and_transitive_admin_independently() {
516        let roles = [
517            "parent",
518            "middle",
519            "leaf",
520            "noinherit",
521            "admin",
522            "admin_leaf",
523        ]
524        .into_iter()
525        .map(|name| (name.into(), role(name)))
526        .chain([("uqa".into(), RoleDefinition::bootstrap())])
527        .collect::<BTreeMap<_, _>>();
528        let mut memberships = BTreeMap::new();
529        membership(&mut memberships, "parent", "middle", false, true, false);
530        membership(&mut memberships, "middle", "leaf", false, true, true);
531        membership(&mut memberships, "parent", "noinherit", false, false, true);
532        membership(&mut memberships, "parent", "admin", true, false, false);
533        membership(&mut memberships, "admin", "admin_leaf", false, false, false);
534
535        assert!(pg_has_role_privilege(
536            &roles,
537            &memberships,
538            Some("leaf"),
539            Some("parent"),
540            RolePrivilegeCheck::Member
541        ));
542        assert!(pg_has_role_privilege(
543            &roles,
544            &memberships,
545            Some("leaf"),
546            Some("parent"),
547            RolePrivilegeCheck::Usage
548        ));
549        assert!(!pg_has_role_privilege(
550            &roles,
551            &memberships,
552            Some("leaf"),
553            Some("parent"),
554            RolePrivilegeCheck::Set
555        ));
556        assert!(!pg_has_role_privilege(
557            &roles,
558            &memberships,
559            Some("noinherit"),
560            Some("parent"),
561            RolePrivilegeCheck::Usage
562        ));
563        assert!(pg_has_role_privilege(
564            &roles,
565            &memberships,
566            Some("noinherit"),
567            Some("parent"),
568            RolePrivilegeCheck::Set
569        ));
570        assert!(pg_has_role_privilege(
571            &roles,
572            &memberships,
573            Some("admin_leaf"),
574            Some("parent"),
575            RolePrivilegeCheck::Admin
576        ));
577        assert!(!pg_has_role_privilege(
578            &roles,
579            &memberships,
580            Some("parent"),
581            Some("parent"),
582            RolePrivilegeCheck::Admin
583        ));
584        assert!(pg_has_role_privilege(
585            &roles,
586            &memberships,
587            Some("uqa"),
588            None,
589            RolePrivilegeCheck::Member
590        ));
591    }
592}
uqa_sql/catalog/roles/memberships.rs

uqa_sql/catalog/roles/
memberships.rs
