runique 3.0.2

//! CSRF Middleware: generates and stores the token in session, validates mutating requests.
use crate::auth::session::is_authenticated;
use crate::context::RequestExtensions;
use crate::forms::extractor::{csrf_required, is_csrf_exempt};
use crate::utils::{
    aliases::AEngine,
    constante::{session::CSRF_TOKEN_KEY, session_key::session::SESSION_USER_ID_KEY},
    csrf::{CsrfContext, CsrfToken},
};
use axum::{
    body::Body,
    extract::State,
    http::{HeaderValue, Method, Request, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
};
use subtle::ConstantTimeEq;
use tera::{Function, Kwargs, TeraResult, Value, escape_html};
use tower_sessions::Session;

/// Tera function backing the `{{ csrf_token(...) }}` template call: renders
/// the current request's CSRF token as a hidden `<input>`, HTML-escaping it
/// before emitting the unescaped tag markup.
pub struct CsrfTokenFunction;

// `tera::State` is aliased locally: `axum::extract::State` is already imported above.
impl Function<TeraResult<Value>> for CsrfTokenFunction {
    fn is_safe(&self) -> bool {
        true
    }

    fn call(&self, kwargs: Kwargs, _: &tera::State) -> TeraResult<Value> {
        let token = kwargs.get::<&str>(CSRF_TOKEN_KEY)?.unwrap_or("");

        // Escaped even though the generator only emits attribute-safe characters:
        // the tag is emitted unescaped, so it must not depend on a guarantee made
        // elsewhere. Mirrors the `csrf_field` filter.
        let mut escaped = Vec::new();
        escape_html(token, &mut escaped)
            .map_err(|e| tera::Error::chain("csrf_token: failed to escape the token", e))?;
        let token = String::from_utf8(escaped)?;

        Ok(Value::safe_string(&format!(
            r#"<input type="hidden" name="csrf_token" value="{token}">"#
        )))
    }
}

/// Issues/refreshes the session's CSRF token — always, on every request,
/// exempt or not — and enforces it on every non-safe method (`csrf_required`
/// — everything except GET/HEAD, so also OPTIONS/TRACE, fail-closed), unless
/// the path is in `csrf_exempt_paths` (`is_csrf_exempt`). Both checks are
/// shared with `Request::form()`'s own guard (`forms/extractor.rs`) so this
/// middleware and the extractor can never silently diverge. OPTIONS
/// preflight never actually reaches this far when CORS is configured — see
/// `SLOT_CORS`'s doc comment.
///
/// `csrf_exempt_paths` only skips *validation* — the token is still
/// generated and injected into request extensions for exempt paths too, so
/// `Request`/`RuniqueContext` (which unconditionally expect a `CsrfToken`
/// extension) keep working for a handler that's exempt from CSRF but still
/// wants everything else `Request` provides (session, template context, …).
///
/// An `X-CSRF-Token` header is checked with a constant-time comparison
/// (`ct_eq`) against the session token, so a wrong guess can't be
/// distinguished from a right one by response timing. Requests without that
/// header are only allowed through when they are form submissions
/// (urlencoded/multipart) — Prisme then validates the `csrf_token` form
/// field itself; a JSON request with no header is rejected outright. The
/// token exposed back to the client via the `X-CSRF-Token` response header
/// is masked per response (`session_token.masked()`) as a BREACH mitigation,
/// and a GET/HEAD request's own `csrf_token` query parameter is stripped via
/// redirect so a token never lingers in browser history or referrer headers.
pub async fn csrf_middleware(
    State(engine): State<AEngine>,
    session: Session,
    mut req: Request<Body>,
    next: Next,
) -> Response {
    // Strip csrf_token from GET query params to avoid token exposure in URLs
    if matches!(req.method(), &Method::GET | &Method::HEAD) {
        let uri = req.uri();
        if let Some(query) = uri.query()
            && query.split('&').any(|p| p.starts_with("csrf_token="))
        {
            let clean_query: String = query
                .split('&')
                .filter(|p| !p.starts_with("csrf_token="))
                .collect::<Vec<_>>()
                .join("&");

            let new_uri = if clean_query.is_empty() {
                uri.path().to_string()
            } else {
                format!("{}?{}", uri.path(), clean_query)
            };

            if let Ok(location) = HeaderValue::from_str(&new_uri) {
                let mut res = (StatusCode::FOUND, "").into_response();
                res.headers_mut()
                    .insert(axum::http::header::LOCATION, location);
                return res;
            }
        }
    }

    // Exempt paths (webhooks with their own signature verification) still get a token
    // generated and injected below — only the *validation* is skipped further down.
    // Bailing out here entirely used to also skip token generation, which meant
    // `Request`/`RuniqueContext` (which unconditionally expect a `CsrfToken` extension)
    // failed with a 500 on any exempt handler that used them for anything else
    // (session access, template context, …) — exempting a path from CSRF checks
    // shouldn't also break an unrelated feature.
    let exempt = is_csrf_exempt(req.uri().path(), &engine.csrf_exempt_paths);

    let secret = &engine.config.server.secret_key;

    let Some(session_token) = session_csrf_token(&session, secret).await else {
        return (StatusCode::INTERNAL_SERVER_ERROR, "Session write error").into_response();
    };

    if !exempt && csrf_required(req.method()) {
        let has_header = req.headers().contains_key("X-CSRF-Token");

        // If header present, we validate (AJAX request)
        if has_header {
            let header_token = req
                .headers()
                .get("X-CSRF-Token")
                .and_then(|h| h.to_str().ok())
                .and_then(|masked| CsrfToken::unmasked(masked).ok());

            match header_token {
                Some(token)
                    if token
                        .as_str()
                        .as_bytes()
                        // ct_eq: constant-time comparison — prevents an attacker
                        // from guessing the token byte by byte via response time
                        .ct_eq(session_token.as_str().as_bytes())
                        .into() =>
                {
                    // OK, continue
                }
                _ => {
                    return (StatusCode::FORBIDDEN, "Invalid CSRF token").into_response();
                }
            }
        } else {
            // No CSRF header: allowed only for HTML form submissions
            // (urlencoded / multipart). JSON requests without header are blocked.
            let ct = req
                .headers()
                .get("content-type")
                .and_then(|v| v.to_str().ok())
                .unwrap_or("");
            let is_form = ct.starts_with("application/x-www-form-urlencoded")
                || ct.starts_with("multipart/form-data");
            if !is_form {
                return (StatusCode::FORBIDDEN, "CSRF token required").into_response();
            }
            // Otherwise, we let Prisme validate the form field
        }
    }

    // Inject the raw session token for the request pipeline (form() validation).
    let extensions = RequestExtensions::new().with_csrf_token(session_token.clone());
    extensions.inject_request(&mut req);

    let mut res = next.run(req).await;

    // Read the token again: the handler may have rotated it (`login()` and
    // `logout()` do), and handing the old one back would leave the
    // frontend with a token the session no longer accepts. If the session
    // can't be written, no header at all rather than a stale token.
    let exposed = session_csrf_token(&session, secret).await;

    // Expose the token to the frontend, masked per-response (BREACH mitigation).
    // If masking ever fails (would only happen on a non-hex token, which the
    // generator never produces), skip the header rather than leaking the raw
    // unmasked session token in the response.
    if let Some(token) = exposed
        && let Ok(masked) = token.masked()
        && let Ok(hv) = HeaderValue::from_str(masked.as_str())
    {
        res.headers_mut().insert("X-CSRF-Token", hv);
    }

    res
}

/// Drops the session's CSRF token so the next one is a new one. Call it
/// whenever the privilege level changes: `login()` and `logout()` do. The CSRF
/// middleware issues the replacement before the response leaves, so a custom
/// auth flow only has to call this.
pub async fn rotate_csrf_token(session: &Session) -> Result<(), tower_sessions::session::Error> {
    session.remove_value(CSRF_TOKEN_KEY).await.map(|_| ())
}

/// The session's CSRF token, created and stored first when the session has
/// none: a new visitor, or a session whose token was just rotated. A new
/// token is bound to the user when someone is logged in, to the session
/// otherwise. `None` when the session can't be written.
async fn session_csrf_token(session: &Session, secret: &str) -> Option<CsrfToken> {
    if let Some(token) = session
        .get::<CsrfToken>(CSRF_TOKEN_KEY)
        .await
        .ok()
        .flatten()
    {
        session.insert(CSRF_TOKEN_KEY, &token).await.ok()?;
        return Some(token);
    }
    let token = if is_authenticated(session).await {
        let user_id: crate::utils::pk::Pk = session
            .get::<crate::utils::pk::Pk>(SESSION_USER_ID_KEY)
            .await
            .ok()
            .flatten()
            .unwrap_or_default();
        CsrfToken::generate_with_context(&CsrfContext::Authenticated { user_id }, secret)
    } else {
        let session_id = session.id().map(|id| id.to_string()).unwrap_or_default();
        CsrfToken::generate_with_context(
            &CsrfContext::Anonymous {
                session_id: &session_id,
            },
            secret,
        )
    };
    session.insert(CSRF_TOKEN_KEY, &token).await.ok()?;
    Some(token)
}