alibi 0.10.0

//! Axum router, session layer and session extractors. Enable the `axum` feature.

use super::dispatch::DispatchSupervisor;
use super::shared::{
    CurrentSession, OptionalSession, auth_request, check_content_length, max_body_bytes,
    payload_too_large, request_headers,
};
use crate::Alibi;
use alibi_core::endpoint::EndpointOptions;
use alibi_core::{
    AuthError, AuthRequest, AuthResponse, AuthResult, AuthSchema, Headers, core_paths,
};
use alibi_plugins::{GetSessionQuery, GetSessionResponse, SessionManagementPlugin};
use axum::{
    Router,
    extract::{FromRef, FromRequestParts, OptionalFromRequestParts, OriginalUri, Request, State},
    http::{HeaderMap, HeaderName, HeaderValue, header, request::Parts},
    response::{IntoResponse, Response},
    routing::any,
};
use std::{
    collections::HashSet,
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tower::{Layer, Service};

/// Mount an initialized auth instance in an Axum application.
pub trait AxumIntegration {
    type Schema: AuthSchema;

    /// Create a router with every authentication route.
    ///
    /// Fully buffered requests continue dispatch after client disconnect. Router
    /// drop drains accepted work while its Tokio runtime remains alive. Server
    /// graceful shutdown alone need not await disconnected work; runtime/process
    /// shutdown can cancel it. Only the framework request context and tracing
    /// span are carried into dispatch, not arbitrary caller task-local values.
    fn axum_router(self) -> Router<Arc<Alibi<Self::Schema>>>;

    /// Create a router to nest into an application with its own state type.
    fn axum_router_with_state<S>(self) -> Router<S>
    where
        Self: Sized,
        Arc<Alibi<Self::Schema>>: FromRef<S>,
        S: Clone + Send + Sync + 'static;
}

impl<T: AuthSchema> AxumIntegration for Arc<Alibi<T>> {
    type Schema = T;

    fn axum_router(self) -> Router<Self> {
        self.axum_router_with_state::<Self>()
    }

    fn axum_router_with_state<S>(self) -> Router<S>
    where
        Self: FromRef<S>,
        S: Clone + Send + Sync + 'static,
    {
        // Dispatch owns disabled paths, method matching and trailing-slash policy.
        // Axum's method router would add Allow/HEAD behavior upstream lacks.
        let supervisor = DispatchSupervisor::new(render);
        let mut paths = HashSet::from([
            core_paths::OK.to_owned(),
            core_paths::ERROR.to_owned(),
            core_paths::UPDATE_USER.to_owned(),
        ]);
        for plugin in self.plugins() {
            paths.extend(plugin.routes().into_iter().map(|route| route.path));
        }
        let mut router = Router::new();
        for path in paths {
            router = router.route(&path, any(handler::<T>(supervisor.clone())));
        }
        router.fallback(handler::<T>(supervisor))
    }
}

impl<S, T> FromRequestParts<S> for CurrentSession<T>
where
    T: AuthSchema,
    Arc<Alibi<T>>: FromRef<S>,
    S: Send + Sync,
{
    type Rejection = Response;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let auth = Arc::<Alibi<T>>::from_ref(state);
        Self::resolve(&auth, parts.uri.path(), &parts.headers)
            .await
            .map_err(IntoResponse::into_response)
    }
}

impl<S, T> FromRequestParts<S> for OptionalSession<T>
where
    T: AuthSchema,
    Arc<Alibi<T>>: FromRef<S>,
    S: Send + Sync,
{
    type Rejection = Response;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        Ok(Self(
            CurrentSession::<T>::from_request_parts(parts, state)
                .await
                .ok(),
        ))
    }
}

