AkurAI Build
Menu

Public repository / fa37e2544d31a1199954a395f9e5b3ebb237b17f

AkurAI-Framework / crates/cli/src/auth.rs

Latest change 0bb96d8dc2b9ebd83dff1e8f321a26b5367af139 - feat(cli): cookie-based session auth + login flow (Batch 4 F3) by Ólafur Búi Ólafsson

//! Cookie-based session authentication on the embedded B+tree.
//!
//! This is the framework's reference auth feature: a durable session store, a
//! tiny users table, a plain-HTML login form, a session-cookie middleware
//! helper, a protected page, and logout. It mirrors the guestbook form-actions
//! demo (validate → store → 303 Post/Redirect/Get) and the record store
//! (`Arc<Mutex<BTree>>`, durable on commit), adding only what auth needs: a
//! `Cookie:` request-header reader and a `Set-Cookie:` response-header writer.
//!
//! ## Security caveats (read these — this is demo-grade, not production)
//!
//! * **PRNG:** session ids are derived from `std` only — `SystemTime` nanos
//!   mixed with a per-process atomic counter and the pid, run through a
//!   SplitMix64 diffuser. This is a framework-demo CSPRNG-lite. It is *not* a
//!   cryptographically secure RNG; do not rely on it to resist a determined
//!   attacker guessing session ids.
//! * **Password hash:** a salted, stretched SHA-1 (reusing the WebSocket SHA-1).
//!   It is salted and iterated, but SHA-1 is fast and not memory-hard — this is
//!   **not bcrypt/argon2** and must not be presented as such.
//! * **Transport:** the cookie is `HttpOnly`, `Path=/`, `SameSite=Lax`. It is
//!   not marked `Secure` here because TLS terminates at the edge (nginx); the
//!   HTTPS layer there is what actually protects the cookie in transit.

use crate::cmd_serve::{build_engine, error_page, load_context};
use akurai_json::Value;
use akurai_storage::BTree;
use akurai_ws::sha1::sha1;
use std::io;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};

use akurai_http::{form, Method, Reply, Request, Response};

/// The seeded demo user so the login flow can be exercised out of the box.
pub const DEMO_USERNAME: &str = "demo";
/// The seeded demo user's password. Demo-grade; documented, not a secret.
pub const DEMO_PASSWORD: &str = "akurai";

/// The session cookie name.
const COOKIE_NAME: &str = "session";

// ---------------------------------------------------------------------------
// State
// ---------------------------------------------------------------------------

/// Auth handler state: a durable session store and a durable users table. Both
/// are cheap to clone (each is an `Arc`), so the whole struct clones freely into
/// the server's request closure.
#[derive(Clone)]
pub struct State {
    sessions: SessionStore,
    users: UserStore,
}

impl State {
    /// Open (or create) the session and user stores under `data_dir`, seeding
    /// the demo user if the users table has no entry for it yet.
    pub fn open(data_dir: &Path) -> io::Result<State> {
        let sessions = SessionStore::open(data_dir.join("sessions.db"))?;
        let users = UserStore::open(data_dir.join("users.db"))?;
        users.ensure_user(DEMO_USERNAME, DEMO_PASSWORD)?;
        Ok(State { sessions, users })
    }

    /// The user named by a request's session cookie, if it maps to a live
    /// session. This is the session middleware/helper: read `Cookie`, look up
    /// the session, yield the current user (or `None`).
    pub fn current_user(&self, req: &Request) -> Option<String> {
        let cookie = req.header("Cookie")?;
        let id = parse_cookie(cookie, COOKIE_NAME)?;
        self.sessions.lookup(&id).ok().flatten().map(|s| s.username)
    }
}

// ---------------------------------------------------------------------------
// Handlers
// ---------------------------------------------------------------------------

/// Route `GET`/`POST /login`.
pub fn login(state: &State, root: &Path, req: &Request) -> Reply {
    // Already signed in? Skip the form and go to the protected page.
    if state.current_user(req).is_some() {
        return Reply::Response(redirect("/account"));
    }
    let response = match req.method {
        Method::Post => login_submit(state, root, req),
        _ => render_login(root, "", ""),
    };
    Reply::Response(response)
}

