AkurAI-Drive / backend/src/session.rs

Latest change feb03c4c2e9a6b0e869a63e52329c938e68dc8e8 - backend: device-token auth, SigV4 presigning, and file API for the sync client by olibuijr

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

use crate::db::with_db;

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

/// A live, non-expired session loaded from the `sessions` table.
pub struct Session {
    pub sub: String,
    pub email: 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 new session row and return its random id (to set as `drive_session`).
pub fn create_session(
    sub: &str,
    email: &str,
    groups: &[String],
    tenant_id: Option<&str>,
) -> String {
    let id = crate::util::token();
    let now = chrono::Utc::now().timestamp();
    let expires = now + SESSION_TTL_SECS;
    let groups_json = serde_json::to_string(groups).unwrap_or_else(|_| "[]".to_string());

    with_db(|conn| {
        conn.execute(
            "INSERT INTO sessions (id, sub, email, groups, tenant_id, created, expires)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
            rusqlite::params![id, sub, email, groups_json, tenant_id, now, expires],
        )
        .expect("failed to insert session");
    });

    id
}

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

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

/// Remember where to send the browser after a successful login, keyed on the OAuth
/// `state`. Used by `/auth/device` so it can resume once a session exists.
pub fn store_oauth_return(state: &str, return_to: &str) {
    let now = chrono::Utc::now().timestamp();
    with_db(|conn| {
        conn.execute(
            "INSERT OR REPLACE INTO oauth_states (state, return_to, created) VALUES (?1, ?2, ?3)",
            rusqlite::params![state, return_to, now],
        )
        .ok();
    });
}

/// Consume the stored `return_to` for an OAuth `state` (one-shot — the row is
/// deleted whether or not it existed).
pub fn take_oauth_return(state: &str) -> Option<String> {
    with_db(|conn| {
        let return_to: Option<String> = conn
            .query_row(
                "SELECT return_to FROM oauth_states WHERE state = ?1",
                rusqlite::params![state],
                |row| row.get(0),
            )
            .ok();
        conn.execute(
            "DELETE FROM oauth_states WHERE state = ?1",
            rusqlite::params![state],
        )
        .ok();
        return_to
    })
}