/// Resolves the session of each request through get-session, as Better Auth's
/// `getSessionFromCtx` does for its own endpoints.
///
/// This applies the cookie cache, renews an aging session, clears an ended one
/// and appends the resulting `Set-Cookie` headers to the response, including
/// rejections. Handlers read the outcome with [`ResolvedSession`]. A failed
/// lookup counts as no session.
///
/// ```rust,no_run
/// use alibi::{Alibi, AuthSchema};
/// use alibi::integrations::axum::{ResolvedSession, SessionLayer};
/// use axum::{Json, Router, routing::get};
/// use std::sync::Arc;
///
/// async fn profile(ResolvedSession(session): ResolvedSession) -> Json<serde_json::Value> {
///     Json(serde_json::json!({ "id": session.user.id }))
/// }
///
/// fn router<S: AuthSchema>(auth: Arc<Alibi<S>>) -> Router {
///     Router::new()
///         .route("/profile", get(profile))
///         .layer(SessionLayer::new(auth))
/// }
/// ```
pub struct SessionLayer<S: AuthSchema> {
    auth: Arc<Alibi<S>>,
    query: Option<GetSessionQuery>,
}

impl<S: AuthSchema> SessionLayer<S> {
    pub const fn new(auth: Arc<Alibi<S>>) -> Self {
        Self { auth, query: None }
    }

    /// Bypass the cookie cache or suppress refresh for every request.
    #[must_use]
    pub const fn query(mut self, query: GetSessionQuery) -> Self {
        self.query = Some(query);
        self
    }

    async fn get_session(
        &self,
        headers: &HeaderMap,
    ) -> AuthResult<(Option<GetSessionResponse>, Headers)> {
        let mut endpoint = SessionManagementPlugin::get_session_endpoint();
        if let Some(query) = &self.query {
            endpoint = endpoint.with_query(query)?;
        }
        let options = EndpointOptions {
            headers: Some(request_headers(headers)),
            ..EndpointOptions::default()
        };
        let output = Box::pin(self.auth.dispatch_endpoint(endpoint, options))
            .await
            .map_err(|error| error.error)?;
        Ok((output.decode()?, output.headers().clone()))
    }
}

impl<S: AuthSchema> Clone for SessionLayer<S> {
    fn clone(&self) -> Self {
        Self {
            auth: Arc::clone(&self.auth),
            query: self.query.clone(),
        }
    }
}

impl<S: AuthSchema, I> Layer<I> for SessionLayer<S> {
    type Service = SessionService<S, I>;

    fn layer(&self, inner: I) -> Self::Service {
        SessionService {
            layer: self.clone(),
            inner,
        }
    }
}

/// The service produced by [`SessionLayer`].
pub struct SessionService<S: AuthSchema, I> {
    layer: SessionLayer<S>,
    inner: I,
}

impl<S: AuthSchema, I: Clone> Clone for SessionService<S, I> {
    fn clone(&self) -> Self {
        Self {
            layer: self.layer.clone(),
            inner: self.inner.clone(),
        }
    }
}

impl<S, I> Service<Request> for SessionService<S, I>
where
    S: AuthSchema,
    I: Service<Request, Response = Response> + Clone + Send + 'static,
    I::Future: Send,
{
    type Response = Response;
    type Error = I::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Response, I::Error>> + Send>>;

    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(context)
    }

    fn call(&mut self, mut request: Request) -> Self::Future {
        // Call the instance that was polled ready and keep a fresh clone.
        let ready = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, ready);
        let layer = self.layer.clone();
        Box::pin(async move {
            // An outer layer already resolved this request; reuse its result.
            if request.extensions().get::<Resolution>().is_some() {
                return inner.call(request).await;
            }
            let (session, headers) =
                layer
                    .get_session(request.headers())
                    .await
                    .unwrap_or_else(|error| {
                        tracing::warn!(%error, "Failed to get session");
                        (None, Headers::new())
                    });
            _ = request.extensions_mut().insert(Resolution(session));
            let mut response = inner.call(request).await?;
            merge_session_headers(response.headers_mut(), headers);
            Ok(response)
        })
    }
}

