eric / rustwarden

//! CSRF Protection Middleware
//!
//! Protects against Cross-Site Request Forgery (CSRF) attacks via
//! session-bound tokens for state-changing operations.
//!
//! Two entry points:
//! - `csrf_protection` — default behavior. Accepts the token only in the
//!   `X-CSRF-Token` header. Suited to JS clients / HTMX.
//! - `csrf_protection_with_config` — explicit `CsrfProtectionConfig`.
//!   When `accept_form_body = true`, also looks for `csrf_token=...` in a
//!   `application/x-www-form-urlencoded` body and buffers it up to
//!   `max_form_body_bytes`. Suited to SSR apps without JS (e.g. Gitrust).
//!
//! Security properties (inherited):
//! - SEC-H4: constant-time token comparison.
//! - SEC-C3: blocks state-changing requests without a token, except those
//!   carrying an `Authorization: Bearer ...` header (API clients, not
//!   browser-driven, so not CSRF-vulnerable).

use axum::{
    body::Body,
    extract::{Request, State},
    http::{Method, StatusCode, Uri, header::CONTENT_TYPE},
    middleware::Next,
    response::{IntoResponse, Response},
};
use rand::RngCore;
use sea_orm::DatabaseConnection;
use std::sync::Arc;
use subtle::ConstantTimeEq;
use tower_sessions::Session;

use crate::error::AppError;

/// Default buffer cap for form bodies that may carry a CSRF token: 1 MiB.
pub const DEFAULT_MAX_FORM_BODY_BYTES: usize = 1024 * 1024;

/// Runtime configuration for [`csrf_protection_with_config`].
///
/// The defaults match the behavior of [`csrf_protection`] (header-only, no
/// body parsing), so an app that opts in to the configurable wrapper without
/// changing any field sees no behavior change.
#[derive(Debug, Clone)]
pub struct CsrfProtectionConfig {
    /// When true, if the `X-CSRF-Token` header is absent and the request
    /// declares `Content-Type: application/x-www-form-urlencoded`, the body
    /// is buffered and scanned for a `csrf_token` field.
    ///
    /// Multipart bodies are never buffered (too risky on large uploads).
    pub accept_form_body: bool,

    /// Maximum number of bytes buffered when `accept_form_body` is true.
    /// Requests that exceed this cap are rejected with HTTP 413.
    pub max_form_body_bytes: usize,
}

impl Default for CsrfProtectionConfig {
    fn default() -> Self {
        Self {
            accept_form_body: false,
            max_form_body_bytes: DEFAULT_MAX_FORM_BODY_BYTES,
        }
    }
}

/// Session key for CSRF token
const CSRF_TOKEN_KEY: &str = "csrf_token";

/// Generate a random CSRF token
fn generate_csrf_token() -> String {
    let mut rng = rand::rng();
    let mut bytes = [0u8; 32];
    rng.fill_bytes(&mut bytes);
    hex::encode(bytes)
}

/// Get or create CSRF token from session
///
/// # Arguments
///
/// * `session` - Tower session
///
/// # Returns
///
/// CSRF token string
pub async fn get_or_create_csrf_token(session: &Session) -> Result<String, AppError> {
    // Try to get existing token from session
    match session.get::<String>(CSRF_TOKEN_KEY).await {
        Ok(Some(token)) => return Ok(token),
        Ok(None) => {
            // Token doesn't exist, create one
        }
        Err(e) => {
            tracing::warn!("Failed to get CSRF token from session: {}", e);
            // Continue to create new token
        }
    }

    // Generate new token if not exists
    let token = generate_csrf_token();
    session.insert(CSRF_TOKEN_KEY, &token).await.map_err(|e| {
        tracing::error!("Failed to store CSRF token in session: {}", e);
        AppError::Internal(format!("Failed to store CSRF token: {}", e))
    })?;
    Ok(token)
}

/// Extract the CSRF token from the `X-CSRF-Token` header only.
///
/// Body-based extraction is handled separately by [`extract_csrf_from_form`]
/// when the middleware is configured with `accept_form_body = true`.
fn extract_header_token(request: &Request) -> Option<String> {
    request
        .headers()
        .get("X-CSRF-Token")
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string())
}

