neiam /fieri-forever

//! Cookie-session authentication middleware.
//!
//! Two layers are provided:
//! * [`web_auth_middleware`] — for protected pages; redirects anonymous
//!   visitors to `/login`.
//! * [`optional_auth_middleware`] — for pages that render differently when
//!   signed in but are also reachable anonymously (e.g. the landing page).
//!
//! Both populate a [`UserContext`] request extension, which handlers pull in
//! via the `UserContext` / `Option<UserContext>` extractors.
use axum::{
    extract::{FromRequestParts, OptionalFromRequestParts, Request, State},
    http::{header::COOKIE, request::Parts, StatusCode},
    middleware::Next,
    response::{IntoResponse, Redirect, Response},
};
use sea_orm::DatabaseConnection;
use std::convert::Infallible;
use std::sync::Arc;

use crate::auth;
use crate::entities::{user_sessions, users};

/// The authenticated user and their session, attached to a request.
#[derive(Clone, Debug)]
pub struct UserContext {
    pub user: users::Model,
    pub session: user_sessions::Model,
}

impl<S: Send + Sync> FromRequestParts<S> for UserContext {
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        parts
            .extensions
            .get::<UserContext>()
            .cloned()
            .ok_or(StatusCode::UNAUTHORIZED)
    }
}

impl<S: Send + Sync> OptionalFromRequestParts<S> for UserContext {
    type Rejection = Infallible;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Option<Self>, Self::Rejection> {
        Ok(parts.extensions.get::<UserContext>().cloned())
    }
}

/// Pull the `session_token` cookie out of a request, if present.
fn extract_session_token(request: &Request) -> Option<String> {
    let cookies = request.headers().get(COOKIE)?.to_str().ok()?;
    for cookie in cookies.split(';') {
        if let Some(token) = cookie.trim().strip_prefix("session_token=") {
            return Some(token.to_string());
        }
    }
    None
}

/// Attach a `UserContext` when a valid session cookie is present, but never
/// reject — anonymous requests pass straight through.
pub async fn optional_auth_middleware(
    State(db): State<Arc<DatabaseConnection>>,
    mut request: Request,
    next: Next,
) -> Response {
    if let Some(token) = extract_session_token(&request) {
        if let Ok(Some((user, session))) = auth::validate_session(&db, &token).await {
            request.extensions_mut().insert(UserContext { user, session });
        }
    }
    next.run(request).await
}

/// Require a valid session; redirect to `/login` otherwise.
pub async fn web_auth_middleware(
    State(db): State<Arc<DatabaseConnection>>,
    mut request: Request,
    next: Next,
) -> Response {
    let Some(token) = extract_session_token(&request) else {
        return Redirect::temporary("/login").into_response();
    };
    match auth::validate_session(&db, &token).await {
        Ok(Some((user, session))) => {
            request.extensions_mut().insert(UserContext { user, session });
            next.run(request).await
        }
        Ok(None) | Err(_) => Redirect::temporary("/login").into_response(),
    }
}

/// Convenience accessors for the `UserContext` extension on a `Request`.
pub trait RequestUserExt {
    fn user_context(&self) -> Option<&UserContext>;
    fn current_user(&self) -> Option<&users::Model>;
    fn current_session(&self) -> Option<&user_sessions::Model>;
}

impl RequestUserExt for Request {
    fn user_context(&self) -> Option<&UserContext> {
        self.extensions().get::<UserContext>()
    }
    fn current_user(&self) -> Option<&users::Model> {
        self.user_context().map(|c| &c.user)
    }
    fn current_session(&self) -> Option<&user_sessions::Model> {
        self.user_context().map(|c| &c.session)
    }
}

/// The user behind an API request, from its `Authorization: Bearer` token.
/// Separate from [`UserContext`] because there's no cookie session behind
/// it; the bearer token itself is kept for the app's "sign out".
#[derive(Clone, Debug)]
pub struct ApiUser {
    pub user: users::Model,
    pub token: String,
}

impl<S: Send + Sync> FromRequestParts<S> for ApiUser {
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        parts.extensions.get::<ApiUser>().cloned().ok_or(StatusCode::UNAUTHORIZED)
    }
}

fn extract_bearer_token(request: &Request) -> Option<String> {
    let header = request.headers().get(axum::http::header::AUTHORIZATION)?.to_str().ok()?;
    let token = header.strip_prefix("Bearer ").or_else(|| header.strip_prefix("bearer "))?.trim();
    (!token.is_empty()).then(|| token.to_string())
}

/// For `/api/*`: require a valid bearer token, else 401. No redirects and no
/// cookies — this is for the Android app and scripts.
pub async fn api_auth_middleware(
    State(db): State<Arc<DatabaseConnection>>,
    mut request: Request,
    next: Next,
) -> Response {
    let Some(token) = extract_bearer_token(&request) else {
        return StatusCode::UNAUTHORIZED.into_response();
    };
    match auth::api_tokens::validate(&db, &token).await {
        Ok(Some(user)) => {
            request.extensions_mut().insert(ApiUser { user, token });
            next.run(request).await
        }
        Ok(None) => StatusCode::UNAUTHORIZED.into_response(),
        Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    }
}