use axum::body::Body;
use axum::http::{HeaderName, Method, Request, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
pub const CSRF_COOKIE_NAME: &str = "csrf_token";
pub const CSRF_HEADER_NAME: &str = "x-csrf-token";
pub const DEFAULT_PUBLIC_PATHS: &[&str] = &[
"/health",
"/metrics",
"/api/v1/auth/login",
"/api/v1/auth/refresh",
];
pub fn is_safe_method(method: &Method) -> bool {
matches!(*method, Method::GET | Method::HEAD | Method::OPTIONS)
}
pub fn is_public_path(path: &str, public_paths: &[&str]) -> bool {
public_paths.contains(&path)
}
pub fn generate_token() -> String {
use rand::RngCore;
let mut bytes = [0u8; 32];
rand::rngs::OsRng.fill_bytes(&mut bytes);
base64_encode(&bytes)
}
fn base64_encode(bytes: &[u8]) -> String {
use base64::Engine;
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
}
#[tracing::instrument(skip(req, next))]
pub async fn csrf_middleware(req: Request<Body>, next: Next) -> Response {
let method = req.method().clone();
let path = req.uri().path().to_string();
if is_safe_method(&method) {
return next.run(req).await;
}
if is_public_path(&path, DEFAULT_PUBLIC_PATHS) {
return next.run(req).await;
}
let cookie_token = extract_cookie_value(req.headers().get("cookie"), CSRF_COOKIE_NAME);
let header_token = req
.headers()
.get(HeaderName::from_static(CSRF_HEADER_NAME))
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
match (cookie_token, header_token) {
(Some(cookie), Some(header)) if constant_time_eq(cookie.as_bytes(), header.as_bytes()) => {
next.run(req).await
}
_ => {
tracing::warn!(
method = %method,
path = %path,
"CSRF 校验失败:Cookie 或 Header token 缺失/不匹配"
);
(StatusCode::FORBIDDEN, "CSRF token 校验失败").into_response()
}
}
}
pub fn extract_cookie_value(
cookie_header: Option<&axum::http::HeaderValue>,
name: &str,
) -> Option<String> {
let header = cookie_header?;
let header_str = header.to_str().ok()?;
for pair in header_str.split(';') {
let pair = pair.trim();
if let Some((k, v)) = pair.split_once('=') {
if k.trim() == name {
return Some(v.trim().to_string());
}
}
}
None
}
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
if a.len() != b.len() {
return false;
}
let mut result = 0u8;
for (x, y) in a.iter().zip(b.iter()) {
result |= x ^ y;
}
result == 0
}
#[cfg(test)]
mod tests {
use super::*;
use axum::body::Body;
use axum::http::{HeaderValue, Method, Request};
use axum::routing::post;
use axum::Router;
use tower::ServiceExt;
fn make_router() -> Router {
Router::new()
.route("/api/data", post(|| async { "ok" }))
.layer(axum::middleware::from_fn(csrf_middleware))
}
async fn send(
method: &str,
path: &str,
cookie: Option<&str>,
header: Option<&str>,
) -> Response {
let mut builder = Request::builder().method(method).uri(path);
if let Some(c) = cookie {
builder = builder.header("cookie", c);
}
if let Some(h) = header {
builder = builder.header(CSRF_HEADER_NAME, h);
}
let req = builder.body(Body::empty()).unwrap();
make_router().oneshot(req).await.unwrap()
}
#[test]
fn test_generate_token_length() {
let token = generate_token();
assert!(token.len() >= 40, "token too short: {}", token.len());
}
#[test]
fn test_generate_token_uniqueness() {
let t1 = generate_token();
let t2 = generate_token();
assert_ne!(t1, t2, "tokens must be unique");
}
#[test]
fn test_is_safe_method() {
assert!(is_safe_method(&Method::GET));
assert!(is_safe_method(&Method::HEAD));
assert!(is_safe_method(&Method::OPTIONS));
assert!(!is_safe_method(&Method::POST));
assert!(!is_safe_method(&Method::PUT));
assert!(!is_safe_method(&Method::DELETE));
}
#[test]
fn test_is_public_path() {
assert!(is_public_path("/health", DEFAULT_PUBLIC_PATHS));
assert!(is_public_path("/metrics", DEFAULT_PUBLIC_PATHS));
assert!(is_public_path("/api/v1/auth/login", DEFAULT_PUBLIC_PATHS));
assert!(is_public_path("/api/v1/auth/refresh", DEFAULT_PUBLIC_PATHS));
assert!(!is_public_path("/health/ready", DEFAULT_PUBLIC_PATHS));
assert!(!is_public_path("/health_evil", DEFAULT_PUBLIC_PATHS));
assert!(!is_public_path(
"/api/v1/auth/login_anything",
DEFAULT_PUBLIC_PATHS
));
assert!(!is_public_path("/api/v1/data", DEFAULT_PUBLIC_PATHS));
}
#[test]
fn test_extract_cookie_value() {
let header = HeaderValue::from_static("csrf_token=abc123; other=value");
assert_eq!(
extract_cookie_value(Some(&header), "csrf_token"),
Some("abc123".to_string())
);
assert_eq!(
extract_cookie_value(Some(&header), "other"),
Some("value".to_string())
);
assert_eq!(extract_cookie_value(Some(&header), "missing"), None);
}
#[test]
fn test_constant_time_eq() {
assert!(constant_time_eq(b"abc", b"abc"));
assert!(!constant_time_eq(b"abc", b"abd"));
assert!(!constant_time_eq(b"abc", b"ab"));
assert!(constant_time_eq(b"", b""));
}
#[tokio::test]
async fn test_get_method_bypasses_csrf() {
let resp = send("GET", "/api/data", None, None).await;
assert_ne!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn test_post_without_csrf_returns_403() {
let resp = send("POST", "/api/data", None, None).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn test_post_with_mismatched_tokens_returns_403() {
let resp = send("POST", "/api/data", Some("csrf_token=abc"), Some("xyz")).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn test_post_with_matching_tokens_passes() {
let token = "valid_token_123";
let resp = send(
"POST",
"/api/data",
Some(&format!("csrf_token={}", token)),
Some(token),
)
.await;
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_post_with_only_cookie_returns_403() {
let resp = send("POST", "/api/data", Some("csrf_token=abc"), None).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn test_post_with_only_header_returns_403() {
let resp = send("POST", "/api/data", None, Some("abc")).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn test_public_path_bypasses_csrf() {
let resp = send("POST", "/health", None, None).await;
assert_ne!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn test_options_method_bypasses_csrf() {
let resp = send("OPTIONS", "/api/data", None, None).await;
assert_ne!(resp.status(), StatusCode::FORBIDDEN);
}
}