AkurAI Build
Menu

akurai-tasks

public

Latest change 18f4b1ebc4782165ffc2d599bc25ac07a245be76 - Harden task mutations and repair web flows by Ólafur Búi Ólafsson

use super::*;
use sha2::{Digest, Sha256};
use std::fs::File;
use std::io::Read;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Permission {
    Read,
    Plan,
    Comment,
    Execute,
    Review,
    Release,
    Admin,
}

impl Store {
    pub fn bootstrap(
        &self,
        principal: &str,
        display: &str,
        actor: &str,
        idempotency: &str,
    ) -> Result<Value> {
        validate_component("principal", principal)?;
        required("display name", display)?;
        let request = object(vec![
            ("command", string("bootstrap")),
            ("principal", string(principal)),
            ("display", string(display.trim())),
        ]);
        self.mutate(actor, idempotency, &request, |db, _| {
            if get_value(db, "meta/bootstrap-complete")?.is_some()
                || !db
                    .range(b"role/", &upper_bound(b"role/"))
                    .map_err(storage)?
                    .is_empty()
            {
                return Err(Error::Conflict(
                    "authorization bootstrap already completed".into(),
                ));
            }
            let now = now();
            let principal_record = object(vec![
                ("id", string(principal)),
                ("kind", string("human")),
                ("display", string(display.trim())),
                ("active", Value::Bool(true)),
                ("revision", Value::Int(1)),
                ("createdAt", Value::Int(now)),
            ]);
            put(db, &format!("principal/{principal}"), &principal_record)?;
            let binding = role_binding(principal, "*", "administrator", actor, now);
            put(db, &format!("role/*/{principal}/administrator"), &binding)?;
            put(
                db,
                "meta/bootstrap-complete",
                &object(vec![
                    ("principal", string(principal)),
                    ("createdAt", Value::Int(now)),
                ]),
            )?;
            Ok(principal_record)
        })
    }

    pub fn upsert_principal(
        &self,
        id: &str,
        kind: &str,
        display: &str,
        active: bool,
        actor: &str,
        idempotency: &str,
    ) -> Result<Value> {
        validate_component("principal", id)?;
        if !matches!(kind, "human" | "agent" | "service" | "import") {
            return Err(Error::Invalid(
                "principal kind must be human, agent, service, or import".into(),
            ));
        }
        required("display name", display)?;
        let request = object(vec![
            ("command", string("upsert_principal")),
            ("id", string(id)),
            ("kind", string(kind)),
            ("active", Value::Bool(active)),
        ]);
        self.mutate(actor, idempotency, &request, |db, _| {
            let revision = get_value(db, &format!("principal/{id}"))?
                .and_then(|v| v.get("revision").and_then(Value::as_i64))
                .unwrap_or(0)
                + 1;
            let record = object(vec![
                ("id", string(id)),
                ("kind", string(kind)),
                ("display", string(display.trim())),
                ("active", Value::Bool(active)),
                ("revision", Value::Int(revision)),
                ("updatedAt", Value::Int(now())),
            ]);
            put(db, &format!("principal/{id}"), &record)?;
            Ok(record)
        })
    }

    pub fn grant_role(
        &self,
        principal: &str,
        scope: &str,
        role: &str,
        actor: &str,
        idempotency: &str,
    ) -> Result<Value> {
        validate_component("principal", principal)?;
        validate_scope(scope)?;
        validate_role(role)?;
        let request = object(vec![
            ("command", string("grant_role")),
            ("principal", string(principal)),
            ("scope", string(scope)),
            ("role", string(role)),
        ]);
        self.mutate(actor, idempotency, &request, |db, _| {
            let target = get_value(db, &format!("principal/{principal}"))?
                .ok_or_else(|| Error::NotFound(format!("principal {principal}")))?;
            if target.get("active").and_then(Value::as_bool) != Some(true) {
                return Err(Error::Invalid("cannot bind an inactive principal".into()));
            }
            let binding = role_binding(principal, scope, role, actor, now());
            put(db, &format!("role/{scope}/{principal}/{role}"), &binding)?;
            Ok(binding)
        })
    }

