//! Stateless CSRF protection (signed double-submit cookie).
//!
//! Authentication is cookie-based (the `auth_token` JWT cookie, see
//! [`crate::shared::guard`]), so any state-changing request the browser can be
//! tricked into making carries the victim's credentials. `SameSite=Lax` on the
//! auth cookie already blocks the cross-site *form* POST case; this layer is the
//! defense-in-depth on top of it.
//!
//! ## How it works
//! On every safe request we ensure the browser holds a `csrf_token` cookie whose
//! value is `<random>.<hmac>` — the HMAC is keyed by the app's JWT secret, so the
//! server can recognise its own tokens without storing any per-session state
//! (no session table, survives restarts and multiple instances). On every unsafe
//! request ([`protect`]) the same token must be echoed back, either as the
//! `X-CSRF-Token` header (htmx requests) or a `_csrf` form field (native
//! `<form>` submits, which cannot set a custom header). Because a cross-origin
//! attacker can neither read the cookie nor forge a valid signature, they cannot
//! produce a matching echo.
//!
//! The browser side lives in the two base templates (`base.html`,
//! `admin/base.html`): an `htmx:configRequest` hook adds the header and a
//! `submit` hook injects the hidden `_csrf` field.
//!
//! `/api/*` is exempt: that subtree is the token-authenticated JSON API and the
//! OAuth2 callback, neither of which is driven by the browser session cookie.

use axum::{
    body::{to_bytes, Body},
    extract::{Request, State},
    http::{header, Method, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
};
use axum_extra::extract::cookie::{Cookie, CookieJar, SameSite};
use bytes::Bytes;
use hmac::{Hmac, Mac};
use loco_rs::prelude::*;
use sha2::Sha256;
use subtle::ConstantTimeEq;
use time::Duration as TimeDuration;
use uuid::Uuid;

/// Cookie that holds the signed token. Deliberately *not* `HttpOnly`: the page
/// JS has to read it to echo it back in the header / hidden field.
pub const CSRF_COOKIE: &str = "csrf_token";
/// Header carrying the echoed token on htmx (and any scripted) requests.
pub const CSRF_HEADER: &str = "x-csrf-token";
/// Hidden form field carrying the echoed token on native `<form>` submits.
pub const CSRF_FIELD: &str = "_csrf";

/// Cookie lifetime. Long enough to outlast a normal browsing session; the token
/// only needs to be stable, not short-lived (it is not a credential on its own).
const COOKIE_MAX_AGE_SECS: i64 = 60 * 60 * 24 * 14;
/// Upper bound on a body we will buffer to find the `_csrf` field. Covers the
/// largest native multipart submit (a 10 MiB image upload plus the other
/// fields); anything bigger is rejected rather than buffered.
const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;

tokio::task_local! {
    /// CSRF token bound to the in-flight request. [`protect`] sets it so the
    /// `csrf_token()` Tera function can render it into pages (the `hx-headers`
    /// on `<body>` and the `ui::csrf_field()` hidden input) without every
    /// controller having to thread it through the view context.
    static REQUEST_TOKEN: String;
}

/// The CSRF token for the current request task, if one is bound. Returns `None`
/// outside a request (e.g. a mailer rendering a template). Used by the
/// `csrf_token()` Tera function registered in the view engine.
#[must_use]
pub fn current_token() -> Option<String> {
    REQUEST_TOKEN.try_with(String::clone).ok()
}

type HmacSha256 = Hmac<Sha256>;

fn to_hex(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        s.push_str(&format!("{b:02x}"));
    }
    s
}

/// HMAC-SHA256 of the random part, keyed by the app secret.
fn sign(secret: &str, random: &str) -> String {
    let mut mac =
        HmacSha256::new_from_slice(secret.as_bytes()).expect("HMAC accepts a key of any length");
    mac.update(random.as_bytes());
    to_hex(&mac.finalize().into_bytes())
}