/// Validate credentials; on success create a session, set the cookie, and
/// 303-redirect to the protected page. On failure re-render with an error.
fn login_submit(state: &State, root: &Path, req: &Request) -> Response {
    let pairs = form::parse_urlencoded(&req.body_str());
    let username = form::field(&pairs, "username")
        .unwrap_or("")
        .trim()
        .to_string();
    let password = form::field(&pairs, "password").unwrap_or("").to_string();

    let ok = state.users.verify(&username, &password).unwrap_or(false);
    if !ok {
        return render_login(root, "Invalid username or password.", &username);
    }

    let id = match state.sessions.create(&username) {
        Ok(id) => id,
        Err(e) => {
            return Response::new(500)
                .with_html(&error_page(&format!("could not start session: {e}")))
        }
    };
    redirect("/account").with_header("Set-Cookie", &set_cookie(&id))
}

/// Render the protected page, or redirect to `/login` when unauthenticated.
pub fn account(state: &State, root: &Path, req: &Request) -> Reply {
    let Some(username) = state.current_user(req) else {
        return Reply::Response(redirect("/login"));
    };
    let engine = match build_engine(root) {
        Ok(engine) => engine,
        Err(msg) => return Reply::Response(Response::new(500).with_html(&error_page(&msg))),
    };
    let mut context = load_context(root);
    if let Value::Object(pairs) = &mut context {
        pairs.push(("username".into(), Value::Str(username)));
    }
    let response = match engine.render("account", &context) {
        Ok(html) => Response::ok().with_html(&html),
        Err(e) => Response::new(500).with_html(&error_page(&e.message)),
    };
    Reply::Response(response)
}

/// Clear the session and expire the cookie, then redirect to `/login`.
pub fn logout(state: &State, req: &Request) -> Reply {
    if let Some(cookie) = req.header("Cookie") {
        if let Some(id) = parse_cookie(cookie, COOKIE_NAME) {
            let _ = state.sessions.delete(&id);
        }
    }
    Reply::Response(redirect("/login").with_header("Set-Cookie", &expired_cookie()))
}

/// Render the login template with an optional error and the username the user
/// already typed (so a failed attempt does not clear the field).
fn render_login(root: &Path, error: &str, username: &str) -> Response {
    let engine = match build_engine(root) {
        Ok(engine) => engine,
        Err(msg) => return Response::new(500).with_html(&error_page(&msg)),
    };
    let mut context = load_context(root);
    if let Value::Object(pairs) = &mut context {
        pairs.push(("error".into(), Value::Str(error.to_string())));
        pairs.push(("form_username".into(), Value::Str(username.to_string())));
    }
    match engine.render("login", &context) {
        Ok(html) => Response::ok().with_html(&html),
        Err(e) => Response::new(500).with_html(&error_page(&e.message)),
    }
}

/// A 303 See Other redirect (Post/Redirect/Get) to `location`.
fn redirect(location: &str) -> Response {
    Response::new(303).with_header("Location", location)
}

// ---------------------------------------------------------------------------
// Cookie helpers
// ---------------------------------------------------------------------------

/// Read a single cookie value out of a `Cookie:` request header. The header is
/// a `; `-separated list of `name=value` pairs; returns the first match for
/// `name`, or `None`.
pub fn parse_cookie(header: &str, name: &str) -> Option<String> {
    header.split(';').find_map(|pair| {
        let (k, v) = pair.split_once('=')?;
        if k.trim() == name {
            Some(v.trim().to_string())
        } else {
            None
        }
    })
}

/// The `Set-Cookie` value that installs a session: `HttpOnly` (no JS access),
/// `Path=/` (whole site), `SameSite=Lax` (sent on top-level navigations only).
/// Not `Secure` — TLS terminates at the edge (see the module caveats).
pub fn set_cookie(id: &str) -> String {
    format!("{COOKIE_NAME}={id}; HttpOnly; Path=/; SameSite=Lax")
}