/// Merge get-session headers as upstream response headers do: session cookies
/// come before the handler's, and other handler headers win. Cache headers
/// belong to `/get-session` itself and are dropped.
fn merge_session_headers(response: &mut HeaderMap, session: Headers) {
    let handler_cookies: Vec<_> = response
        .get_all(header::SET_COOKIE)
        .iter()
        .cloned()
        .collect();
    _ = response.remove(header::SET_COOKIE);
    for (name, value) in session {
        let (Ok(name), Ok(value)) = (HeaderName::try_from(name), HeaderValue::try_from(value))
        else {
            continue;
        };
        if name == header::SET_COOKIE {
            _ = response.append(name, value);
        } else if name != header::CACHE_CONTROL
            && name != header::PRAGMA
            && !response.contains_key(&name)
        {
            _ = response.insert(name, value);
        }
    }
    for cookie in handler_cookies {
        _ = response.append(header::SET_COOKIE, cookie);
    }
}

#[derive(Clone)]
struct Resolution(Option<GetSessionResponse>);

/// The session [`SessionLayer`] resolved for the request.
///
/// Rejects with `401` when there is no valid session; extract
/// `Option<ResolvedSession>` for routes that also serve guests. Both reject with
/// `500` when the layer is not installed.
#[derive(Debug, Clone)]
pub struct ResolvedSession(pub GetSessionResponse);

impl ResolvedSession {
    fn from_parts(parts: &Parts) -> AuthResult<Option<Self>> {
        parts
            .extensions
            .get::<Resolution>()
            .map(|Resolution(session)| session.clone().map(Self))
            .ok_or_else(|| AuthError::internal("SessionLayer is not installed"))
    }
}

impl<T: Send + Sync> FromRequestParts<T> for ResolvedSession {
    type Rejection = AuthError;

    fn from_request_parts(
        parts: &mut Parts,
        _state: &T,
    ) -> impl Future<Output = AuthResult<Self>> + Send {
        std::future::ready(Self::from_parts(parts).and_then(|session| {
            session.ok_or(AuthError::Upstream {
                status: 401,
                code: "UNAUTHORIZED",
                message: "Unauthorized",
            })
        }))
    }
}

impl<T: Send + Sync> OptionalFromRequestParts<T> for ResolvedSession {
    type Rejection = AuthError;

    fn from_request_parts(
        parts: &mut Parts,
        _state: &T,
    ) -> impl Future<Output = AuthResult<Option<Self>>> + Send {
        std::future::ready(Self::from_parts(parts))
    }
}

type HandlerFuture = Pin<Box<dyn Future<Output = Response> + Send>>;

fn handler<T: AuthSchema>(
    supervisor: DispatchSupervisor<Response>,
) -> impl Fn(State<Arc<Alibi<T>>>, Request) -> HandlerFuture + Clone {
    move |State(auth): State<Arc<Alibi<T>>>, request: Request| {
        let supervisor = supervisor.clone();
        Box::pin(async move {
            match convert_request(&auth, request).await {
                Ok(request) => supervisor.dispatch(auth, request).await,
                Err(error) => error.into_response(),
            }
        })
    }
}

async fn convert_request<T: AuthSchema>(
    auth: &Alibi<T>,
    request: Request,
) -> AuthResult<AuthRequest> {
    let max_bytes = max_body_bytes(auth);
    let (parts, body) = request.into_parts();
    check_content_length(&parts.headers, max_bytes)?;
    let body = axum::body::to_bytes(body, max_bytes)
        .await
        .map_err(|error| {
            if is_length_limit_error(&error) {
                return payload_too_large(max_bytes);
            }
            tracing::warn!(%error, "Failed to read request body");
            AuthError::bad_request("Failed to read request body")
        })?;
    let original_uri = parts
        .extensions
        .get::<OriginalUri>()
        .map_or(&parts.uri, |original| &original.0);
    auth_request(
        &parts.method,
        &parts.uri,
        original_uri,
        &parts.headers,
        body.to_vec(),
        &auth.config().base_url,
    )
}

/// Distinguish the read limit from transport failures such as malformed
/// chunked framing or a client disconnect.
fn is_length_limit_error(error: &axum::Error) -> bool {
    std::iter::successors(std::error::Error::source(error), |source| source.source())
        .any(<dyn std::error::Error>::is::<http_body_util::LengthLimitError>)
}

fn render(result: Result<AuthResponse, AuthError>) -> Response {
    result
        .unwrap_or_else(AuthError::to_auth_response)
        .into_response()
}