AkurAI-Build / src/access.rs

Latest change cd8019fddc46769e797dc54194decfd9d4a98f2d - OLI-270: branch protection — rules per branch, pre-receive check, merge gate by Óli

//! Collaborator roles, deploy keys and scoped expiring personal access tokens.
//! Only SHA-256 digests of tokens are stored; the plaintext is returned once.

use anyhow::{Context, Result, bail, ensure};
use rusqlite::{OptionalExtension, params};
use serde::Serialize;
use sha2::{Digest, Sha256};

use crate::db::Database;

const TOKEN_PREFIX: &str = "abt_";
const TOKEN_RANDOM_BYTES: usize = 32;
const MAX_NAME: usize = 128;
const MAX_SUBJECT: usize = 256;
/// Longest lifetime a token may be granted, in seconds (one year).
pub const MAX_TTL_SECS: i64 = 365 * 24 * 3600;

#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    Read,
    Write,
    Admin,
}

impl Role {
    pub fn parse(value: &str) -> Result<Self> {
        match value {
            "read" => Ok(Self::Read),
            "write" => Ok(Self::Write),
            "admin" => Ok(Self::Admin),
            other => bail!("unknown role {other:?}"),
        }
    }

    pub fn as_str(self) -> &'static str {
        match self {
            Self::Read => "read",
            Self::Write => "write",
            Self::Admin => "admin",
        }
    }

    pub fn allows(self, required: Role) -> bool {
        self >= required
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Kind {
    Pat,
    DeployKey,
}

impl Kind {
    fn as_str(self) -> &'static str {
        match self {
            Self::Pat => "pat",
            Self::DeployKey => "deploy_key",
        }
    }

    fn parse(value: &str) -> Result<Self> {
        match value {
            "pat" => Ok(Self::Pat),
            "deploy_key" => Ok(Self::DeployKey),
            other => bail!("unknown token kind {other:?}"),
        }
    }
}

#[derive(Clone, Debug, Serialize)]
pub struct Collaborator {
    pub subject: String,
    pub role: Role,
    pub created_at: i64,
}

#[derive(Clone, Debug, Serialize)]
pub struct TokenInfo {
    pub id: i64,
    pub kind: Kind,
    pub name: String,
    pub subject: Option<String>,
    pub repository: Option<String>,
    pub role: Role,
    pub expires_at: Option<i64>,
    pub revoked_at: Option<i64>,
    pub last_used_at: Option<i64>,
    pub created_at: i64,
}

pub struct NewToken<'a> {
    pub kind: Kind,
    pub name: &'a str,
    pub subject: Option<&'a str>,
    pub repository: Option<&'a str>,
    pub role: Role,
    pub ttl_secs: Option<i64>,
}

/// An authenticated non-master credential.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Principal {
    pub token_id: i64,
    pub kind: Kind,
    pub subject: Option<String>,
    pub repository_id: Option<i64>,
    pub role_cap: Role,
}

fn hash_token(plaintext: &str) -> String {
    Sha256::digest(plaintext.as_bytes())
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect()
}

fn generate_token() -> Result<String> {
    let mut bytes = [0_u8; TOKEN_RANDOM_BYTES];
    getrandom::fill(&mut bytes).context("read operating-system randomness")?;
    let body: String = bytes.iter().map(|byte| format!("{byte:02x}")).collect();
    Ok(format!("{TOKEN_PREFIX}{body}"))
}

fn validate_text(label: &str, value: &str, max: usize) -> Result<()> {
    ensure!(
        !value.trim().is_empty() && value.len() <= max,
        "{label} must be 1 to {max} bytes"
    );
    Ok(())
}

fn repository_id(connection: &rusqlite::Connection, name: &str) -> Result<i64> {
    connection
        .query_row(
            "SELECT id FROM repositories WHERE name = ?1",
            [name],
            |row| row.get(0),
        )
        .optional()?
        .with_context(|| format!("unknown repository {name}"))
}

pub fn set_collaborator(
    database: &Database,
    repository: &str,
    subject: &str,
    role: Role,
) -> Result<()> {
    validate_text("subject", subject, MAX_SUBJECT)?;
    let connection = database.connection()?;
    let id = repository_id(&connection, repository)?;
    connection.execute(
        "INSERT INTO repository_collaborators(repository_id, subject, role) VALUES (?1, ?2, ?3)
         ON CONFLICT(repository_id, subject) DO UPDATE SET role = excluded.role",
        params![id, subject, role.as_str()],
    )?;
    Ok(())
}