/// The `Set-Cookie` value that expires the session cookie (empty value,
/// `Max-Age=0`), telling the browser to drop it immediately.
pub fn expired_cookie() -> String {
    format!("{COOKIE_NAME}=; HttpOnly; Path=/; SameSite=Lax; Max-Age=0")
}

// ---------------------------------------------------------------------------
// Session store
// ---------------------------------------------------------------------------

/// A live session: who it belongs to and when it started (unix seconds).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Session {
    pub username: String,
    pub created: i64,
}

/// A durable, clonable session store keyed by opaque session id.
#[derive(Clone)]
struct SessionStore {
    tree: Arc<Mutex<BTree>>,
}

impl SessionStore {
    fn open(path: impl AsRef<Path>) -> io::Result<SessionStore> {
        Ok(SessionStore {
            tree: Arc::new(Mutex::new(BTree::open(path)?)),
        })
    }

    /// Create a session for `username`, returning its new opaque id. Durable on
    /// return (it commits).
    fn create(&self, username: &str) -> io::Result<String> {
        let id = random_id();
        let created = now_secs();
        let record = Value::Object(vec![
            ("username".into(), Value::Str(username.to_string())),
            ("created".into(), Value::Int(created)),
        ]);
        let mut tree = self.tree.lock().unwrap();
        tree.insert(id.as_bytes(), record.to_json().as_bytes())?;
        tree.commit()?;
        Ok(id)
    }

    /// Look up a session by id, or `None` if there is no such session.
    fn lookup(&self, id: &str) -> io::Result<Option<Session>> {
        let mut tree = self.tree.lock().unwrap();
        let Some(bytes) = tree.get(id.as_bytes())? else {
            return Ok(None);
        };
        Ok(decode_session(&bytes))
    }

    /// Delete a session by id. Durable on return. Returns whether it existed.
    fn delete(&self, id: &str) -> io::Result<bool> {
        let mut tree = self.tree.lock().unwrap();
        let removed = tree.delete(id.as_bytes())?;
        tree.commit()?;
        Ok(removed)
    }
}

/// Decode a stored session JSON blob into a [`Session`], or `None` if malformed.
fn decode_session(bytes: &[u8]) -> Option<Session> {
    let text = std::str::from_utf8(bytes).ok()?;
    let value = akurai_json::parse(text).ok()?;
    let username = value.get("username")?.as_str()?.to_string();
    let created = match value.get("created") {
        Some(Value::Int(n)) => *n,
        _ => 0,
    };
    Some(Session { username, created })
}

// ---------------------------------------------------------------------------
// User store
// ---------------------------------------------------------------------------

/// A durable, clonable users table keyed by username, storing a salted password
/// hash (never the plaintext).
#[derive(Clone)]
struct UserStore {
    tree: Arc<Mutex<BTree>>,
}

impl UserStore {
    fn open(path: impl AsRef<Path>) -> io::Result<UserStore> {
        Ok(UserStore {
            tree: Arc::new(Mutex::new(BTree::open(path)?)),
        })
    }

    /// Create `username` with `password` if it does not already exist. Idempotent
    /// — an existing user is left untouched (so the demo password is not reset on
    /// every boot).
    fn ensure_user(&self, username: &str, password: &str) -> io::Result<()> {
        let mut tree = self.tree.lock().unwrap();
        if tree.get(username.as_bytes())?.is_some() {
            return Ok(());
        }
        let salt = random_id();
        let record = Value::Object(vec![
            ("salt".into(), Value::Str(salt.clone())),
            ("hash".into(), Value::Str(hash_password(&salt, password))),
        ]);
        tree.insert(username.as_bytes(), record.to_json().as_bytes())?;
        tree.commit()?;
        Ok(())
    }