    pub fn issue_credential(
        &self,
        principal: &str,
        label: &str,
        expires_at: Option<i64>,
        actor: &str,
        idempotency: &str,
    ) -> Result<Value> {
        validate_component("principal", principal)?;
        required("credential label", label)?;
        let token = random_token(32)?;
        let digest = digest(&token);
        let id = random_token(12)?;
        let request = object(vec![
            ("command", string("issue_credential")),
            ("principal", string(principal)),
            ("label", string(label.trim())),
            (
                "expiresAt",
                expires_at.map(Value::Int).unwrap_or(Value::Null),
            ),
        ]);
        let mut result = self.mutate(actor, idempotency, &request, |db, _| {
            let target = get_value(db, &format!("principal/{principal}"))?
                .ok_or_else(|| Error::NotFound(format!("principal {principal}")))?;
            if target.get("active").and_then(Value::as_bool) != Some(true) {
                return Err(Error::Invalid("principal is inactive".into()));
            }
            let record = object(vec![
                ("id", string(&id)),
                ("principal", string(principal)),
                ("label", string(label.trim())),
                ("hash", string(&digest)),
                (
                    "expiresAt",
                    expires_at.map(Value::Int).unwrap_or(Value::Null),
                ),
                ("revoked", Value::Bool(false)),
                ("revision", Value::Int(1)),
                ("createdAt", Value::Int(now())),
            ]);
            put(db, &format!("credential/{id}"), &record)?;
            Ok(record)
        })?;
        update(&mut result, "token", string(&token));
        update(&mut result, "hash", Value::Null);
        Ok(result)
    }

    pub fn revoke_credential(&self, id: &str, actor: &str, idempotency: &str) -> Result<Value> {
        validate_component("credential", id)?;
        let request = object(vec![
            ("command", string("revoke_credential")),
            ("id", string(id)),
        ]);
        self.mutate(actor, idempotency, &request, |db, _| {
            let mut record = get_value(db, &format!("credential/{id}"))?
                .ok_or_else(|| Error::NotFound(format!("credential {id}")))?;
            update(&mut record, "revoked", Value::Bool(true));
            bump(&mut record);
            put(db, &format!("credential/{id}"), &record)?;
            update(&mut record, "hash", Value::Null);
            Ok(record)
        })
    }

    pub fn authenticate_bearer(&self, token: &str) -> Result<Option<String>> {
        if token.len() < 32 {
            return Ok(None);
        }
        let wanted = digest(token);
        let mut db = self
            .db
            .lock()
            .map_err(|_| Error::Storage("lock poisoned".into()))?;
        for (_, bytes) in db
            .range(b"credential/", &upper_bound(b"credential/"))
            .map_err(storage)?
        {
            let record = decode(&bytes)?;
            let expired = record
                .get("expiresAt")
                .and_then(Value::as_i64)
                .is_some_and(|expiry| expiry <= now());
            if record.get("revoked").and_then(Value::as_bool) == Some(true) || expired {
                continue;
            }
            if record
                .get("hash")
                .and_then(Value::as_str)
                .is_some_and(|stored| constant_eq(stored.as_bytes(), wanted.as_bytes()))
            {
                let principal = field(&record, "principal").unwrap_or("");
                let active = get_value(&mut db, &format!("principal/{principal}"))?
                    .and_then(|p| p.get("active").and_then(Value::as_bool))
                    .unwrap_or(false);
                return Ok(active.then(|| principal.to_string()));
            }
        }
        Ok(None)
    }

    pub fn principal_display(&self, principal: &str) -> Result<Option<String>> {
        let mut db = self
            .db
            .lock()
            .map_err(|_| Error::Storage("lock poisoned".into()))?;
        Ok(get_value(&mut db, &format!("principal/{principal}"))?
            .and_then(|value| field(&value, "display").map(str::to_string)))
    }

