AkurAI-Monitor / src/auth.rs

Latest change d255d9a040a989aa4961a614ddfb0d20568a758f - Restart expired OIDC login flows by Ólafur Búi Ólafsson

use axum::{
    extract::{Query, Request},
    http::{StatusCode, header},
    middleware::Next,
    response::{IntoResponse, Redirect, Response},
};
use serde::Deserialize;
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::sync::LazyLock;
use std::sync::Mutex;

use crate::config;

const SESSION_COOKIE: &str = "monitor_session";
const SESSION_MAX_AGE: u64 = 60 * 60 * 24; // 24 hours

// In-memory session store: session_id -> (email, created_at_unix)
static SESSIONS: LazyLock<Mutex<HashMap<String, (String, i64)>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

// Abandoned OIDC flows expire after 10 minutes
const PENDING_FLOW_TTL: i64 = 600;

fn generate_session_id() -> String {
    use rand::Rng;
    let bytes: [u8; 32] = rand::rng().random();
    hex::encode(bytes)
}

fn generate_state() -> String {
    use rand::Rng;
    let bytes: [u8; 16] = rand::rng().random();
    hex::encode(bytes)
}

// PKCE S256 code verifier + challenge
fn generate_pkce() -> (String, String) {
    use rand::Rng;
    let bytes: [u8; 32] = rand::rng().random();
    let verifier = hex::encode(bytes);
    let hash = Sha256::digest(verifier.as_bytes());
    let challenge = base64_url_encode(&hash);
    (verifier, challenge)
}

fn base64_url_encode(data: &[u8]) -> String {
    use base64_encode::encode;
    encode(data)
        .trim_end_matches('=')
        .replace('+', "-")
        .replace('/', "_")
}

// Simple base64 without pulling in the base64 crate
mod base64_encode {
    const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

    pub fn encode(data: &[u8]) -> String {
        let mut result = String::new();
        for chunk in data.chunks(3) {
            let b0 = chunk[0] as u32;
            let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
            let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
            let triple = (b0 << 16) | (b1 << 8) | b2;
            result.push(CHARS[((triple >> 18) & 0x3F) as usize] as char);
            result.push(CHARS[((triple >> 12) & 0x3F) as usize] as char);
            if chunk.len() > 1 {
                result.push(CHARS[((triple >> 6) & 0x3F) as usize] as char);
            } else {
                result.push('=');
            }
            if chunk.len() > 2 {
                result.push(CHARS[(triple & 0x3F) as usize] as char);
            } else {
                result.push('=');
            }
        }
        result
    }
}

// Pending OIDC flows: state -> (code_verifier, created_at_unix)
static PENDING_FLOWS: LazyLock<Mutex<HashMap<String, (String, i64)>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

fn get_session_email(headers: &axum::http::HeaderMap) -> Option<String> {
    let cookie_header = headers.get(header::COOKIE)?.to_str().ok()?;
    let session_id = cookie_header.split(';').find_map(|pair| {
        let pair = pair.trim();
        let (k, v) = pair.split_once('=')?;
        if k.trim() == SESSION_COOKIE {
            Some(v.trim().to_string())
        } else {
            None
        }
    })?;

    let mut sessions = SESSIONS.lock().unwrap();
    match sessions.get(&session_id) {
        Some((email, created))
            if chrono::Utc::now().timestamp() - created < SESSION_MAX_AGE as i64 =>
        {
            Some(email.clone())
        }
        Some(_) => {
            // Expired server-side — drop it so the session can't outlive the cookie.
            sessions.remove(&session_id);
            None
        }
        None => None,
    }
}

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

pub async fn auth_middleware(request: Request, next: Next) -> Response {
    let path = request.uri().path().to_string();

    // Allow health, log ingestion (token-authed in the handler), auth routes,
    // browser icon assets, and theme assets (needed before session exists so the
    // login redirect page renders with the correct theme).
    if path == "/api/health"
        || path == "/api/ingest"
        || path.starts_with("/auth/")
        || path == "/favicon.svg"
        || path == "/apple-touch-icon.png"
        || path == "/theme.js"
        || path == "/themes.css"
        || path == "/themes.json"
    {
        return next.run(request).await;
    }

    if path.starts_with("/api/") {
        let token = config::get().api_token.as_bytes();
        let provided = request
            .headers()
            .get(header::AUTHORIZATION)
            .and_then(|value| value.to_str().ok())
            .and_then(|value| value.strip_prefix("Bearer "));

        if !token.is_empty()
            && provided.map(|value| constant_time_eq(value.as_bytes(), token)) == Some(true)
        {
            return next.run(request).await;
        }
        if provided.is_some() {
            return StatusCode::UNAUTHORIZED.into_response();
        }
    }

    if get_session_email(request.headers()).is_some() {
        return next.run(request).await;
    }

    // Not authenticated — redirect to login
    let cfg = config::get();
    let state = generate_state();
    let (verifier, challenge) = generate_pkce();

    {
        let mut flows = PENDING_FLOWS.lock().unwrap();
        let now = chrono::Utc::now().timestamp();
        // Prune abandoned flows so crawlers hitting protected paths can't grow the map unboundedly.
        flows.retain(|_, (_, created)| now - *created < PENDING_FLOW_TTL);
        flows.insert(state.clone(), (verifier, now));
    }

    let authorize_url = format!(
        "{}/authorize?client_id={}&redirect_uri={}&response_type=code&scope=openid+email&state={}&code_challenge={}&code_challenge_method=S256",
        cfg.oidc_issuer,
        cfg.oidc_client_id,
        urlenc(&cfg.oidc_redirect_uri),
        state,
        challenge,
    );

    Redirect::temporary(&authorize_url).into_response()
}