    /// Verify a `username`/`password` pair in constant-ish time against the
    /// stored salted hash. A missing user returns `Ok(false)`.
    fn verify(&self, username: &str, password: &str) -> io::Result<bool> {
        let mut tree = self.tree.lock().unwrap();
        let Some(bytes) = tree.get(username.as_bytes())? else {
            return Ok(false);
        };
        let Some((salt, hash)) = decode_user(&bytes) else {
            return Ok(false);
        };
        Ok(constant_time_eq(&hash_password(&salt, password), &hash))
    }
}

/// Decode a stored user JSON blob into `(salt, hash)`, or `None` if malformed.
fn decode_user(bytes: &[u8]) -> Option<(String, String)> {
    let text = std::str::from_utf8(bytes).ok()?;
    let value = akurai_json::parse(text).ok()?;
    let salt = value.get("salt")?.as_str()?.to_string();
    let hash = value.get("hash")?.as_str()?.to_string();
    Some((salt, hash))
}

// ---------------------------------------------------------------------------
// Password hashing (demo-grade — see module caveats)
// ---------------------------------------------------------------------------

/// Hash stretch rounds. SHA-1 is fast, so we iterate to slow brute force a
/// little; this is *not* a memory-hard KDF and must not be sold as bcrypt.
const HASH_ROUNDS: usize = 4096;

/// A salted, stretched SHA-1 of `password`, returned as lowercase hex. Same
/// `(salt, password)` always yields the same hash; a different salt or password
/// yields a different one.
pub fn hash_password(salt: &str, password: &str) -> String {
    let mut digest = sha1(format!("{salt}:{password}").as_bytes());
    // Re-mix the salt into every round so the salt is not just a first-round
    // prefix an attacker can strip.
    let salt_bytes = salt.as_bytes();
    for _ in 0..HASH_ROUNDS {
        let mut buf = Vec::with_capacity(digest.len() + salt_bytes.len());
        buf.extend_from_slice(&digest);
        buf.extend_from_slice(salt_bytes);
        digest = sha1(&buf);
    }
    to_hex(&digest)
}

/// Compare two hex hashes without an early-exit short circuit on the first
/// differing byte (length differences still short-circuit, which is fine here:
/// our hashes are always the same length).
fn constant_time_eq(a: &str, b: &str) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff = 0u8;
    for (x, y) in a.bytes().zip(b.bytes()) {
        diff |= x ^ y;
    }
    diff == 0
}

/// Lowercase hex of a byte slice.
fn to_hex(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        s.push_str(&format!("{b:02x}"));
    }
    s
}

// ---------------------------------------------------------------------------
// CSPRNG-lite id generation (demo-grade — see module caveats)
// ---------------------------------------------------------------------------

/// A monotonic per-process counter mixed into every id so two ids minted within
/// the same `SystemTime` tick still differ.
static ID_COUNTER: AtomicU64 = AtomicU64::new(0);

/// A 128-bit opaque id as 32 hex chars, derived from `std` only:
/// `SystemTime` nanos + pid + an atomic counter, diffused through SplitMix64.
/// Demo-grade randomness — see the module caveats.
fn random_id() -> String {
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos() as u64)
        .unwrap_or(0);
    let pid = std::process::id() as u64;
    let count = ID_COUNTER.fetch_add(1, Ordering::Relaxed);

    let mut seed_a = nanos ^ pid.rotate_left(32);
    let mut seed_b = count
        .wrapping_add(nanos.rotate_left(17))
        .wrapping_mul(0x9E37_79B9_7F4A_7C15);
    let a = split_mix64(&mut seed_a);
    let b = split_mix64(&mut seed_b);
    format!("{a:016x}{b:016x}")
}

/// One SplitMix64 step (Vigna): advance `state` and return a well-diffused word.
fn split_mix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

