AkurAI-Drive / backend/src/auth.rs

Latest change 6e3a0824f4ec4d2bea2524a20b4960c7ab9e2766 - backend: akurai-drive-api (OIDC RP + per-tenant S3 provisioning) by olibuijr

use axum::{
    extract::Query,
    http::{HeaderMap, StatusCode, header},
    response::{IntoResponse, Response},
};
use serde::Deserialize;

use crate::config;
use crate::session;

/// `GET /auth/login` — mint a random state, stash it in a short-lived cookie, and
/// 302 to the IDP `/authorize` endpoint.
pub async fn login() -> Response {
    let cfg = config::get();
    let state = crate::util::token();

    // Build the authorize query. form_urlencoded encodes spaces as '+'; the only
    // value containing spaces is `scope`, so swapping '+' -> %20 yields exactly the
    // `openid%20profile%20email%20groups` form without touching anything else.
    let query = form_urlencoded::Serializer::new(String::new())
        .append_pair("client_id", &cfg.oidc_client_id)
        .append_pair("redirect_uri", &cfg.oidc_redirect_uri)
        .append_pair("response_type", "code")
        .append_pair("scope", "openid profile email groups")
        .append_pair("state", &state)
        .finish()
        .replace('+', "%20");

    let location = format!("{}/authorize?{query}", cfg.idp_issuer);
    let cookie =
        format!("drive_oauth_state={state}; Path=/; HttpOnly; SameSite=Lax; Secure; Max-Age=600");

    (
        StatusCode::FOUND,
        [(header::SET_COOKIE, cookie), (header::LOCATION, location)],
    )
        .into_response()
}

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

/// `GET /auth/callback` — verify state, exchange the code, create a session, set
/// the `drive_session` cookie and 302 to `/app`. Any failure 302s to `/?auth_error=1`.
pub async fn callback(headers: HeaderMap, Query(q): Query<CallbackQuery>) -> Response {
    let (code, state) = match (q.code, q.state) {
        (Some(c), Some(s)) if !c.is_empty() && !s.is_empty() => (c, s),
        _ => return auth_error(),
    };

    // CSRF: the state in the query must match the cookie we set at /auth/login.
    match session::get_cookie(&headers, "drive_oauth_state") {
        Some(cookie_state) if cookie_state == state => {}
        _ => return auth_error(),
    }

    // Back-channel token + userinfo (blocking ureq) on the blocking pool.
    let userinfo =
        match tokio::task::spawn_blocking(move || crate::idp::exchange_code_for_userinfo(&code))
            .await
        {
            Ok(Ok(info)) => info,
            Ok(Err(e)) => {
                tracing::warn!("oidc exchange failed: {e}");
                return auth_error();
            }
            Err(e) => {
                tracing::error!("oidc exchange task panicked: {e}");
                return auth_error();
            }
        };

    let session_id = session::create_session(
        &userinfo.sub,
        &userinfo.email,
        &userinfo.groups,
        userinfo.tenant_id.as_deref(),
    );

    // ~30 days = 2592000s.
    let cookie = format!(
        "drive_session={session_id}; Path=/; HttpOnly; SameSite=Lax; Secure; Max-Age=2592000"
    );

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

/// `GET /auth/logout` — drop the session row, clear the cookie, 302 to `/`.
pub async fn logout(headers: HeaderMap) -> Response {
    if let Some(id) = session::get_cookie(&headers, "drive_session") {
        session::delete_session(&id);
    }
    let clear = "drive_session=; Path=/; HttpOnly; SameSite=Lax; Secure; Max-Age=0".to_string();

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

fn auth_error() -> Response {
    (
        StatusCode::FOUND,
        [(header::LOCATION, "/?auth_error=1".to_string())],
    )
        .into_response()
}