1//! Running as a relation's owner, as a SECURITY DEFINER function does.
2//!
3//! The answer cache is readable and writable only by its owner: a role
4//! that could write it could forge answers, as `jev.mock_response` would
5//! (contract *Cost and safety*). A role granted EXECUTE on `jev_prob`
6//! must still use it, so the scan reaches it as the owner.
7
8use std::ffi::CStr;
9
10use pgrx::{pg_sys, spi::Spi};
11
12/// Runs `f` as the owner of `relation`. The previous user is restored on
13/// return and on unwind; an ERROR's abort restores it too.
14pub fn as_owner_of<T>(relation: &CStr, f: impl FnOnce() -> T) -> Result<T, String> {
15    let name = relation.to_str().map_err(|e| e.to_string())?;
16    let owner = Spi::get_one_with_args::<pg_sys::Oid>(
17        "SELECT relowner FROM pg_class WHERE oid = to_regclass($1)",
18        &[name.into()],
19    )
20    .map_err(|e| e.to_string())?
21    .ok_or_else(|| format!("{name} does not exist"))?;
22    let mut previous = pg_sys::InvalidOid;
23    let mut context = 0;
24    // SAFETY: the pair a SECURITY DEFINER call makes (fmgr.c), on the
25    // backend thread; `Restore` undoes it exactly once.
26    unsafe {
27        pg_sys::GetUserIdAndSecContext(&mut previous, &mut context);
28        pg_sys::SetUserIdAndSecContext(owner, context | pg_sys::SECURITY_LOCAL_USERID_CHANGE as i32);
29    }
30    let _restore = Restore { previous, context };
31    Ok(f())
32}
33
34struct Restore {
35    previous: pg_sys::Oid,
36    context: i32,
37}
38
39impl Drop for Restore {
40    fn drop(&mut self) {
41        // SAFETY: restores what `as_owner_of` read.
42        unsafe { pg_sys::SetUserIdAndSecContext(self.previous, self.context) }
43    }
44}