phonograph/phono-backends/src/pg_role.rs

use chrono::{DateTime, Utc};
use sqlx::{postgres::types::Oid, prelude::FromRow, query_as};
use thiserror::Error;
use uuid::Uuid;

use crate::{client::WorkspaceClient, rolnames::ROLE_PREFIX_USER};

#[derive(Clone, Debug, Eq, Hash, FromRow, PartialEq)]
pub struct PgRole {
    /// ID of role
    pub oid: Oid,

    /// Role name
    pub rolname: String,

    /// Role has superuser privileges
    pub rolsuper: bool,

    /// Role automatically inherits privileges of roles it is a member of
    pub rolinherit: bool,

    /// Role can create more roles
    pub rolcreaterole: bool,

    /// Role can create databases
    pub rolcreatedb: bool,

    /// Role can log in. That is, this role can be given as the initial session authorization identifier
    pub rolcanlogin: bool,

    /// Role is a replication role. A replication role can initiate replication connections and create and drop replication slots.
    pub rolreplication: bool,

    /// For roles that can log in, this sets maximum number of concurrent connections this role can make. -1 means no limit.
    pub rolconnlimit: i32,

    /// Password expiry time (only used for password authentication); null if no expiration
    pub rolvaliduntil: Option<DateTime<Utc>>,

    /// Role bypasses every row-level security policy, see Section 5.9 for more information.
    pub rolbypassrls: bool,
}

impl PgRole {
    pub fn with_name_in(names: Vec<String>) -> WithNameInQuery {
        WithNameInQuery { names }
    }

    pub fn with_name_starting_with(prefix: String) -> WithNameStartingWithQuery {
        WithNameStartingWithQuery { prefix }
    }
}

#[derive(Clone, Debug)]
pub struct WithNameInQuery {
    names: Vec<String>,
}

impl WithNameInQuery {
    pub async fn fetch_all(
        &self,
        client: &mut WorkspaceClient,
    ) -> Result<Vec<PgRole>, sqlx::Error> {
        query_as!(
            PgRole,
            r#"
select
    oid as "oid!",
    rolname as "rolname!",
    rolsuper as "rolsuper!",
    rolinherit as "rolinherit!",
    rolcreaterole as "rolcreaterole!",
    rolcreatedb as "rolcreatedb!",
    rolcanlogin as "rolcanlogin!",
    rolreplication as "rolreplication!",
    rolconnlimit as "rolconnlimit!",
    rolvaliduntil,
    rolbypassrls as "rolbypassrls!"
from pg_roles
where rolname = any($1)
"#,
            self.names.as_slice()
        )
        .fetch_all(client.get_conn())
        .await
    }
}

#[derive(Clone, Debug)]
pub struct WithNameStartingWithQuery {
    prefix: String,
}

impl WithNameStartingWithQuery {
    pub async fn fetch_all(
        &self,
        client: &mut WorkspaceClient,
    ) -> Result<Vec<PgRole>, sqlx::Error> {
        query_as!(
            PgRole,
            r#"
select
    oid as "oid!",
    rolname as "rolname!",
    rolsuper as "rolsuper!",
    rolinherit as "rolinherit!",
    rolcreaterole as "rolcreaterole!",
    rolcreatedb as "rolcreatedb!",
    rolcanlogin as "rolcanlogin!",
    rolreplication as "rolreplication!",
    rolconnlimit as "rolconnlimit!",
    rolvaliduntil,
    rolbypassrls as "rolbypassrls!"
from pg_roles
where starts_with(rolname, $1)
"#,
            self.prefix
        )
        .fetch_all(client.get_conn())
        .await
    }
}

/// A representation of role grants, starting from a single role and traversing
/// the full set (as visible via the current DB connection) of either more
/// specific roles which are granted to it, or more general roles to which it
/// is granted. This is useful, for example, for enumerating permissions which
/// are inherited rather than granted directly to a specific user.
#[derive(Clone, Debug)]
pub struct RoleTree {
    pub role: PgRole,
    pub branches: Vec<RoleTree>,
    pub inherit: bool,
}

#[derive(Clone, Debug, FromRow)]
struct RoleTreeRow {
    #[sqlx(flatten)]
    role: PgRole,
    branch: Option<Oid>,
    inherit: bool,
}

impl RoleTree {
    pub fn members_of_oid(role_oid: Oid) -> MembersOfOidQuery {
        MembersOfOidQuery { role: role_oid }
    }

    pub fn members_of_rolname(rolname: &str) -> MembersOfRolnameQuery {
        MembersOfRolnameQuery {
            role: rolname.to_owned(),
        }
    }

    pub fn granted_to(role_oid: Oid) -> GrantedToQuery {
        GrantedToQuery { role_oid }
    }

    pub fn granted_to_rolname<'a>(rolname: &'a str) -> GrantedToRolnameQuery<'a> {
        GrantedToRolnameQuery { rolname }
    }

    pub fn flatten_inherited(self) -> Vec<PgRole> {
        [
            vec![self.role],
            self.branches
                .into_iter()
                .filter(|member| member.inherit)
                .map(|member| member.flatten_inherited())
                .collect::<Vec<_>>()
                .concat(),
        ]
        .concat()
    }
}

