//! CSRF (Cross-Site Request Forgery) protection.
//!
//! Synchronizer token pattern: generate random token on session start,
//! include in meta tag, validate on state-changing requests via X-CSRF-Token header.

use axum::{
    extract::Request,
    http::StatusCode,
    middleware::Next,
    response::{IntoResponse, Response},
};
use rand::Rng;
use tower_sessions::Session;

const CSRF_SESSION_KEY: &str = "csrf_token";
const CSRF_TOKEN_LENGTH: usize = 32;

/// Generate a new random CSRF token (64-char hex string).
pub fn generate_token() -> String {
    let mut token = [0u8; CSRF_TOKEN_LENGTH];
    rand::rng().fill_bytes(&mut token);
    hex::encode(token)
}

/// Get the existing CSRF token from the session, or create and store a new one.
///
/// Returns a 500 response if the new token could not be written to the session
/// store. Rendering a page around a token the session does not hold produces a
/// form whose next POST fails validation for no reason the user can see, which
/// is worse than an error: the middleware compares against the stored token, so
/// an unstored token is a guaranteed silent failure later. Failing here keeps
/// the failure where it happened.
pub async fn get_or_create_token(session: &Session) -> Result<String, Response> {
    if let Ok(Some(token)) = session.get::<String>(CSRF_SESSION_KEY).await {
        return Ok(token);
    }
    let token = generate_token();
    match session.insert(CSRF_SESSION_KEY, &token).await {
        Ok(()) => Ok(token),
        Err(e) => {
            tracing::error!(error = ?e, "failed to store CSRF token in session");
            Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                "Session store unavailable",
            )
                .into_response())
        }
    }
}

/// Constant-time comparison to prevent timing attacks.
pub fn constant_time_compare(a: &str, b: &str) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.bytes()
        .zip(b.bytes())
        .fold(0u8, |acc, (x, y)| acc | (x ^ y))
        == 0
}

/// Extract a single field value from an `application/x-www-form-urlencoded`
/// body, decoding `+` and `%XX` escapes. Returns the first match. Kept
/// dependency-free; only used on the no-JS CSRF fallback path.
fn extract_form_field(bytes: &[u8], key: &str) -> Option<String> {
    for pair in bytes.split(|&b| b == b'&') {
        let mut it = pair.splitn(2, |&b| b == b'=');
        let k = it.next().unwrap_or(b"");
        if k == key.as_bytes() {
            return Some(percent_decode_form(it.next().unwrap_or(b"")));
        }
    }
    None
}

/// Decode a urlencoded form value: `+` → space, `%XX` → byte, lossy UTF-8.
fn percent_decode_form(input: &[u8]) -> String {
    let mut out = Vec::with_capacity(input.len());
    let mut i = 0;
    while i < input.len() {
        match input[i] {
            b'+' => out.push(b' '),
            b'%' if i + 2 < input.len() => {
                let hi = (input[i + 1] as char).to_digit(16);
                let lo = (input[i + 2] as char).to_digit(16);
                if let (Some(hi), Some(lo)) = (hi, lo) {
                    out.push((hi * 16 + lo) as u8);
                    i += 3;
                    continue;
                }
                out.push(b'%');
            }
            b => out.push(b),
        }
        i += 1;
    }
    String::from_utf8_lossy(&out).into_owned()
}

/// Max body size we'll buffer to recover a CSRF token from a form field on the
/// no-JS fallback path. Forms are tiny; anything larger isn't a urlencoded form
/// we'd be parsing a token out of.
const MAX_FORM_FALLBACK_BYTES: usize = 64 * 1024;