/// Current unix time in whole seconds (0 on a pre-epoch clock, which never
/// happens in practice).
fn now_secs() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// The scratch base: `CARGO_TARGET_TMPDIR` when present (integration tests),
    /// otherwise a dir under this crate's own `target/` — never `/tmp`.
    fn scratch_base() -> PathBuf {
        match std::env::var_os("CARGO_TARGET_TMPDIR") {
            Some(dir) => PathBuf::from(dir),
            None => PathBuf::from(env!("CARGO_MANIFEST_DIR"))
                .join("target")
                .join("auth-test-tmp"),
        }
    }

    /// A unique scratch file path under [`scratch_base`] (never `/tmp`).
    fn tmp(name: &str) -> PathBuf {
        let base = scratch_base();
        std::fs::create_dir_all(&base).unwrap();
        let p = base.join(format!(
            "akurai-auth-{name}-{}-{}.db",
            std::process::id(),
            ID_COUNTER.fetch_add(1, Ordering::Relaxed)
        ));
        let _ = std::fs::remove_file(&p);
        p
    }

    // --- session store ----------------------------------------------------

    #[test]
    fn session_create_lookup_delete_roundtrip() {
        let path = tmp("session-roundtrip");
        let store = SessionStore::open(&path).unwrap();
        let id = store.create("ada").unwrap();

        let found = store.lookup(&id).unwrap().expect("session should exist");
        assert_eq!(found.username, "ada");
        assert!(found.created > 0);

        assert!(store.delete(&id).unwrap(), "delete reports it existed");
        assert!(store.lookup(&id).unwrap().is_none(), "gone after delete");
        assert!(!store.delete(&id).unwrap(), "second delete is a no-op");
    }

    #[test]
    fn unknown_session_id_is_none() {
        let path = tmp("session-unknown");
        let store = SessionStore::open(&path).unwrap();
        assert!(store.lookup("does-not-exist").unwrap().is_none());
    }

    #[test]
    fn sessions_persist_across_reopen() {
        let path = tmp("session-persist");
        let id = {
            let store = SessionStore::open(&path).unwrap();
            store.create("grace").unwrap()
        };
        // Reopen from disk: the session survives.
        let store = SessionStore::open(&path).unwrap();
        let found = store.lookup(&id).unwrap().expect("session survives reopen");
        assert_eq!(found.username, "grace");
    }

    #[test]
    fn ids_are_unique_even_in_a_tight_loop() {
        let path = tmp("session-unique");
        let store = SessionStore::open(&path).unwrap();
        let mut ids = std::collections::HashSet::new();
        for _ in 0..1000 {
            assert!(ids.insert(store.create("u").unwrap()), "no id collisions");
        }
    }

    // --- password hashing -------------------------------------------------

    #[test]
    fn same_input_same_hash() {
        assert_eq!(hash_password("salt", "pw"), hash_password("salt", "pw"));
    }

    #[test]
    fn different_password_different_hash() {
        assert_ne!(hash_password("salt", "pw"), hash_password("salt", "other"));
    }

    #[test]
    fn different_salt_different_hash() {
        assert_ne!(hash_password("salt-a", "pw"), hash_password("salt-b", "pw"));
    }

    #[test]
    fn hash_is_hex_and_never_the_plaintext() {
        let h = hash_password("s", "hunter2");
        assert_eq!(h.len(), 40, "SHA-1 hex is 40 chars");
        assert!(h.bytes().all(|b| b.is_ascii_hexdigit()));
        assert!(!h.contains("hunter2"));
    }

    #[test]
    fn constant_time_eq_matches_plain_eq() {
        assert!(constant_time_eq("abc", "abc"));
        assert!(!constant_time_eq("abc", "abd"));
        assert!(!constant_time_eq("abc", "abcd"));
    }

    // --- user store -------------------------------------------------------

    #[test]
    fn user_verify_accepts_correct_rejects_wrong() {
        let path = tmp("user-verify");
        let users = UserStore::open(&path).unwrap();
        users.ensure_user("demo", "akurai").unwrap();

        assert!(users.verify("demo", "akurai").unwrap(), "right password");
        assert!(!users.verify("demo", "nope").unwrap(), "wrong password");
        assert!(!users.verify("ghost", "akurai").unwrap(), "no such user");
    }

    #[test]
    fn ensure_user_is_idempotent() {
        let path = tmp("user-idempotent");
        let users = UserStore::open(&path).unwrap();
        users.ensure_user("demo", "first").unwrap();
        // A second call with a different password must NOT overwrite the user.
        users.ensure_user("demo", "second").unwrap();
        assert!(users.verify("demo", "first").unwrap());
        assert!(!users.verify("demo", "second").unwrap());
    }

    // --- cookie helpers ---------------------------------------------------

    #[test]
    fn parse_cookie_finds_named_value() {
        let header = "theme=dark; session=abc123; lang=en";
        assert_eq!(parse_cookie(header, "session").as_deref(), Some("abc123"));
        assert_eq!(parse_cookie(header, "theme").as_deref(), Some("dark"));
        assert_eq!(parse_cookie(header, "missing"), None);
    }

    #[test]
    fn parse_cookie_handles_single_pair_and_spacing() {
        assert_eq!(
            parse_cookie("session=xyz", "session").as_deref(),
            Some("xyz")
        );
        assert_eq!(
            parse_cookie("  session = spaced ", "session").as_deref(),
            Some("spaced")
        );
    }

    #[test]
    fn set_and_expire_cookie_have_the_right_attributes() {
        let set = set_cookie("tok");
        assert!(set.starts_with("session=tok"));
        assert!(set.contains("HttpOnly"));
        assert!(set.contains("Path=/"));
        assert!(set.contains("SameSite=Lax"));

        let clear = expired_cookie();
        assert!(clear.starts_with("session=;"));
        assert!(clear.contains("Max-Age=0"));
    }

    // --- login state machine (function level, no socket) ------------------

    fn open_state(name: &str) -> State {
        let dir = scratch_base().join(format!(
            "akurai-auth-state-{name}-{}-{}",
            std::process::id(),
            ID_COUNTER.fetch_add(1, Ordering::Relaxed)
        ));
        std::fs::create_dir_all(&dir).unwrap();
        // Clean any prior run's files so seeding is deterministic.
        for f in ["sessions.db", "users.db"] {
            let _ = std::fs::remove_file(dir.join(f));
        }
        State::open(&dir).unwrap()
    }

    fn req_with_cookie(cookie: Option<&str>) -> Request {
        let head = match cookie {
            Some(c) => format!("GET /account HTTP/1.1\r\nCookie: {c}\r\n\r\n"),
            None => "GET /account HTTP/1.1\r\n\r\n".to_string(),
        };
        Request::parse_head(head.as_bytes()).unwrap()
    }

    #[test]
    fn full_login_cycle_via_state() {
        let state = open_state("cycle");

        // No cookie → not authenticated.
        assert!(state.current_user(&req_with_cookie(None)).is_none());

        // Create a session for the seeded demo user (what login_submit does).
        let id = state.sessions.create(DEMO_USERNAME).unwrap();
        let cookie = set_cookie(&id);
        // The browser would send back just `name=value`; emulate that.
        let sent = format!("{COOKIE_NAME}={id}");
        assert_eq!(
            state.current_user(&req_with_cookie(Some(&sent))).as_deref(),
            Some(DEMO_USERNAME)
        );
        assert!(cookie.contains(&id));

        // Logout deletes the session; the same cookie now resolves to nobody.
        assert!(state.sessions.delete(&id).unwrap());
        assert!(state.current_user(&req_with_cookie(Some(&sent))).is_none());
    }

    #[test]
    fn account_redirects_when_unauthenticated() {
        let state = open_state("account-guard");
        let root = Path::new("nonexistent-frontend");
        let Reply::Response(resp) = account(&state, root, &req_with_cookie(None)) else {
            panic!("expected a buffered response");
        };
        assert_eq!(resp.status, 303);
        assert_eq!(
            resp.headers
                .iter()
                .find(|(k, _)| k.eq_ignore_ascii_case("location"))
                .map(|(_, v)| v.as_str()),
            Some("/login")
        );
    }

    #[test]
    fn demo_credentials_are_seeded() {
        let state = open_state("seed");
        assert!(state.users.verify(DEMO_USERNAME, DEMO_PASSWORD).unwrap());
    }
}