AkurAI-Drive / backend/src/auth.rs

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

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

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

#[derive(Deserialize)]
pub struct LoginQuery {
    /// Optional local path (must start with '/') to land on after login, carried
    /// across the OAuth round-trip via the `state`. Off-site values are ignored.
    return_to: Option<String>,
}

/// `GET /auth/login` — start an OIDC authorization-code + PKCE login.
pub async fn login(Query(q): Query<LoginQuery>) -> Response {
    begin_login(q.return_to.as_deref())
}

/// Only ever resume to a path on our own server — never an attacker-supplied
/// absolute or protocol-relative URL (open-redirect guard).
pub fn safe_return_to(return_to: Option<&str>) -> String {
    match return_to {
        Some(rt) if rt.starts_with('/') && !rt.starts_with("//") && !rt.contains('\\') => {
            rt.to_string()
        }
        _ => "/app".to_string(),
    }
}

/// Shared login kickoff used by `/auth/login` and `/auth/device`: store the
/// nonce + PKCE verifier server-side under a fresh `state`, bind the state to
/// this browser with a cookie, and 302 to the IdP authorize endpoint.
fn begin_login(return_to: Option<&str>) -> Response {
    let state = crate::util::token();
    let login = session::LoginState {
        nonce: crate::util::token(),
        verifier: crate::util::token(),
        return_to: safe_return_to(return_to),
    };
    session::store_login_state(&state, &login);
    let location = oidc::authorize_url(config::get(), &state, &login.nonce, &login.verifier);

    (
        StatusCode::FOUND,
        [
            (header::SET_COOKIE, session::state_cookie(&state)),
            (header::LOCATION, location),
        ],
    )
        .into_response()
}

#[derive(Deserialize)]
pub struct DeviceQuery {
    /// Loopback URL the CLI listens on; the device token is appended as `?token=`.
    redirect: Option<String>,
}

/// `GET /auth/device?redirect=<loopback-url>` — the CLI login bridge.
///
/// `redirect` must be a loopback URL. With a live web session we mint a long-lived
/// device token and 302 to `<redirect>?token=<token>`. Without one we kick off the
/// normal OIDC login, carrying this device URL as `return_to` so the browser comes
/// straight back here (now authenticated) to collect the token.
pub async fn device(headers: HeaderMap, Query(q): Query<DeviceQuery>) -> Response {
    let redirect = match q.redirect {
        Some(r) if !r.is_empty() => r,
        _ => return bad_request("redirect is required"),
    };
    if !crate::device::is_loopback_redirect(&redirect) {
        return bad_request(
            "redirect must be a loopback URL (http://127.0.0.1:* or http://localhost:*)",
        );
    }

    match session::lookup_session(&headers) {
        Some(s) => {
            let token = crate::device::create_device_token(&s.sub, &s.email);
            let sep = if redirect.contains('?') { '&' } else { '?' };
            let location = format!("{redirect}{sep}token={token}");
            (StatusCode::FOUND, [(header::LOCATION, location)]).into_response()
        }
        None => {
            // No session yet: log in, then return to this same /auth/device URL.
            let return_to = format!(
                "/auth/device?{}",
                form_urlencoded::Serializer::new(String::new())
                    .append_pair("redirect", &redirect)
                    .finish()
            );
            begin_login(Some(&return_to))
        }
    }
}

fn bad_request(msg: &str) -> Response {
    (StatusCode::BAD_REQUEST, Json(json!({ "error": msg }))).into_response()
}

#[derive(Deserialize)]
pub struct PasswordBody {
    email: Option<String>,
    password: Option<String>,
}

/// `POST /auth/password` `{email, password}` — headless native sign-in for the
/// desktop/CLI clients, so they never have to open a browser.
///
/// The IDP has no password grant, so we run its browser login flow server-side
/// over localhost (see [`crate::idp::password_login`]): login → authorize → token
/// → userinfo. On success we mint a device token (exactly as `GET /auth/device`)
/// and hand it back. Bad credentials → 401; an MFA-protected account → 409 (use
/// the browser flow). The password is never logged.
pub async fn password(Json(body): Json<PasswordBody>) -> Response {
    let email = body.email.unwrap_or_default().trim().to_string();
    let password = body.password.unwrap_or_default();
    if email.is_empty() || password.is_empty() {
        return bad_request("email and password are required");
    }

    let outcome =
        tokio::task::spawn_blocking(move || crate::idp::password_login(&email, &password)).await;

    match outcome {
        Ok(Ok(crate::idp::PasswordLogin::Ok(info))) => {
            let token = crate::device::create_device_token(&info.sub, &info.email);
            (
                StatusCode::OK,
                Json(json!({ "token": token, "email": info.email })),
            )
                .into_response()
        }
        Ok(Ok(crate::idp::PasswordLogin::InvalidCredentials)) => (
            StatusCode::UNAUTHORIZED,
            Json(json!({ "error": "invalid email or password" })),
        )
            .into_response(),
        Ok(Ok(crate::idp::PasswordLogin::MfaRequired)) => (
            StatusCode::CONFLICT,
            Json(json!({
                "error": "mfa_required",
                "detail": "This account has MFA — use browser sign-in.",
            })),
        )
            .into_response(),
        Ok(Err(e)) => {
            tracing::warn!("password login failed: {e}");
            (
                StatusCode::BAD_GATEWAY,
                Json(json!({ "error": "login failed" })),
            )
                .into_response()
        }
        Err(e) => {
            tracing::error!("password login task panicked: {e}");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": "login failed" })),
            )
                .into_response()
        }
    }
}

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