/// Middleware: validate the CSRF token on POST/PUT/PATCH/DELETE.
///
/// Two delivery paths, in order:
///  1. `X-CSRF-Token` header (set by mt.js for every fetch/HTMX request). This
///     is the fast path, the request body is never touched.
///  2. A hidden `csrf_token` form field, for graceful degradation when mt.js
///     didn't run (JS disabled, asset failure). Only urlencoded bodies are
///     inspected, and only when the header is absent; multipart uploads remain
///     header-only. The body is buffered, validated, and re-attached so the
///     downstream handler still sees it.
///
/// `/auth/` is deliberately NOT exempt: its only mutating routes (`logout`,
/// `refresh`) are same-origin forms that carry the token, so they get CSRF
/// protection like everything else. `login`/`callback` are GET and so never
/// reach this check. `/_test/` is only ever mounted by the integration harness
/// (never in production); `/api/health` is GET-only.
pub async fn csrf_middleware(request: Request, next: Next) -> Response {
    let method = request.method().clone();

    if !["POST", "PUT", "PATCH", "DELETE"].contains(&method.as_str()) {
        return next.run(request).await;
    }

    let path = request.uri().path().to_string();

    let exempt_prefixes = ["/api/health", "/_test/"];
    if exempt_prefixes.iter().any(|p| path.starts_with(p)) {
        return next.run(request).await;
    }

    let session = match request.extensions().get::<Session>() {
        Some(s) => s.clone(),
        None => {
            tracing::warn!("CSRF check failed: no session");
            return (StatusCode::FORBIDDEN, "CSRF validation failed").into_response();
        }
    };

    let session_token: Option<String> = session.get(CSRF_SESSION_KEY).await.ok().flatten();

    // Fast path: token in the header (mt.js). Body untouched.
    if let Some(header_token) = request
        .headers()
        .get("X-CSRF-Token")
        .and_then(|v| v.to_str().ok())
        .map(std::string::ToString::to_string)
    {
        return match session_token {
            Some(ref expected) if constant_time_compare(expected, &header_token) => {
                next.run(request).await
            }
            _ => {
                tracing::warn!(path = %path, "CSRF token mismatch");
                (StatusCode::FORBIDDEN, "Invalid CSRF token").into_response()
            }
        };
    }

    // Fallback path: no header. Accept a hidden `csrf_token` form field from a
    // urlencoded body so the site degrades gracefully without JS.
    let is_form = request
        .headers()
        .get(axum::http::header::CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .is_some_and(|ct| ct.starts_with("application/x-www-form-urlencoded"));

    if !is_form {
        tracing::warn!(path = %path, "CSRF token missing");
        return (StatusCode::FORBIDDEN, "CSRF token required").into_response();
    }

    let (parts, body) = request.into_parts();
    let Ok(bytes) = axum::body::to_bytes(body, MAX_FORM_FALLBACK_BYTES).await else {
        tracing::warn!(path = %path, "CSRF fallback: body too large or unreadable");
        return (StatusCode::FORBIDDEN, "CSRF token required").into_response();
    };

    let form_token = extract_form_field(&bytes, "csrf_token");

    let valid = matches!(
        (&session_token, &form_token),
        (Some(expected), Some(provided)) if constant_time_compare(expected, provided)
    );

    if !valid {
        tracing::warn!(path = %path, "CSRF token mismatch (form fallback)");
        return (StatusCode::FORBIDDEN, "Invalid CSRF token").into_response();
    }

    next.run(Request::from_parts(parts, axum::body::Body::from(bytes)))
        .await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn token_length_and_hex() {
        let token = generate_token();
        assert_eq!(token.len(), 64);
        assert!(token.chars().all(|c| c.is_ascii_hexdigit()));
    }

    #[test]
    fn tokens_are_unique() {
        let a = generate_token();
        let b = generate_token();
        assert_ne!(a, b);
    }

    #[test]
    fn constant_time_compare_works() {
        assert!(constant_time_compare("abc", "abc"));
        assert!(!constant_time_compare("abc", "abd"));
        assert!(!constant_time_compare("abc", "abcd"));
        assert!(!constant_time_compare("", "a"));
        assert!(constant_time_compare("", ""));
    }

    #[test]
    fn extract_form_field_finds_token() {
        let body = b"username=alice&csrf_token=deadbeef&duration=1h";
        assert_eq!(
            extract_form_field(body, "csrf_token").as_deref(),
            Some("deadbeef")
        );
        assert_eq!(
            extract_form_field(body, "username").as_deref(),
            Some("alice")
        );
        assert_eq!(extract_form_field(body, "missing"), None);
    }

    #[test]
    fn extract_form_field_decodes_escapes() {
        let body = b"reason=a+b%2Fc&csrf_token=abc123";
        assert_eq!(extract_form_field(body, "reason").as_deref(), Some("a b/c"));
        assert_eq!(
            extract_form_field(body, "csrf_token").as_deref(),
            Some("abc123")
        );
    }

    #[test]
    fn extract_form_field_handles_empty_and_valueless() {
        assert_eq!(extract_form_field(b"", "csrf_token"), None);
        assert_eq!(
            extract_form_field(b"csrf_token=", "csrf_token").as_deref(),
            Some("")
        );
        assert_eq!(
            extract_form_field(b"csrf_token", "csrf_token").as_deref(),
            Some("")
        );
    }
}