/// True iff the request declares `Content-Type: application/x-www-form-urlencoded`.
fn is_form_urlencoded(request: &Request) -> bool {
    request
        .headers()
        .get(CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .map(|v| {
            v.to_ascii_lowercase()
                .starts_with("application/x-www-form-urlencoded")
        })
        .unwrap_or(false)
}

/// Extract the `csrf_token` field from a URL-encoded form body.
///
/// Returns `None` when:
/// - the body is not valid UTF-8,
/// - no `csrf_token` field is present,
/// - the URL-decoded value is invalid.
///
/// Exposed as `pub` so apps that run their own buffering can still reuse the
/// framework's parser (matches proposition B in
/// [`docs/UPSTREAM_RUSTWARDEN.md`](../../../../docs/UPSTREAM_RUSTWARDEN.md)).
pub fn extract_csrf_from_form(body: &[u8]) -> Option<String> {
    let text = std::str::from_utf8(body).ok()?;
    text.split('&')
        .filter_map(|pair| pair.split_once('='))
        .find(|(k, _)| *k == "csrf_token")
        .and_then(|(_, v)| urlencoding::decode(v).ok())
        .map(|cow| cow.into_owned())
}

/// Middleware that validates CSRF tokens for state-changing operations
///
/// This middleware:
/// - Allows GET/HEAD/OPTIONS requests without CSRF token
/// - Requires valid CSRF token for POST/PUT/DELETE/PATCH requests
/// - Validates token against session-stored token
///
/// # Arguments
///
/// * `State(db)` - Database connection
/// * `session` - Tower session
/// * `uri` - Request URI
/// * `request` - HTTP request
/// * `next` - Next middleware or handler
///
/// # Returns
///
/// * `200 OK` if valid or safe method (passes to handler)
/// * `403 Forbidden` if CSRF token missing or invalid
///
/// SEC-C3: Check if the request uses Bearer authentication (API client)
/// API clients using Authorization: Bearer are not vulnerable to CSRF
/// because the token is not sent automatically by the browser.
fn is_bearer_api_request(request: &Request) -> bool {
    request
        .headers()
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .map(|v| v.to_lowercase().starts_with("bearer "))
        .unwrap_or(false)
}

/// SEC-H4: Constant-time comparison of CSRF tokens
fn csrf_tokens_match(a: &str, b: &str) -> bool {
    a.as_bytes().ct_eq(b.as_bytes()).into()
}

/// Outcome of running the CSRF check against an incoming request.
///
/// `Pass` carries the (possibly reconstructed) request forward so body-buffered
/// flows can replay the consumed bytes to downstream extractors.
#[derive(Debug)]
enum CsrfOutcome {
    /// Request is safe, Bearer-authed, or token validated — forward to `next`.
    Pass(Request),
    /// Token missing, mismatched, or body cap exceeded — short-circuit.
    Reject(Response),
}

/// Core CSRF check with configurable body buffering.
///
/// Flow:
/// 1. Safe method (`GET`/`HEAD`/`OPTIONS`) → pass.
/// 2. `Authorization: Bearer ...` → pass (API client, not CSRF-vulnerable).
/// 3. `X-CSRF-Token` header present → validate and decide.
/// 4. Else if `cfg.accept_form_body` and `Content-Type: application/x-www-form-urlencoded` →
///    buffer (capped at `cfg.max_form_body_bytes`), scan for `csrf_token`, validate.
///    The buffered body is reattached to the request so the downstream handler
///    can still decode its `Form<T>`.
/// 5. Else → reject (403 CSRF token required).
async fn csrf_check(
    session: &Session,
    uri: &Uri,
    request: Request,
    cfg: &CsrfProtectionConfig,
) -> Result<CsrfOutcome, AppError> {
    let method = request.method().clone();

    // SEC-H4: TRACE removed — only safe idempotent methods are exempt.
    if matches!(method, Method::GET | Method::HEAD | Method::OPTIONS) {
        return Ok(CsrfOutcome::Pass(request));
    }

    // SEC-C3: Skip CSRF for Bearer API clients.
    if is_bearer_api_request(&request) {
        tracing::debug!("CSRF skip: Bearer API request for {} {}", method, uri);
        return Ok(CsrfOutcome::Pass(request));
    }

    tracing::debug!("CSRF protection: validating token for {} {}", method, uri);
    let session_token = get_or_create_csrf_token(session).await?;

    // Prefer the header: cheap and non-destructive.
    if let Some(token) = extract_header_token(&request) {
        return Ok(validate_or_reject(
            &session_token,
            &token,
            &method,
            uri,
            request,
        ));
    }

    // Fall back to the form body when explicitly enabled.
    if cfg.accept_form_body && is_form_urlencoded(&request) {
        let (parts, body) = request.into_parts();
        let bytes = match axum::body::to_bytes(body, cfg.max_form_body_bytes).await {
            Ok(b) => b,
            Err(err) => {
                tracing::warn!(
                    cap = cfg.max_form_body_bytes,
                    error = %err,
                    "CSRF body buffering failed (cap exceeded or stream error) for {} {}",
                    method,
                    uri,
                );
                return Ok(CsrfOutcome::Reject(
                    (
                        StatusCode::PAYLOAD_TOO_LARGE,
                        "Form body too large for CSRF check",
                    )
                        .into_response(),
                ));
            }
        };
        let body_token = extract_csrf_from_form(&bytes);
        // Reconstruct the request *before* deciding, so that a valid token
        // still has a body available for the downstream handler.
        let reconstructed = Request::from_parts(parts, Body::from(bytes));
        return Ok(match body_token {
            Some(t) => validate_or_reject(&session_token, &t, &method, uri, reconstructed),
            None => {
                tracing::warn!("CSRF token missing from form body for {} {}", method, uri);
                CsrfOutcome::Reject((StatusCode::FORBIDDEN, "CSRF token required").into_response())
            }
        });
    }

    // SEC-C3: No token at all → block.
    tracing::warn!("CSRF token missing for {} {}", method, uri);
    Ok(CsrfOutcome::Reject(
        (StatusCode::FORBIDDEN, "CSRF token required").into_response(),
    ))
}

fn validate_or_reject(
    session_token: &str,
    incoming: &str,
    method: &Method,
    uri: &Uri,
    request: Request,
) -> CsrfOutcome {
    if csrf_tokens_match(incoming, session_token) {
        tracing::debug!("CSRF token validated for {} {}", method, uri);
        CsrfOutcome::Pass(request)
    } else {
        tracing::warn!(
            "CSRF token mismatch for {} {} (expected: {}..., got: {}...)",
            method,
            uri,
            &session_token[..8.min(session_token.len())],
            &incoming[..8.min(incoming.len())]
        );
        CsrfOutcome::Reject((StatusCode::FORBIDDEN, "Invalid CSRF token").into_response())
    }
}

/// Default header-only CSRF middleware.
///
/// Wire as `axum::middleware::from_fn_with_state(db, csrf_protection)` on a
/// router whose routes already reject without a token. For SSR apps that only
/// have the token in a form body, use [`csrf_protection_with_config`] instead.
pub async fn csrf_protection(
    State(_db): State<DatabaseConnection>,
    session: Session,
    uri: Uri,
    request: Request,
    next: Next,
) -> Result<Response, AppError> {
    let cfg = CsrfProtectionConfig::default();
    match csrf_check(&session, &uri, request, &cfg).await? {
        CsrfOutcome::Pass(req) => Ok(next.run(req).await),
        CsrfOutcome::Reject(resp) => Ok(resp),
    }
}

/// CSRF middleware with explicit [`CsrfProtectionConfig`].
///
/// Wire with `axum::middleware::from_fn_with_state((db, Arc::new(cfg)),
/// csrf_protection_with_config)` when you need `accept_form_body = true` for
/// SSR HTML forms (no JS, token in a hidden input).
pub async fn csrf_protection_with_config(
    State((_db, cfg)): State<(DatabaseConnection, Arc<CsrfProtectionConfig>)>,
    session: Session,
    uri: Uri,
    request: Request,
    next: Next,
) -> Result<Response, AppError> {
    match csrf_check(&session, &uri, request, cfg.as_ref()).await? {
        CsrfOutcome::Pass(req) => Ok(next.run(req).await),
        CsrfOutcome::Reject(resp) => Ok(resp),
    }
}

/// Internal implementation that takes db by reference instead of via `State`.
///
/// Behaviorally identical to [`csrf_protection`]; kept as a separate entry
/// point for call sites that already have a `&DatabaseConnection` in scope.
pub async fn csrf_protection_impl(
    _db: &DatabaseConnection,
    session: Session,
    uri: Uri,
    request: Request,
    next: Next,
) -> Result<Response, AppError> {
    let cfg = CsrfProtectionConfig::default();
    match csrf_check(&session, &uri, request, &cfg).await? {
        CsrfOutcome::Pass(req) => Ok(next.run(req).await),
        CsrfOutcome::Reject(resp) => Ok(resp),
    }
}

/// Validate CSRF token from form data
///
/// # Arguments
///
/// * `session` - Tower session
/// * `form_token` - CSRF token from form
///
/// # Returns
///
/// Ok(()) if valid, Err(AppError) if invalid
pub async fn validate_csrf_token(session: &Session, form_token: &str) -> Result<(), AppError> {
    let session_token = get_or_create_csrf_token(session).await?;

    // SEC-H4: Constant-time comparison to prevent timing attacks
    if !csrf_tokens_match(form_token, &session_token) {
        tracing::warn!(
            "CSRF token mismatch (expected: {}..., got: {}...)",
            &session_token[..8.min(session_token.len())],
            &form_token[..8.min(form_token.len())]
        );
        return Err(AppError::Forbidden);
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request as HttpRequest;
    use std::sync::Arc;
    use tower_sessions::MemoryStore;

    // ---------- Form body extraction (pure function) ----------

    #[test]
    fn extract_csrf_from_form_returns_token_value() {
        let body = b"csrf_token=deadbeef";
        assert_eq!(extract_csrf_from_form(body), Some("deadbeef".to_string()));
    }

    #[test]
    fn extract_csrf_from_form_handles_multiple_fields() {
        let body = b"foo=bar&csrf_token=abc123&baz=qux";
        assert_eq!(extract_csrf_from_form(body), Some("abc123".to_string()));
    }

    #[test]
    fn extract_csrf_from_form_url_decodes_token() {
        // 61 = '=' so %3D is a decoded '='
        let body = b"csrf_token=abc%3D%3D";
        assert_eq!(extract_csrf_from_form(body), Some("abc==".to_string()));
    }

    #[test]
    fn extract_csrf_from_form_returns_none_when_absent() {
        let body = b"foo=bar&baz=qux";
        assert_eq!(extract_csrf_from_form(body), None);
    }

    #[test]
    fn extract_csrf_from_form_returns_none_on_invalid_utf8() {
        let body = &[0xff, 0xfe, 0xfd];
        assert_eq!(extract_csrf_from_form(body), None);
    }

    #[test]
    fn test_generate_csrf_token() {
        let token1 = generate_csrf_token();
        let token2 = generate_csrf_token();

        assert_eq!(token1.len(), 64); // 32 bytes = 64 hex chars
        assert_ne!(token1, token2); // Tokens should be different
    }

    // ---------- T12: constant-time token comparison ----------

    #[test]
    fn csrf_tokens_match_is_true_for_equal_tokens() {
        let token = generate_csrf_token();
        assert!(csrf_tokens_match(&token, &token));
    }

    #[test]
    fn csrf_tokens_match_is_false_for_different_tokens() {
        let a = generate_csrf_token();
        let b = generate_csrf_token();
        assert!(!csrf_tokens_match(&a, &b));
    }

    #[test]
    fn csrf_tokens_match_is_false_for_length_mismatch() {
        // `ct_eq` over byte slices of differing lengths must not match.
        assert!(!csrf_tokens_match("abc", "abcd"));
    }

    fn test_session() -> Session {
        Session::new(None, Arc::new(MemoryStore::default()), None)
    }

    fn build_request(method: Method, token_header: Option<&str>, bearer: bool) -> Request {
        let mut builder = HttpRequest::builder().method(method).uri("/resource");
        if let Some(t) = token_header {
            builder = builder.header("X-CSRF-Token", t);
        }
        if bearer {
            builder = builder.header(axum::http::header::AUTHORIZATION, "Bearer some.jwt.token");
        }
        builder.body(Body::empty()).unwrap()
    }

    fn build_form_request(body: &str, content_type: Option<&str>) -> Request {
        let ct = content_type.unwrap_or("application/x-www-form-urlencoded");
        HttpRequest::builder()
            .method(Method::POST)
            .uri("/resource")
            .header(CONTENT_TYPE, ct)
            .body(Body::from(body.to_owned()))
            .unwrap()
    }

    fn uri() -> Uri {
        "/resource".parse().unwrap()
    }

    fn assert_rejected(outcome: CsrfOutcome, expected: StatusCode) {
        match outcome {
            CsrfOutcome::Reject(resp) => assert_eq!(resp.status(), expected),
            CsrfOutcome::Pass(_) => panic!("expected Reject({expected}), got Pass"),
        }
    }

    fn default_cfg() -> CsrfProtectionConfig {
        CsrfProtectionConfig::default()
    }

    fn form_body_cfg() -> CsrfProtectionConfig {
        CsrfProtectionConfig {
            accept_form_body: true,
            max_form_body_bytes: 4096,
        }
    }

    #[tokio::test]
    async fn check_passes_safe_method() {
        let session = test_session();
        let req = build_request(Method::GET, None, false);
        let outcome = csrf_check(&session, &uri(), req, &default_cfg())
            .await
            .unwrap();
        assert!(matches!(outcome, CsrfOutcome::Pass(_)));
    }

    #[tokio::test]
    async fn check_passes_bearer_auth() {
        let session = test_session();
        let req = build_request(Method::POST, None, true);
        let outcome = csrf_check(&session, &uri(), req, &default_cfg())
            .await
            .unwrap();
        assert!(matches!(outcome, CsrfOutcome::Pass(_)));
    }

    #[tokio::test]
    async fn check_accepts_matching_header_token() {
        let session = test_session();
        let token = get_or_create_csrf_token(&session).await.unwrap();
        let req = build_request(Method::POST, Some(&token), false);
        let outcome = csrf_check(&session, &uri(), req, &default_cfg())
            .await
            .unwrap();
        assert!(matches!(outcome, CsrfOutcome::Pass(_)));
    }

    #[tokio::test]
    async fn check_rejects_mismatched_header_token() {
        let session = test_session();
        let _ = get_or_create_csrf_token(&session).await.unwrap();
        let req = build_request(Method::POST, Some("bogus"), false);
        let outcome = csrf_check(&session, &uri(), req, &default_cfg())
            .await
            .unwrap();
        assert_rejected(outcome, StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn check_rejects_missing_token() {
        let session = test_session();
        let req = build_request(Method::POST, None, false);
        let outcome = csrf_check(&session, &uri(), req, &default_cfg())
            .await
            .unwrap();
        assert_rejected(outcome, StatusCode::FORBIDDEN);
    }

    // ---------- Form-body (accept_form_body = true) ----------

    #[tokio::test]
    async fn check_accepts_matching_form_body_token_and_preserves_body() {
        let session = test_session();
        let token = get_or_create_csrf_token(&session).await.unwrap();
        let form = format!("csrf_token={token}&other=value");
        let req = build_form_request(&form, None);

        let outcome = csrf_check(&session, &uri(), req, &form_body_cfg())
            .await
            .unwrap();

        let forwarded = match outcome {
            CsrfOutcome::Pass(r) => r,
            CsrfOutcome::Reject(_) => panic!("expected Pass for valid form token"),
        };
        // Body must survive the buffering so downstream handlers can Form-decode it.
        let (_, body) = forwarded.into_parts();
        let bytes = axum::body::to_bytes(body, 4096).await.unwrap();
        assert_eq!(std::str::from_utf8(&bytes).unwrap(), form);
    }

    #[tokio::test]
    async fn check_rejects_form_body_when_accept_form_body_is_false() {
        let session = test_session();
        let token = get_or_create_csrf_token(&session).await.unwrap();
        let form = format!("csrf_token={token}");
        let req = build_form_request(&form, None);

        // Default config does NOT buffer the body → token unreachable → 403.
        let outcome = csrf_check(&session, &uri(), req, &default_cfg())
            .await
            .unwrap();
        assert_rejected(outcome, StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn check_rejects_form_body_with_mismatched_token() {
        let session = test_session();
        let _ = get_or_create_csrf_token(&session).await.unwrap();
        let req = build_form_request("csrf_token=wrong-token", None);

        let outcome = csrf_check(&session, &uri(), req, &form_body_cfg())
            .await
            .unwrap();
        assert_rejected(outcome, StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn check_rejects_form_body_exceeding_cap() {
        let session = test_session();
        let huge_body = "a=".to_string() + &"x".repeat(1024);
        let cfg = CsrfProtectionConfig {
            accept_form_body: true,
            max_form_body_bytes: 16, // intentionally tiny
        };
        let req = build_form_request(&huge_body, None);

        let outcome = csrf_check(&session, &uri(), req, &cfg).await.unwrap();
        assert_rejected(outcome, StatusCode::PAYLOAD_TOO_LARGE);
    }

    #[tokio::test]
    async fn check_rejects_multipart_body_without_header() {
        let session = test_session();
        let _ = get_or_create_csrf_token(&session).await.unwrap();
        // Multipart is never buffered — security choice, avoids large uploads.
        let req = build_form_request("ignored", Some("multipart/form-data; boundary=----abc"));

        let outcome = csrf_check(&session, &uri(), req, &form_body_cfg())
            .await
            .unwrap();
        assert_rejected(outcome, StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn default_config_matches_header_only_behavior() {
        let cfg = CsrfProtectionConfig::default();
        assert!(!cfg.accept_form_body);
        assert_eq!(cfg.max_form_body_bytes, DEFAULT_MAX_FORM_BODY_BYTES);
    }
}