/// Check the callback's `state` against the login cookie; returns the code.
pub fn check_callback(
    q: CallbackQuery,
    cookie_state: Option<&str>,
) -> Result<(String, String), &'static str> {
    let state = q.state.filter(|s| !s.is_empty()).ok_or("missing state")?;
    if cookie_state != Some(state.as_str()) {
        return Err("state does not match the login cookie");
    }
    let code = q.code.filter(|c| !c.is_empty()).ok_or("missing code")?;
    Ok((code, state))
}

/// `GET /auth/callback` — verify state, run the back channel, create a session
/// and resume to the stored `return_to`.
pub async fn callback(headers: HeaderMap, Query(q): Query<CallbackQuery>) -> Response {
    let cookie_state = session::get_cookie(&headers, session::STATE_COOKIE);
    let (code, state) = match check_callback(q, cookie_state.as_deref()) {
        Ok(v) => v,
        Err(msg) => return bad_request(msg),
    };
    let Some(login) = session::take_login_state(&state) else {
        return bad_request("unknown or expired state");
    };

    let (verifier, nonce) = (login.verifier, login.nonce);
    let outcome =
        tokio::task::spawn_blocking(move || oidc::complete_login(&code, &verifier, &nonce)).await;
    let identity = match outcome {
        Ok(Ok(identity)) => identity,
        Ok(Err(oidc::LoginError::Denied)) => {
            tracing::warn!("login denied: no allowed group");
            return (
                StatusCode::FORBIDDEN,
                "Your account is not permitted to use AkurAI Drive.",
            )
                .into_response();
        }
        Ok(Err(oidc::LoginError::Failed(e))) => {
            tracing::warn!("oidc callback rejected: {e}");
            return login_failed();
        }
        Err(e) => {
            tracing::error!("oidc callback task panicked: {e}");
            return login_failed();
        }
    };

    let session_id = session::create_session(&identity);
    let mut response = (StatusCode::FOUND, [(header::LOCATION, login.return_to)]).into_response();
    let jar = response.headers_mut();
    for c in [
        session::session_cookie(&session_id),
        session::cookie(session::STATE_COOKIE, None, 0),
    ] {
        if let Ok(v) = header::HeaderValue::from_str(&c) {
            jar.append(header::SET_COOKIE, v);
        }
    }
    response
}

/// `GET /auth/logout` — drop the local session, clear the cookie, and send the
/// browser to the IdP's `/logout` to end the SSO session too.
pub async fn logout(headers: HeaderMap) -> Response {
    if let Some(id) = session::get_cookie(&headers, session::SESSION_COOKIE) {
        session::delete_session(&id);
    }
    (
        StatusCode::FOUND,
        [
            (
                header::SET_COOKIE,
                session::cookie(session::SESSION_COOKIE, None, 0),
            ),
            (header::LOCATION, oidc::logout_url(config::get())),
        ],
    )
        .into_response()
}

fn login_failed() -> Response {
    (
        StatusCode::BAD_GATEWAY,
        Json(json!({ "error": "login could not be completed" })),
    )
        .into_response()
}

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

    fn q(code: Option<&str>, state: Option<&str>) -> CallbackQuery {
        CallbackQuery {
            code: code.map(String::from),
            state: state.map(String::from),
        }
    }

    #[test]
    fn return_to_is_local_only() {
        assert_eq!(safe_return_to(Some("/auth/device?x=1")), "/auth/device?x=1");
        assert_eq!(safe_return_to(None), "/app");
        assert_eq!(safe_return_to(Some("https://evil")), "/app");
        assert_eq!(safe_return_to(Some("//evil")), "/app");
        assert_eq!(safe_return_to(Some("/\\evil")), "/app");
    }

    #[test]
    fn callback_state_must_match_cookie() {
        assert_eq!(
            check_callback(q(Some("c"), Some("s")), Some("s")),
            Ok(("c".into(), "s".into()))
        );
        assert_eq!(
            check_callback(q(Some("c"), None), Some("s")),
            Err("missing state")
        );
        assert_eq!(
            check_callback(q(Some("c"), Some("s")), Some("t")),
            Err("state does not match the login cookie")
        );
        assert_eq!(
            check_callback(q(Some("c"), Some("s")), None),
            Err("state does not match the login cookie")
        );
        assert_eq!(
            check_callback(q(None, Some("s")), Some("s")),
            Err("missing code")
        );
    }
}