#[derive(Deserialize)]
pub struct CallbackQuery {
    pub code: Option<String>,
    pub state: Option<String>,
    pub error: Option<String>,
    pub error_description: Option<String>,
}

fn restart_login(reason: &'static str) -> Response {
    tracing::warn!(reason, "restarting expired OIDC login flow");
    Redirect::to("/").into_response()
}

pub async fn auth_callback(Query(q): Query<CallbackQuery>) -> impl IntoResponse {
    // Handle IDP error responses
    if let Some(error) = &q.error {
        let desc = q.error_description.as_deref().unwrap_or("Unknown error");
        tracing::error!(error, desc, "OIDC authorization error");
        return (StatusCode::UNAUTHORIZED, format!("Login failed: {desc}")).into_response();
    }

    let Some(code) = &q.code else {
        return (StatusCode::BAD_REQUEST, "Missing authorization code").into_response();
    };

    let Some(state) = &q.state else {
        return restart_login("missing state");
    };

    let cfg = config::get();

    // Retrieve and remove the pending flow
    let verifier = {
        let mut flows = PENDING_FLOWS.lock().unwrap();
        flows.remove(state)
    };

    let Some((verifier, _)) = verifier else {
        return restart_login("unknown or expired state");
    };

    // Exchange authorization code for tokens
    let client = reqwest::Client::new();
    let token_url = format!("{}/token", cfg.oidc_issuer);

    let res = client
        .post(&token_url)
        .form(&[
            ("grant_type", "authorization_code"),
            ("code", code),
            ("redirect_uri", &cfg.oidc_redirect_uri),
            ("client_id", &cfg.oidc_client_id),
            ("client_secret", &cfg.oidc_client_secret),
            ("code_verifier", &verifier),
        ])
        .send()
        .await;

    let Ok(res) = res else {
        return (StatusCode::BAD_GATEWAY, "Failed to contact IDP").into_response();
    };

    if !res.status().is_success() {
        let body = res.text().await.unwrap_or_default();
        tracing::error!(body, "token exchange failed");
        return (StatusCode::UNAUTHORIZED, "Token exchange failed").into_response();
    }

    let Ok(token_response) = res.json::<serde_json::Value>().await else {
        return (StatusCode::BAD_GATEWAY, "Invalid token response").into_response();
    };

    let Some(access_token) = token_response
        .get("access_token")
        .and_then(|value| value.as_str())
    else {
        return (
            StatusCode::BAD_GATEWAY,
            "Token response omitted access token",
        )
            .into_response();
    };
    let userinfo = client
        .get(format!("{}/userinfo", cfg.oidc_issuer))
        .bearer_auth(access_token)
        .send()
        .await;
    let Ok(userinfo) = userinfo else {
        return (StatusCode::BAD_GATEWAY, "Failed to verify identity").into_response();
    };
    if !userinfo.status().is_success() {
        return (StatusCode::UNAUTHORIZED, "Identity verification failed").into_response();
    }
    let Ok(claims) = userinfo.json::<serde_json::Value>().await else {
        return (StatusCode::BAD_GATEWAY, "Invalid userinfo response").into_response();
    };
    let Some(email) = claims
        .get("email")
        .and_then(|value| value.as_str())
        .filter(|email| !email.is_empty())
        .map(str::to_owned)
    else {
        return (StatusCode::UNAUTHORIZED, "Verified identity has no email").into_response();
    };

    // Create session
    let session_id = generate_session_id();
    SESSIONS.lock().unwrap().insert(
        session_id.clone(),
        (email.clone(), chrono::Utc::now().timestamp()),
    );

    tracing::info!(email, "user authenticated via OIDC");

    let cookie = format!(
        "{SESSION_COOKIE}={session_id}; Path=/; HttpOnly; Secure; SameSite=Lax; Max-Age={SESSION_MAX_AGE}"
    );

    (
        StatusCode::SEE_OTHER,
        [
            (header::SET_COOKIE, cookie),
            (header::LOCATION, "/".to_string()),
        ],
    )
        .into_response()
}

pub async fn auth_logout() -> impl IntoResponse {
    let cookie = format!("{SESSION_COOKIE}=; Path=/; HttpOnly; Secure; SameSite=Lax; Max-Age=0");
    (
        StatusCode::SEE_OTHER,
        [
            (header::SET_COOKIE, cookie),
            (header::LOCATION, "/".to_string()),
        ],
    )
}

fn urlenc(s: &str) -> String {
    let mut result = String::new();
    for b in s.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                result.push(b as char);
            }
            _ => {
                result.push_str(&format!("%{:02X}", b));
            }
        }
    }
    result
}

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

    #[tokio::test]
    async fn stale_oidc_state_restarts_login() {
        let response = auth_callback(Query(CallbackQuery {
            code: Some("stale-code".to_string()),
            state: Some("stale-state".to_string()),
            error: None,
            error_description: None,
        }))
        .await
        .into_response();

        assert_eq!(response.status(), StatusCode::SEE_OTHER);
        assert_eq!(
            response
                .headers()
                .get(header::LOCATION)
                .expect("redirect includes location"),
            "/"
        );
    }
}