pub fn remove_collaborator(database: &Database, repository: &str, subject: &str) -> Result<bool> {
    let connection = database.connection()?;
    let id = repository_id(&connection, repository)?;
    Ok(connection.execute(
        "DELETE FROM repository_collaborators WHERE repository_id = ?1 AND subject = ?2",
        params![id, subject],
    )? > 0)
}

pub fn collaborators(database: &Database, repository: &str) -> Result<Vec<Collaborator>> {
    let connection = database.connection()?;
    let id = repository_id(&connection, repository)?;
    let mut statement = connection.prepare(
        "SELECT subject, role, created_at FROM repository_collaborators
         WHERE repository_id = ?1 ORDER BY subject",
    )?;
    let rows = statement.query_map([id], |row| {
        Ok((
            row.get::<_, String>(0)?,
            row.get::<_, String>(1)?,
            row.get(2)?,
        ))
    })?;
    rows.map(|row| {
        let (subject, role, created_at) = row?;
        Ok(Collaborator {
            subject,
            role: Role::parse(&role)?,
            created_at,
        })
    })
    .collect()
}

pub fn collaborator_role(
    database: &Database,
    repository: &str,
    subject: &str,
) -> Result<Option<Role>> {
    let connection = database.connection()?;
    connection
        .query_row(
            "SELECT c.role FROM repository_collaborators c
             JOIN repositories r ON r.id = c.repository_id
             WHERE r.name = ?1 AND c.subject = ?2",
            params![repository, subject],
            |row| row.get::<_, String>(0),
        )
        .optional()?
        .map(|role| Role::parse(&role))
        .transpose()
}

/// Creates a token and returns its metadata with the one-time plaintext.
pub fn create_token(
    database: &Database,
    new: &NewToken<'_>,
    now: i64,
) -> Result<(TokenInfo, String)> {
    validate_text("name", new.name, MAX_NAME)?;
    if let Some(subject) = new.subject {
        validate_text("subject", subject, MAX_SUBJECT)?;
    }
    if let Some(ttl) = new.ttl_secs {
        ensure!(
            (1..=MAX_TTL_SECS).contains(&ttl),
            "ttl must be between 1 and {MAX_TTL_SECS} seconds"
        );
    }
    match new.kind {
        Kind::Pat => ensure!(
            new.subject.is_some(),
            "a personal access token needs a subject"
        ),
        Kind::DeployKey => {
            ensure!(new.repository.is_some(), "a deploy key needs a repository");
            ensure!(new.role != Role::Admin, "a deploy key cannot be admin");
        }
    }
    let plaintext = generate_token()?;
    let expires_at = new.ttl_secs.map(|ttl| now + ttl);
    let connection = database.connection()?;
    let repository_id = new
        .repository
        .map(|name| repository_id(&connection, name))
        .transpose()?;
    connection.execute(
        "INSERT INTO access_tokens(kind, name, subject, repository_id, role, token_hash, expires_at, created_at)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
        params![
            new.kind.as_str(),
            new.name,
            new.subject,
            repository_id,
            new.role.as_str(),
            hash_token(&plaintext),
            expires_at,
            now
        ],
    )?;
    let id = connection.last_insert_rowid();
    drop(connection);
    let info = token(database, id)?.context("token vanished after insert")?;
    Ok((info, plaintext))
}

const TOKEN_SELECT: &str = "SELECT t.id, t.kind, t.name, t.subject, r.name, t.role,
    t.expires_at, t.revoked_at, t.last_used_at, t.created_at
    FROM access_tokens t LEFT JOIN repositories r ON r.id = t.repository_id";

fn map_token(row: &rusqlite::Row<'_>) -> rusqlite::Result<(TokenInfo, String, String)> {
    let kind: String = row.get(1)?;
    let role: String = row.get(5)?;
    Ok((
        TokenInfo {
            id: row.get(0)?,
            kind: Kind::Pat,
            name: row.get(2)?,
            subject: row.get(3)?,
            repository: row.get(4)?,
            role: Role::Read,
            expires_at: row.get(6)?,
            revoked_at: row.get(7)?,
            last_used_at: row.get(8)?,
            created_at: row.get(9)?,
        },
        kind,
        role,
    ))
}

fn finish_token((mut info, kind, role): (TokenInfo, String, String)) -> Result<TokenInfo> {
    info.kind = Kind::parse(&kind)?;
    info.role = Role::parse(&role)?;
    Ok(info)
}

fn token(database: &Database, id: i64) -> Result<Option<TokenInfo>> {
    let connection = database.connection()?;
    connection
        .query_row(&format!("{TOKEN_SELECT} WHERE t.id = ?1"), [id], map_token)
        .optional()?
        .map(finish_token)
        .transpose()
}

