1use std::collections::{BTreeMap, BTreeSet};
10
11use crate::ast::{GrantSchemaStmt, RoleAttribute, SchemaPrivilege, SchemaRevokeBehavior};
12use crate::SQLError;
13use uqa_core::catalog_schema::{SchemaAclEntry, SchemaPrivileges};
14
15use super::SchemaSecurity;
16use crate::catalog::roles::{role_inherits, RoleDefinition, RoleMembership, RoleMembershipKey};
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
19pub enum SchemaAclPrivilege {
20 Usage,
21 Create,
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct SchemaPrivilegeCheck {
26 pub privilege: SchemaAclPrivilege,
27 pub grant_option: bool,
28}
29
30impl SchemaAclPrivilege {
31 const fn mask(self) -> SchemaPrivileges {
32 match self {
33 Self::Usage => SchemaPrivileges {
34 usage: true,
35 create: false,
36 },
37 Self::Create => SchemaPrivileges {
38 usage: false,
39 create: true,
40 },
41 }
42 }
43}
44
45pub fn requested_acl_privileges(
46 requested: &[SchemaPrivilege],
47) -> Result<Vec<SchemaAclPrivilege>, SQLError> {
48 requested
49 .iter()
50 .map(|privilege| match privilege {
51 SchemaPrivilege::Usage => Ok(SchemaAclPrivilege::Usage),
52 SchemaPrivilege::Create => Ok(SchemaAclPrivilege::Create),
53 SchemaPrivilege::Unsupported(name) => Err(SQLError::Routine {
54 sqlstate: "0LP01".into(),
55 message: format!("invalid privilege type {name} for schema"),
56 }),
57 })
58 .collect()
59}
60
61pub fn parse_privilege_checks(value: &str) -> Result<Vec<SchemaPrivilegeCheck>, SQLError> {
62 value
63 .split(',')
64 .map(|item| {
65 let item = item.trim();
66 let upper = item.to_ascii_uppercase();
67 let (name, grant_option) = upper
68 .strip_suffix(" WITH GRANT OPTION")
69 .map_or((upper.as_str(), false), |name| (name.trim_end(), true));
70 let privilege = match name {
71 "USAGE" => SchemaAclPrivilege::Usage,
72 "CREATE" => SchemaAclPrivilege::Create,
73 _ => {
74 return Err(SQLError::Routine {
75 sqlstate: "22023".into(),
76 message: format!("unrecognized privilege type: \"{item}\""),
77 })
78 }
79 };
80 Ok(SchemaPrivilegeCheck {
81 privilege,
82 grant_option,
83 })
84 })
85 .collect()
86}
87
88fn acl_grantor<'a>(entry: &'a SchemaAclEntry, owner: &'a str) -> &'a str {
89 entry.grantor.as_deref().unwrap_or(owner)
90}
91
92fn materialize_acl(security: &mut SchemaSecurity) {
93 if security.acl.is_none() {
94 security.acl = Some(vec![SchemaAclEntry {
95 role: security.role_owner.clone(),
96 grantor: Some(security.role_owner.clone()),
97 privileges: SchemaPrivileges::ALL,
98 grant_options: SchemaPrivileges::default(),
99 }]);
100 }
101}
102
103fn grant_option_roles(
104 security: &SchemaSecurity,
105 privilege: SchemaAclPrivilege,
106) -> BTreeSet<String> {
107 let mut reachable = BTreeSet::from([security.role_owner.clone()]);
108 let Some(acl) = security.acl.as_ref() else {
109 return reachable;
110 };
111 loop {
112 let mut changed = false;
113 for entry in acl {
114 if entry.role != "PUBLIC"
115 && entry.grant_options.intersects(privilege.mask())
116 && reachable.contains(acl_grantor(entry, &security.role_owner))
117 {
118 changed |= reachable.insert(entry.role.clone());
119 }
120 }
121 if !changed {
122 return reachable;
123 }
124 }
125}
126
127pub fn select_acl_grantor(
128 security: &SchemaSecurity,
129 privilege: SchemaAclPrivilege,
130 current_user: &str,
131 roles: &BTreeMap<String, RoleDefinition>,
132 memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
133) -> Option<String> {
134 if role_inherits(roles, memberships, current_user, &security.role_owner) {
135 return Some(security.role_owner.clone());
136 }
137 let grant_options = grant_option_roles(security, privilege);
138 if grant_options.contains(current_user) {
139 return Some(current_user.to_string());
140 }
141 security.acl.as_ref().and_then(|acl| {
142 acl.iter()
143 .filter(|entry| entry.role != "PUBLIC" && grant_options.contains(&entry.role))
144 .find(|entry| role_inherits(roles, memberships, current_user, &entry.role))
145 .map(|entry| entry.role.clone())
146 })
147}
148
149pub fn role_has_schema_privilege(
150 security: &SchemaSecurity,
151 subject: &str,
152 privilege: SchemaAclPrivilege,
153 roles: &BTreeMap<String, RoleDefinition>,
154 memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
155) -> bool {
156 role_has_schema_privilege_check(
157 security,
158 subject,
159 SchemaPrivilegeCheck {
160 privilege,
161 grant_option: false,
162 },
163 roles,
164 memberships,
165 )
166}
167
168pub fn role_has_schema_privilege_check(
169 security: &SchemaSecurity,
170 subject: &str,
171 check: SchemaPrivilegeCheck,
172 roles: &BTreeMap<String, RoleDefinition>,
173 memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
174) -> bool {
175 if roles
176 .get(subject)
177 .is_some_and(|role| role.has(RoleAttribute::Superuser))
178 {
179 return true;
180 }
181 if check.grant_option {
182 return grant_option_roles(security, check.privilege)
183 .iter()
184 .any(|role| role_inherits(roles, memberships, subject, role));
185 }
186 match security.acl.as_ref() {
187 None => role_inherits(roles, memberships, subject, &security.role_owner),
188 Some(acl) => acl.iter().any(|entry| {
189 entry.privileges.intersects(check.privilege.mask())
190 && (entry.role == "PUBLIC"
191 || role_inherits(roles, memberships, subject, &entry.role))
192 }),
193 }
194}
195
196pub fn grant_acl(
197 security: &mut SchemaSecurity,
198 privilege: SchemaAclPrivilege,
199 grantees: &[String],
200 grantor: &str,
201 grant_option: bool,
202) {
203 materialize_acl(security);
204 let owner = security.role_owner.clone();
205 let acl = security.acl.as_mut().expect("schema ACL was materialized");
206 for grantee in grantees {
207 let position = acl
208 .iter()
209 .position(|entry| entry.role == *grantee && acl_grantor(entry, &owner) == grantor)
210 .unwrap_or_else(|| {
211 acl.push(SchemaAclEntry {
212 role: grantee.clone(),
213 grantor: Some(grantor.to_string()),
214 privileges: SchemaPrivileges::default(),
215 grant_options: SchemaPrivileges::default(),
216 });
217 acl.len() - 1
218 });
219 let entry = &mut acl[position];
220 entry.privileges.insert(privilege.mask());
221 if grant_option && grantee != "PUBLIC" && grantee != &owner {
222 entry.grant_options.insert(privilege.mask());
223 }
224 }
225}
226
227pub fn revoke_acl(
228 security: &mut SchemaSecurity,
229 privilege: SchemaAclPrivilege,
230 grantees: &[String],
231 grantor: &str,
232 grant_option_only: bool,
233 cascade: bool,
234) -> Result<(), SQLError> {
235 let before = grant_option_roles(security, privilege);
236 materialize_acl(security);
237 let owner = security.role_owner.clone();
238 let acl = security.acl.as_mut().expect("schema ACL was materialized");
239 for entry in acl
240 .iter_mut()
241 .filter(|entry| grantees.contains(&entry.role) && acl_grantor(entry, &owner) == grantor)
242 {
243 entry.grant_options.remove(privilege.mask());
244 if !grant_option_only {
245 entry.privileges.remove(privilege.mask());
246 }
247 }
248 remove_empty_entries(acl);
249 revoke_dependent_acl(security, privilege, &before, cascade)
250}
251
252fn revoke_dependent_acl(
253 security: &mut SchemaSecurity,
254 privilege: SchemaAclPrivilege,
255 before: &BTreeSet<String>,
256 cascade: bool,
257) -> Result<(), SQLError> {
258 loop {
259 let current = grant_option_roles(security, privilege);
260 let lost = before
261 .difference(¤t)
262 .cloned()
263 .collect::<BTreeSet<_>>();
264 if lost.is_empty() {
265 return Ok(());
266 }
267 let owner = security.role_owner.clone();
268 let dependent = security.acl.as_ref().is_some_and(|acl| {
269 acl.iter().any(|entry| {
270 lost.contains(acl_grantor(entry, &owner))
271 && (entry.privileges.intersects(privilege.mask())
272 || entry.grant_options.intersects(privilege.mask()))
273 })
274 });
275 if !dependent {
276 return Ok(());
277 }
278 if !cascade {
279 return Err(SQLError::Routine {
280 sqlstate: "2BP01".into(),
281 message: "dependent privileges exist".into(),
282 });
283 }
284 let acl = security
285 .acl
286 .as_mut()
287 .expect("dependent schema privileges require an explicit ACL");
288 for entry in acl
289 .iter_mut()
290 .filter(|entry| lost.contains(acl_grantor(entry, &owner)))
291 {
292 entry.privileges.remove(privilege.mask());
293 entry.grant_options.remove(privilege.mask());
294 }
295 remove_empty_entries(acl);
296 }
297}
298
299fn remove_empty_entries(acl: &mut Vec<SchemaAclEntry>) {
300 acl.retain(|entry| !entry.privileges.is_empty() || !entry.grant_options.is_empty());
301}
302
303pub fn schema_security_with_public_privileges(create: bool) -> SchemaSecurity {
304 let role_owner = "uqa".to_string();
305 SchemaSecurity {
306 role_owner: role_owner.clone(),
307 acl: Some(vec![
308 SchemaAclEntry {
309 role: role_owner.clone(),
310 grantor: Some(role_owner.clone()),
311 privileges: SchemaPrivileges::ALL,
312 grant_options: SchemaPrivileges::default(),
313 },
314 SchemaAclEntry {
315 role: "PUBLIC".into(),
316 grantor: Some(role_owner),
317 privileges: SchemaPrivileges {
318 usage: true,
319 create,
320 },
321 grant_options: SchemaPrivileges::default(),
322 },
323 ]),
324 }
325}
326
327pub fn rewrite_schema_acl_owner(security: &mut SchemaSecurity, new_owner: &str) {
328 if let Some(acl) = &mut security.acl {
329 for entry in acl.iter_mut() {
330 if entry.role == security.role_owner {
331 entry.role = new_owner.to_string();
332 }
333 if entry.grantor.as_deref().unwrap_or(&security.role_owner) == security.role_owner {
334 entry.grantor = Some(new_owner.to_string());
335 }
336 }
337 let mut merged: Vec<SchemaAclEntry> = Vec::new();
338 for entry in std::mem::take(acl) {
339 if let Some(previous) = merged
340 .iter_mut()
341 .find(|previous| previous.role == entry.role && previous.grantor == entry.grantor)
342 {
343 previous.privileges.insert(entry.privileges);
344 previous.grant_options.insert(entry.grant_options);
345 } else {
346 merged.push(entry);
347 }
348 }
349 *acl = merged;
350 }
351 security.role_owner = new_owner.to_string();
352}
353
354pub fn apply_schema_acl(
355 statement: &GrantSchemaStmt,
356 grantees: &[String],
357 privileges: &[SchemaAclPrivilege],
358 current_user: &str,
359 roles: &BTreeMap<String, RoleDefinition>,
360 memberships: &BTreeMap<RoleMembershipKey, RoleMembership>,
361 current: &SchemaSecurity,
362) -> Result<(SchemaSecurity, usize), SQLError> {
363 let grantors = privileges
364 .iter()
365 .map(|privilege| {
366 (
367 *privilege,
368 select_acl_grantor(current, *privilege, current_user, roles, memberships),
369 )
370 })
371 .collect::<Vec<_>>();
372 let grantable = grantors
373 .iter()
374 .filter(|(_, grantor)| grantor.is_some())
375 .count();
376 let mut next = current.clone();
377 for (privilege, grantor) in grantors {
378 let Some(grantor) = grantor else {
379 continue;
380 };
381 if statement.is_grant {
382 grant_acl(
383 &mut next,
384 privilege,
385 grantees,
386 &grantor,
387 statement.grant_option,
388 );
389 } else {
390 revoke_acl(
391 &mut next,
392 privilege,
393 grantees,
394 &grantor,
395 statement.grant_option_only,
396 statement.revoke_behavior == SchemaRevokeBehavior::Cascade,
397 )?;
398 }
399 }
400 Ok((next, grantable))
401}
402
403pub fn validate_schema_acl_roles(
404 statement: &GrantSchemaStmt,
405 grantees: &[String],
406 requested_grantor: Option<&str>,
407 current_user: &str,
408 roles: &BTreeMap<String, RoleDefinition>,
409) -> Result<(), SQLError> {
410 for role in grantees {
411 if role != "PUBLIC" && !roles.contains_key(role) {
412 return Err(SQLError::Routine {
413 sqlstate: "42704".into(),
414 message: format!("role \"{role}\" does not exist"),
415 });
416 }
417 }
418 if statement.is_grant && statement.grant_option && grantees.iter().any(|role| role == "PUBLIC")
419 {
420 return Err(SQLError::Routine {
421 sqlstate: "0LP01".into(),
422 message: "grant options can only be granted to roles".into(),
423 });
424 }
425 if let Some(requested_grantor) = requested_grantor {
426 if !roles.contains_key(requested_grantor) {
427 return Err(SQLError::Routine {
428 sqlstate: "42704".into(),
429 message: format!("role \"{requested_grantor}\" does not exist"),
430 });
431 }
432 if requested_grantor != current_user {
433 return Err(SQLError::Routine {
434 sqlstate: "0A000".into(),
435 message: "grantor must be current user".into(),
436 });
437 }
438 }
439 Ok(())
440}
441
442pub fn schema_acl_warning(is_grant: bool, partial: bool, name: &str) -> (&'static str, String) {
443 let message = match (is_grant, partial) {
444 (true, true) => format!("not all privileges were granted for \"{name}\""),
445 (true, false) => format!("no privileges were granted for \"{name}\""),
446 (false, true) => format!("not all privileges could be revoked for \"{name}\""),
447 (false, false) => format!("no privileges could be revoked for \"{name}\""),
448 };
449 ("WARNING", message)
450}
451
452pub fn resolve_schema_grant_targets(
453 registry: &BTreeMap<String, SchemaSecurity>,
454 schemas: &[String],
455) -> Result<Vec<String>, SQLError> {
456 let mut targets = Vec::with_capacity(schemas.len());
457 for schema in schemas {
458 if !registry.contains_key(schema) {
459 return Err(SQLError::Routine {
460 sqlstate: "3F000".into(),
461 message: format!("schema \"{schema}\" does not exist"),
462 });
463 }
464 if !targets.contains(schema) {
465 targets.push(schema.clone());
466 }
467 }
468 Ok(targets)
469}