#[derive(Clone, Debug)]
pub struct MembersOfOidQuery {
    role: Oid,
}
impl MembersOfOidQuery {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        let rows: Vec<RoleTreeRow> = query_as(
            "
        with recursive cte as (
            select $1::regrole::oid as roleid, null::oid as branch, true as inherit
            union all
                select m.member, m.roleid, c.inherit and m.inherit_option
                from cte as c
                join pg_auth_members m on m.roleid = c.roleid
        )
        select pg_roles.*, branch, inherit
        from (
            select roleid, branch, bool_or(inherit) as inherit
            from cte
            group by roleid, branch
        ) as subquery
            join pg_roles on pg_roles.oid = subquery.roleid
        ",
        )
        .bind(self.role)
        .fetch_all(client.get_conn())
        .await?;
        Ok(rows
            .iter()
            .find(|row| row.branch.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                inherit: root_row.inherit,
            }))
    }
}

#[derive(Clone, Debug)]
pub struct MembersOfRolnameQuery {
    role: String,
}

impl MembersOfRolnameQuery {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        // This could almost be a macro to DRY with MembersOfOidQuery, except
        // for the extra ::text:: cast required on the parameter in this query.
        let rows: Vec<RoleTreeRow> = query_as(
            "
        with recursive cte as (
            select $1::text::regrole::oid as roleid, null::oid as branch, true as inherit
            union all
                select m.member, m.roleid, c.inherit and m.inherit_option
                from cte as c
                join pg_auth_members m on m.roleid = c.roleid
        )
        select pg_roles.*, branch, inherit
        from (
            select roleid, branch, bool_or(inherit) as inherit
            from cte
            group by roleid, branch
        ) as subquery
            join pg_roles on pg_roles.oid = subquery.roleid
        ",
        )
        .bind(self.role.as_str() as &str)
        .fetch_all(client.get_conn())
        .await?;
        Ok(rows
            .iter()
            .find(|row| row.branch.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                inherit: root_row.inherit,
            }))
    }
}

#[derive(Clone, Debug)]
pub struct GrantedToQuery {
    role_oid: Oid,
}

impl GrantedToQuery {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        let rows: Vec<RoleTreeRow> = query_as(
            "
with recursive cte as (
    select $1 as roleid, null::oid as branch, true as inherit
    union all
        select m.roleid, m.member as branch, c.inherit and m.inherit_option
        from cte as c
        join pg_auth_members m on m.member = c.roleid
)
select pg_roles.*, branch, inherit
from (
    select roleid, branch, bool_or(inherit) as inherit
    from cte
    group by roleid, branch
) as subquery
    join pg_roles on pg_roles.oid = subquery.roleid
",
        )
        .bind(self.role_oid)
        .fetch_all(client.get_conn())
        .await?;
        Ok(rows
            .iter()
            .find(|row| row.branch.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                inherit: root_row.inherit,
            }))
    }
}

#[derive(Clone, Copy, Debug)]
pub struct GrantedToRolnameQuery<'a> {
    rolname: &'a str,
}

impl<'a> GrantedToRolnameQuery<'a> {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        let rows: Vec<RoleTreeRow> = query_as(
            "
with recursive cte as (
    select $1::regrole::oid as roleid, null::oid as branch, true as inherit
    union all
        select m.roleid, m.member as branch, c.inherit and m.inherit_option
        from cte as c
        join pg_auth_members m on m.member = c.roleid
)
select pg_roles.*, branch, inherit
from (
    select roleid, branch, bool_or(inherit) as inherit
    from cte
    group by roleid, branch
) as subquery
    join pg_roles on pg_roles.oid = subquery.roleid
",
        )
        .bind(self.rolname)
        .fetch_all(client.get_conn())
        .await?;
        Ok(rows
            .iter()
            .find(|row| row.branch.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                inherit: root_row.inherit,
            }))
    }
}

fn compute_members(rows: &Vec<RoleTreeRow>, root: Oid) -> Vec<RoleTree> {
    rows.iter()
        .filter(|row| row.branch == Some(root))
        .map(|row| RoleTree {
            role: row.role.clone(),
            branches: compute_members(rows, row.role.oid),
            inherit: row.inherit,
        })
        .collect()
}

#[derive(Debug, Error)]
pub enum RolnameParseError {
    #[error("rolname does not have phono user prefix")]
    MissingPrefix,
    #[error("unable to parse uuid from rolname: {0}")]
    BadUuid(uuid::Error),
}

pub fn user_id_from_rolname(rolname: &str) -> Result<Uuid, RolnameParseError> {
    if rolname.starts_with(ROLE_PREFIX_USER) {
        let mut rolname = rolname.to_owned();
        rolname.replace_range(0..ROLE_PREFIX_USER.len(), "");
        Uuid::parse_str(&rolname).map_err(RolnameParseError::BadUuid)
    } else {
        Err(RolnameParseError::MissingPrefix)
    }
}