AkurAI-Drive / backend/src/auth.rs

Latest change 2a1ed73acfbc6091a9fcdba5a28f55de9e2d75a4 - feat: native email+password sign-in (no browser) by olibuijr

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

use crate::config;
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` — mint a random state, stash it in a short-lived cookie, and
/// 302 to the IDP `/authorize` endpoint.
pub async fn login(Query(q): Query<LoginQuery>) -> Response {
    begin_login(q.return_to.as_deref())
}

/// Shared OIDC-login kickoff used by `/auth/login` and `/auth/device`. Mints a
/// state, optionally remembers a local `return_to`, sets the CSRF cookie, and 302s
/// to the IDP authorize endpoint.
fn begin_login(return_to: Option<&str>) -> Response {
    let cfg = config::get();
    let state = crate::util::token();

    // Only ever resume to a path on our own server — never an attacker-supplied
    // absolute URL (open-redirect guard).
    if let Some(rt) = return_to
        && rt.starts_with('/')
    {
        session::store_oauth_return(&state, rt);
    }

    // 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 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>,
}

/// `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(),
    );

    // Resume to a remembered local path (e.g. /auth/device?...) if one was carried
    // through the `state`; otherwise the dashboard. Only local paths are honored.
    let location = session::take_oauth_return(&state)
        .filter(|rt| rt.starts_with('/'))
        .unwrap_or_else(|| "/app".to_string());

    // ~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, location)],
    )
        .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()
}