    pub fn create_session(
        &self,
        principal: &str,
        ttl_seconds: i64,
        actor: &str,
        idempotency: &str,
    ) -> Result<Value> {
        validate_component("principal", principal)?;
        let token = random_token(32)?;
        let csrf = random_token(24)?;
        let hash = digest(&token);
        let request = object(vec![
            ("command", string("create_session")),
            ("principal", string(principal)),
            ("ttl", Value::Int(ttl_seconds)),
        ]);
        let mut result = self.mutate(actor, idempotency, &request, |db, _| {
            let p = get_value(db, &format!("principal/{principal}"))?
                .ok_or_else(|| Error::NotFound(format!("principal {principal}")))?;
            if p.get("active").and_then(Value::as_bool) != Some(true) {
                return Err(Error::Invalid("principal is inactive".into()));
            }
            let record = object(vec![
                ("hash", string(&hash)),
                ("principal", string(principal)),
                ("csrf", string(&digest(&csrf))),
                ("expiresAt", Value::Int(now() + ttl_seconds.max(60))),
                ("revoked", Value::Bool(false)),
                ("createdAt", Value::Int(now())),
            ]);
            put(db, &format!("session/{hash}"), &record)?;
            Ok(record)
        })?;
        update(&mut result, "token", string(&token));
        update(&mut result, "csrf", string(&csrf));
        update(&mut result, "hash", Value::Null);
        Ok(result)
    }

    pub fn authenticate_session(
        &self,
        token: &str,
        csrf: Option<&str>,
        mutation: bool,
    ) -> Result<Option<String>> {
        let hash = digest(token);
        let mut db = self
            .db
            .lock()
            .map_err(|_| Error::Storage("lock poisoned".into()))?;
        let Some(record) = get_value(&mut db, &format!("session/{hash}"))? else {
            return Ok(None);
        };
        if record.get("revoked").and_then(Value::as_bool) == Some(true)
            || record.get("expiresAt").and_then(Value::as_i64).unwrap_or(0) <= now()
        {
            return Ok(None);
        }
        if mutation {
            let Some(csrf) = csrf else { return Ok(None) };
            let supplied = digest(csrf);
            let valid = record
                .get("csrf")
                .and_then(Value::as_str)
                .is_some_and(|stored| constant_eq(stored.as_bytes(), supplied.as_bytes()));
            if !valid {
                return Ok(None);
            }
        }
        let principal = field(&record, "principal").unwrap_or("");
        let active = get_value(&mut db, &format!("principal/{principal}"))?
            .and_then(|p| p.get("active").and_then(Value::as_bool))
            .unwrap_or(false);
        Ok(active.then(|| principal.to_string()))
    }

    pub fn revoke_session(&self, token: &str, actor: &str, idempotency: &str) -> Result<Value> {
        let hash = digest(token);
        let request = object(vec![
            ("command", string("revoke_session")),
            ("session", string(&hash)),
        ]);
        self.mutate(actor, idempotency, &request, |db, _| {
            let mut record = get_value(db, &format!("session/{hash}"))?
                .ok_or_else(|| Error::NotFound("session".into()))?;
            update(&mut record, "revoked", Value::Bool(true));
            put(db, &format!("session/{hash}"), &record)?;
            update(&mut record, "hash", Value::Null);
            Ok(record)
        })
    }

    pub fn bootstrap_complete(&self) -> Result<bool> {
        let mut db = self
            .db
            .lock()
            .map_err(|_| Error::Storage("lock poisoned".into()))?;
        Ok(get_value(&mut db, "meta/bootstrap-complete")?.is_some())
    }

    pub fn authorize(
        &self,
        principal: &str,
        project: Option<&str>,
        permission: Permission,
    ) -> Result<bool> {
        let mut db = self
            .db
            .lock()
            .map_err(|_| Error::Storage("lock poisoned".into()))?;
        let active = get_value(&mut db, &format!("principal/{principal}"))?
            .and_then(|p| p.get("active").and_then(Value::as_bool))
            .unwrap_or(false);
        if !active {
            return Ok(false);
        }
        for scope in project.into_iter().chain(std::iter::once("*")) {
            let prefix = format!("role/{scope}/{principal}/");
            for (_, bytes) in db
                .range(prefix.as_bytes(), &upper_bound(prefix.as_bytes()))
                .map_err(storage)?
            {
                let binding = decode(&bytes)?;
                if field(&binding, "role").is_some_and(|role| role_allows(role, permission)) {
                    return Ok(true);
                }
            }
        }
        Ok(false)
    }
}