pub fn list_tokens(
    database: &Database,
    subject: Option<&str>,
    repository: Option<&str>,
    kind: Kind,
) -> Result<Vec<TokenInfo>> {
    let connection = database.connection()?;
    let mut statement = connection.prepare(&format!(
        "{TOKEN_SELECT} WHERE t.kind = ?1 AND (?2 IS NULL OR t.subject = ?2)
         AND (?3 IS NULL OR r.name = ?3) ORDER BY t.id"
    ))?;
    let rows = statement.query_map(params![kind.as_str(), subject, repository], map_token)?;
    rows.map(|row| finish_token(row?)).collect()
}

/// Revokes a token; scoped to `repository` when given so a repo admin cannot
/// touch another repo's keys.
pub fn revoke_token(
    database: &Database,
    id: i64,
    repository: Option<&str>,
    now: i64,
) -> Result<bool> {
    let connection = database.connection()?;
    Ok(connection.execute(
        "UPDATE access_tokens SET revoked_at = ?2
         WHERE id = ?1 AND revoked_at IS NULL
         AND (?3 IS NULL OR repository_id = (SELECT id FROM repositories WHERE name = ?3))",
        params![id, now, repository],
    )? > 0)
}

/// Resolves a plaintext token to a principal; revoked, expired and unknown
/// tokens yield `None`.
pub fn authenticate(database: &Database, plaintext: &str, now: i64) -> Result<Option<Principal>> {
    if !plaintext.starts_with(TOKEN_PREFIX) {
        return Ok(None);
    }
    let connection = database.connection()?;
    let row = connection
        .query_row(
            "SELECT id, kind, subject, repository_id, role FROM access_tokens
             WHERE token_hash = ?1 AND revoked_at IS NULL
             AND (expires_at IS NULL OR expires_at > ?2)",
            params![hash_token(plaintext), now],
            |row| {
                Ok((
                    row.get::<_, i64>(0)?,
                    row.get::<_, String>(1)?,
                    row.get::<_, Option<String>>(2)?,
                    row.get::<_, Option<i64>>(3)?,
                    row.get::<_, String>(4)?,
                ))
            },
        )
        .optional()?;
    let Some((token_id, kind, subject, repository_id, role)) = row else {
        return Ok(None);
    };
    connection.execute(
        "UPDATE access_tokens SET last_used_at = ?2 WHERE id = ?1",
        params![token_id, now],
    )?;
    Ok(Some(Principal {
        token_id,
        kind: Kind::parse(&kind)?,
        subject,
        repository_id,
        role_cap: Role::parse(&role)?,
    }))
}