/// Mint a fresh `<random>.<hmac>` token.
pub fn make_token(secret: &str) -> String {
    let random = Uuid::new_v4().simple().to_string();
    let sig = sign(secret, &random);
    format!("{random}.{sig}")
}

/// True when `token` is well-formed and its signature is the one this server
/// would produce — i.e. it is one of *our* tokens, not an attacker-injected one.
fn signature_valid(secret: &str, token: &str) -> bool {
    let Some((random, sig)) = token.split_once('.') else {
        return false;
    };
    let expected = sign(secret, random);
    expected.as_bytes().ct_eq(sig.as_bytes()).into()
}

/// Constant-time equality of two full tokens (the double-submit comparison).
fn tokens_match(a: &str, b: &str) -> bool {
    a.as_bytes().ct_eq(b.as_bytes()).into()
}

fn issue_cookie(token: String) -> Cookie<'static> {
    Cookie::build((CSRF_COOKIE, token))
        .path("/")
        .http_only(false)
        .same_site(SameSite::Lax)
        .max_age(TimeDuration::seconds(COOKIE_MAX_AGE_SECS))
        .build()
}

fn forbidden(reason: &str) -> Response {
    tracing::debug!(reason, "CSRF check rejected request");
    (StatusCode::FORBIDDEN, "CSRF validation failed").into_response()
}

/// Attach the given token to an outgoing response as the `csrf_token` cookie.
fn attach_cookie(res: &mut Response, token: &str) {
    let cookie = issue_cookie(token.to_string());
    if let Ok(value) = cookie.encoded().to_string().parse() {
        res.headers_mut().append(header::SET_COOKIE, value);
    }
}

/// Pull the `_csrf` value out of an `application/x-www-form-urlencoded` body.
fn field_from_urlencoded(bytes: &[u8]) -> Option<String> {
    form_urlencoded::parse(bytes)
        .find(|(k, _)| k == CSRF_FIELD)
        .map(|(_, v)| v.into_owned())
}

/// Pull the `_csrf` value out of a `multipart/form-data` body. Parses a *copy*
/// of the buffered bytes purely to read the field; the original bytes are still
/// forwarded to the handler unchanged.
async fn field_from_multipart(content_type: &str, bytes: Bytes) -> Option<String> {
    let boundary = multer::parse_boundary(content_type).ok()?;
    let stream = futures_util::stream::once(async move { Ok::<_, std::io::Error>(bytes) });
    let mut multipart = multer::Multipart::new(stream, boundary);
    while let Ok(Some(field)) = multipart.next_field().await {
        if field.name() == Some(CSRF_FIELD) {
            return field.text().await.ok();
        }
    }
    None
}