fn role_binding(principal: &str, scope: &str, role: &str, actor: &str, created: i64) -> Value {
    object(vec![
        ("principal", string(principal)),
        ("scope", string(scope)),
        ("role", string(role)),
        ("grantedBy", string(actor)),
        ("revision", Value::Int(1)),
        ("createdAt", Value::Int(created)),
    ])
}

fn validate_role(role: &str) -> Result<()> {
    if matches!(
        role,
        "viewer"
            | "contributor"
            | "planner"
            | "implementer"
            | "reviewer"
            | "verifier"
            | "releaser"
            | "administrator"
    ) {
        Ok(())
    } else {
        Err(Error::Invalid(format!("unknown role {role}")))
    }
}

fn validate_scope(scope: &str) -> Result<()> {
    if scope == "*" {
        Ok(())
    } else {
        validate_key(scope)
    }
}

fn role_allows(role: &str, permission: Permission) -> bool {
    if role == "administrator" {
        return true;
    }
    match permission {
        Permission::Read => true,
        Permission::Plan => role == "planner",
        Permission::Comment => role == "contributor",
        Permission::Execute => role == "implementer",
        Permission::Review => matches!(role, "reviewer" | "verifier"),
        Permission::Release => role == "releaser",
        Permission::Admin => false,
    }
}

pub(crate) fn digest(value: &str) -> String {
    let bytes = Sha256::digest(value.as_bytes());
    let mut out = String::with_capacity(bytes.len() * 2);
    for byte in bytes {
        out.push_str(&format!("{byte:02x}"));
    }
    out
}

pub(crate) fn random_token(bytes: usize) -> Result<String> {
    let mut input = vec![0u8; bytes];
    File::open("/dev/urandom")
        .and_then(|mut file| file.read_exact(&mut input))
        .map_err(storage)?;
    let mut out = String::with_capacity(bytes * 2);
    for byte in input {
        out.push_str(&format!("{byte:02x}"));
    }
    Ok(out)
}

pub(crate) fn constant_eq(left: &[u8], right: &[u8]) -> bool {
    if left.len() != right.len() {
        return false;
    }
    left.iter()
        .zip(right)
        .fold(0u8, |diff, (a, b)| diff | (a ^ b))
        == 0
}

fn validate_component(label: &str, value: &str) -> Result<()> {
    required(label, value)?;
    if value.len() > 128
        || !value
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'@'))
    {
        return Err(Error::Invalid(format!(
            "{label} contains unsupported characters"
        )));
    }
    Ok(())
}

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

    #[test]
    fn bootstrap_scoped_roles_and_credentials_are_enforced() {
        let path = std::env::temp_dir().join(format!(
            "akurai-tasks-auth-{}-{}.db",
            std::process::id(),
            now()
        ));
        let _ = std::fs::remove_file(&path);
        let store = Store::open(&path).unwrap();
        store
            .bootstrap("admin@example.com", "Admin", "bootstrap", "bootstrap")
            .unwrap();
        assert!(store
            .authorize("admin@example.com", Some("OPS"), Permission::Admin)
            .unwrap());
        assert!(store
            .bootstrap("other@example.com", "Other", "bootstrap", "again")
            .is_err());
        store
            .upsert_principal(
                "agent-1",
                "agent",
                "Agent",
                true,
                "admin@example.com",
                "agent",
            )
            .unwrap();
        store
            .grant_role("agent-1", "OPS", "implementer", "admin@example.com", "role")
            .unwrap();
        assert!(store
            .authorize("agent-1", Some("OPS"), Permission::Execute)
            .unwrap());
        assert!(!store
            .authorize("agent-1", Some("OTHER"), Permission::Execute)
            .unwrap());
        let issued = store
            .issue_credential("agent-1", "mcp", None, "admin@example.com", "credential")
            .unwrap();
        let token = field(&issued, "token").unwrap().to_string();
        assert!(issued.get("hash").is_none() || issued.get("hash") == Some(&Value::Null));
        assert_eq!(
            store.authenticate_bearer(&token).unwrap().as_deref(),
            Some("agent-1")
        );
        store
            .revoke_credential(field(&issued, "id").unwrap(), "admin@example.com", "revoke")
            .unwrap();
        assert_eq!(store.authenticate_bearer(&token).unwrap(), None);
        drop(store);
        let _ = std::fs::remove_file(path);
    }
}