impl Principal {
    /// Effective role on `repository`: a deploy key holds its own role on its
    /// repository only; a PAT holds the lesser of its cap and the subject's
    /// collaborator role, limited to its repository when scoped.
    pub fn role_on(&self, database: &Database, repository: &str) -> Result<Option<Role>> {
        if let Some(scoped) = self.repository_id {
            let connection = database.connection()?;
            let id = repository_id(&connection, repository)?;
            drop(connection);
            if id != scoped {
                return Ok(None);
            }
        }
        match self.kind {
            Kind::DeployKey => Ok(Some(self.role_cap)),
            Kind::Pat => {
                let subject = self.subject.as_deref().unwrap_or_default();
                Ok(collaborator_role(database, repository, subject)?
                    .map(|granted| granted.min(self.role_cap)))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const NOW: i64 = 1_000_000;
    const KEY: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";

    fn database() -> Database {
        let db = Database::memory(KEY).expect("db");
        db.add_repository("demo", "https://example.com/demo.git", "main")
            .expect("repo");
        db.add_repository("other", "https://example.com/other.git", "main")
            .expect("repo");
        db
    }

    fn pat<'a>(role: Role, repository: Option<&'a str>, ttl: Option<i64>) -> NewToken<'a> {
        NewToken {
            kind: Kind::Pat,
            name: "laptop",
            subject: Some("ada@example.com"),
            repository,
            role,
            ttl_secs: ttl,
        }
    }

    #[test]
    fn role_ordering_gates_actions() {
        assert!(Role::Admin.allows(Role::Write));
        assert!(Role::Write.allows(Role::Read));
        assert!(!Role::Read.allows(Role::Write));
        assert!(Role::parse("owner").is_err());
    }

    #[test]
    fn collaborator_roles_upsert_list_and_remove() {
        let db = database();
        set_collaborator(&db, "demo", "ada@example.com", Role::Read).expect("set");
        set_collaborator(&db, "demo", "ada@example.com", Role::Write).expect("upsert");
        let list = collaborators(&db, "demo").expect("list");
        assert_eq!(list.len(), 1);
        assert_eq!(list[0].role, Role::Write);
        assert!(remove_collaborator(&db, "demo", "ada@example.com").expect("rm"));
        assert!(!remove_collaborator(&db, "demo", "ada@example.com").expect("rm again"));
        assert!(set_collaborator(&db, "missing", "a", Role::Read).is_err());
    }

    #[test]
    fn token_plaintext_is_returned_once_and_only_hash_is_stored() {
        let db = database();
        let (info, plaintext) =
            create_token(&db, &pat(Role::Write, None, Some(60)), NOW).expect("create");
        assert!(plaintext.starts_with("abt_"));
        assert_eq!(info.expires_at, Some(NOW + 60));
        let stored: String = db
            .connection()
            .expect("conn")
            .query_row("SELECT token_hash FROM access_tokens", [], |row| row.get(0))
            .expect("hash");
        assert!(!stored.contains(&plaintext));
        assert_eq!(stored, hash_token(&plaintext));
    }

    #[test]
    fn authenticate_rejects_unknown_expired_and_revoked_tokens() {
        let db = database();
        let (info, plaintext) =
            create_token(&db, &pat(Role::Read, None, Some(60)), NOW).expect("create");
        assert!(
            authenticate(&db, &plaintext, NOW + 59)
                .expect("ok")
                .is_some()
        );
        assert!(
            authenticate(&db, &plaintext, NOW + 60)
                .expect("ok")
                .is_none()
        );
        assert!(authenticate(&db, "abt_nope", NOW).expect("ok").is_none());
        assert!(authenticate(&db, "not-a-token", NOW).expect("ok").is_none());
        assert!(revoke_token(&db, info.id, None, NOW).expect("revoke"));
        assert!(authenticate(&db, &plaintext, NOW).expect("ok").is_none());
    }

    #[test]
    fn ttl_is_bounded() {
        let db = database();
        assert!(create_token(&db, &pat(Role::Read, None, Some(0)), NOW).is_err());
        assert!(create_token(&db, &pat(Role::Read, None, Some(MAX_TTL_SECS + 1)), NOW).is_err());
    }

    #[test]
    fn pat_role_is_capped_by_collaborator_role_and_scope() {
        let db = database();
        set_collaborator(&db, "demo", "ada@example.com", Role::Write).expect("set");
        let (_, wide) = create_token(&db, &pat(Role::Admin, None, None), NOW).expect("wide");
        let (_, narrow) = create_token(&db, &pat(Role::Read, None, None), NOW).expect("narrow");
        let (_, scoped) =
            create_token(&db, &pat(Role::Write, Some("demo"), None), NOW).expect("scoped");
        let role = |token: &str, repo: &str| {
            authenticate(&db, token, NOW)
                .expect("auth")
                .expect("principal")
                .role_on(&db, repo)
                .expect("role")
        };
        assert_eq!(role(&wide, "demo"), Some(Role::Write));
        assert_eq!(role(&narrow, "demo"), Some(Role::Read));
        assert_eq!(role(&wide, "other"), None);
        assert_eq!(role(&scoped, "demo"), Some(Role::Write));
        assert_eq!(role(&scoped, "other"), None);
    }

    #[test]
    fn deploy_key_is_bound_to_one_repository_and_never_admin() {
        let db = database();
        let new = |role, repository| NewToken {
            kind: Kind::DeployKey,
            name: "ci",
            subject: None,
            repository,
            role,
            ttl_secs: None,
        };
        assert!(create_token(&db, &new(Role::Admin, Some("demo")), NOW).is_err());
        assert!(create_token(&db, &new(Role::Read, None), NOW).is_err());
        let (_, key) = create_token(&db, &new(Role::Read, Some("demo")), NOW).expect("key");
        let principal = authenticate(&db, &key, NOW).expect("auth").expect("some");
        assert_eq!(principal.role_on(&db, "demo").expect("r"), Some(Role::Read));
        assert_eq!(principal.role_on(&db, "other").expect("r"), None);
        let listed = list_tokens(&db, None, Some("demo"), Kind::DeployKey).expect("list");
        assert_eq!(listed.len(), 1);
    }

    #[test]
    fn revoke_is_scoped_to_repository() {
        let db = database();
        let (info, _) = create_token(
            &db,
            &NewToken {
                kind: Kind::DeployKey,
                name: "ci",
                subject: None,
                repository: Some("demo"),
                role: Role::Read,
                ttl_secs: None,
            },
            NOW,
        )
        .expect("key");
        assert!(!revoke_token(&db, info.id, Some("other"), NOW).expect("wrong repo"));
        assert!(revoke_token(&db, info.id, Some("demo"), NOW).expect("right repo"));
    }
}