AkurAI-Drive / backend/src/session.rs

Latest change 87af7df38e96b7f796055497f19e265e0dff2cdb - Drive logout via IdP /logout; pin akurai-log 3c855de by AkurAI Agent

use axum::http::{HeaderMap, header};

use crate::db::with_db;

const SESSION_TTL_SECS: i64 = 30 * 24 * 60 * 60; // ~30 days
const LOGIN_STATE_TTL_SECS: i64 = 600;

/// Session cookie shared with the other AkurAI apps' SSO login.
pub const SESSION_COOKIE: &str = "akurai_session";
/// Short-lived cookie binding an in-flight login's `state` to this browser.
pub const STATE_COOKIE: &str = "akurai_oauth_state";

/// `Set-Cookie` value for `name`; `None` clears it.
pub fn cookie(name: &str, value: Option<&str>, max_age: i64) -> String {
    let (v, age) = value.map_or(("", 0), |v| (v, max_age));
    format!("{name}={v}; Path=/; HttpOnly; Secure; SameSite=Lax; Max-Age={age}")
}

pub fn session_cookie(id: &str) -> String {
    cookie(SESSION_COOKIE, Some(id), SESSION_TTL_SECS)
}

pub fn state_cookie(state: &str) -> String {
    cookie(STATE_COOKIE, Some(state), LOGIN_STATE_TTL_SECS)
}

/// A live, non-expired session loaded from the `sessions` table.
pub struct Session {
    pub sub: String,
    pub email: String,
    pub role: String,
    pub organization_id: Option<String>,
    pub groups: Vec<String>,
    pub tenant_id: Option<String>,
}

/// Read a named cookie value out of the request `Cookie` header.
pub fn get_cookie(headers: &HeaderMap, name: &str) -> Option<String> {
    let cookies = headers.get(header::COOKIE)?.to_str().ok()?;
    let prefix = format!("{name}=");
    for part in cookies.split(';') {
        let part = part.trim();
        if let Some(value) = part.strip_prefix(&prefix) {
            return Some(value.to_string());
        }
    }
    None
}

/// Persist a session for a signed-in identity and return its random id.
pub fn create_session(id: &crate::oidc::Identity) -> String {
    let sid = crate::util::token();
    let now = chrono::Utc::now().timestamp();
    let groups_json = serde_json::to_string(&id.groups).unwrap_or_else(|_| "[]".to_string());

    with_db(|conn| {
        conn.execute(
            "DELETE FROM sessions WHERE expires <= ?1",
            rusqlite::params![now],
        )
        .ok();
        conn.execute(
            "INSERT INTO sessions
             (id, sub, email, role, organization_id, tenant_id, groups, id_token, created, expires)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
            rusqlite::params![
                sid,
                id.sub,
                id.email,
                id.role,
                id.organization_id,
                id.tenant_id,
                groups_json,
                id.id_token,
                now,
                now + SESSION_TTL_SECS
            ],
        )
        .expect("failed to insert session");
    });

    sid
}

/// Look up the current session from the session cookie. Expired rows are
/// treated as no session.
pub fn lookup_session(headers: &HeaderMap) -> Option<Session> {
    let id = get_cookie(headers, SESSION_COOKIE)?;
    let now = chrono::Utc::now().timestamp();

    with_db(|conn| {
        conn.query_row(
            "SELECT sub, email, role, organization_id, tenant_id, groups
                 FROM sessions WHERE id = ?1 AND expires > ?2",
            rusqlite::params![id, now],
            |row| {
                let groups_json: String = row.get(5)?;
                Ok(Session {
                    sub: row.get(0)?,
                    email: row.get(1)?,
                    role: row.get(2)?,
                    organization_id: row.get(3)?,
                    tenant_id: row.get(4)?,
                    groups: serde_json::from_str(&groups_json).unwrap_or_default(),
                })
            },
        )
        .ok()
    })
}

/// Delete a session row by id (used on logout).
pub fn delete_session(id: &str) {
    with_db(|conn| {
        conn.execute("DELETE FROM sessions WHERE id = ?1", rusqlite::params![id])
            .ok();
    });
}

/// Server-side half of an in-flight login.
#[derive(Debug, PartialEq)]
pub struct LoginState {
    pub nonce: String,
    pub verifier: String,
    pub return_to: String,
}

pub fn store_login_state(state: &str, login: &LoginState) {
    let now = chrono::Utc::now().timestamp();
    with_db(|conn| {
        conn.execute(
            "INSERT OR REPLACE INTO oauth_states (state, nonce, verifier, return_to, created)
             VALUES (?1, ?2, ?3, ?4, ?5)",
            rusqlite::params![state, login.nonce, login.verifier, login.return_to, now],
        )
        .ok();
    });
}

/// Consume a login state (single use); expired rows are swept on the way.
pub fn take_login_state(state: &str) -> Option<LoginState> {
    let cutoff = chrono::Utc::now().timestamp() - LOGIN_STATE_TTL_SECS;
    with_db(|conn| {
        let found = conn
            .query_row(
                "SELECT nonce, verifier, return_to FROM oauth_states
                 WHERE state = ?1 AND created > ?2",
                rusqlite::params![state, cutoff],
                |row| {
                    Ok(LoginState {
                        nonce: row.get(0)?,
                        verifier: row.get(1)?,
                        return_to: row.get(2)?,
                    })
                },
            )
            .ok();
        conn.execute(
            "DELETE FROM oauth_states WHERE state = ?1 OR created <= ?2",
            rusqlite::params![state, cutoff],
        )
        .ok();
        found
    })
}

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

    #[test]
    fn cookies_set_and_clear() {
        assert_eq!(
            state_cookie("s1"),
            "akurai_oauth_state=s1; Path=/; HttpOnly; Secure; SameSite=Lax; Max-Age=600"
        );
        assert_eq!(
            cookie(SESSION_COOKIE, None, 99),
            "akurai_session=; Path=/; HttpOnly; Secure; SameSite=Lax; Max-Age=0"
        );
    }
}