/// CSRF enforcement middleware. Safe methods get a token cookie (minted if
/// missing); unsafe methods must echo a valid, matching token. See the module
/// docs for the full scheme.
pub async fn protect(
    State(ctx): State<AppContext>,
    jar: CookieJar,
    req: Request,
    next: Next,
) -> Response {
    // The token is keyed by the JWT secret. If no secret is configured we cannot
    // validate, so fail open rather than 403 the whole site — the same secret is
    // already required for auth to work at all.
    let Ok(jwt) = ctx.config.get_jwt_config() else {
        return next.run(req).await;
    };
    let secret = jwt.secret.clone();

    // The token the browser currently holds, accepted only if untampered.
    let cookie_token = jar
        .get(CSRF_COOKIE)
        .map(|c| c.value().to_string())
        .filter(|t| signature_valid(&secret, t));

    let is_safe = matches!(
        *req.method(),
        Method::GET | Method::HEAD | Method::OPTIONS | Method::TRACE
    );
    // Token-auth JSON API + OAuth2 callback: not browser-cookie-driven.
    let exempt = req.uri().path().starts_with("/api/");

    if is_safe || exempt {
        // Bind a token for the page to render even on the very first visit, and
        // persist that same token in the cookie so the later submit matches.
        let token = cookie_token.clone().unwrap_or_else(|| make_token(&secret));
        let mut res = REQUEST_TOKEN.scope(token.clone(), next.run(req)).await;
        if cookie_token.is_none() {
            attach_cookie(&mut res, &token);
        }
        return res;
    }

    // ---- unsafe, non-exempt: require a valid double-submit ----
    let Some(expected) = cookie_token else {
        return forbidden("missing or tampered CSRF cookie");
    };

    // htmx and other scripted requests send the token as a header; prefer it so
    // we never have to touch the body.
    let header_token = req
        .headers()
        .get(CSRF_HEADER)
        .and_then(|v| v.to_str().ok())
        .map(str::to_string);
    if let Some(header_token) = header_token {
        if tokens_match(&header_token, &expected) {
            return REQUEST_TOKEN.scope(expected, next.run(req)).await;
        }
        return forbidden("CSRF header did not match cookie");
    }

    // Native <form> submit: the token is a `_csrf` body field. Buffer the body,
    // read the field from a copy, then forward the original bytes untouched.
    let content_type = req
        .headers()
        .get(header::CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("")
        .to_string();

    let (parts, body) = req.into_parts();
    let Ok(bytes) = to_bytes(body, MAX_BODY_BYTES).await else {
        return forbidden("request body too large to validate CSRF");
    };

    let submitted = if content_type.starts_with("application/x-www-form-urlencoded") {
        field_from_urlencoded(&bytes)
    } else if content_type.starts_with("multipart/form-data") {
        field_from_multipart(&content_type, bytes.clone()).await
    } else {
        None
    };

    let valid = submitted
        .as_deref()
        .is_some_and(|t| tokens_match(t, &expected));
    if !valid {
        return forbidden("missing or non-matching CSRF form field");
    }

    let req = Request::from_parts(parts, Body::from(bytes));
    REQUEST_TOKEN.scope(expected, next.run(req)).await
}

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

    const SECRET: &str = "test-secret";

    #[test]
    fn fresh_token_validates() {
        let token = make_token(SECRET);
        assert!(signature_valid(SECRET, &token));
    }

    #[test]
    fn tampered_or_foreign_token_rejected() {
        let token = make_token(SECRET);
        // Flip the random part but keep the old signature.
        let (_, sig) = token.split_once('.').unwrap();
        let forged = format!("{}.{sig}", Uuid::new_v4().simple());
        assert!(!signature_valid(SECRET, &forged));
        // A token minted under a different secret is not ours.
        assert!(!signature_valid(SECRET, &make_token("other-secret")));
        // Malformed input never validates.
        assert!(!signature_valid(SECRET, "no-dot-here"));
    }

    #[test]
    fn double_submit_compare() {
        let token = make_token(SECRET);
        assert!(tokens_match(&token, &token.clone()));
        assert!(!tokens_match(&token, &make_token(SECRET)));
    }

    #[test]
    fn extract_field_from_urlencoded() {
        let body = b"email=a%40b.com&_csrf=abc.def&password=x";
        assert_eq!(field_from_urlencoded(body), Some("abc.def".to_string()));
        assert_eq!(field_from_urlencoded(b"email=x"), None);
    }

    #[tokio::test]
    async fn extract_field_from_multipart() {
        let boundary = "X-BOUNDARY";
        let content_type = format!("multipart/form-data; boundary={boundary}");
        let body = format!(
            "--{b}\r\nContent-Disposition: form-data; name=\"name\"\r\n\r\nWidget\r\n\
             --{b}\r\nContent-Disposition: form-data; name=\"_csrf\"\r\n\r\nabc.def\r\n\
             --{b}--\r\n",
            b = boundary
        );
        let got = field_from_multipart(&content_type, Bytes::from(body)).await;
        assert_eq!(got, Some("abc.def".to_string()));
    }
}
