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