From e1fb225034521a31779c1554a5052d97ea2d45ff Mon Sep 17 00:00:00 2001
From: kqjy <jun@jzwsite.com>
Date: Wed, 22 Apr 2026 23:01:32 +0800
Subject: [PATCH] csrf fixes

---
 crates/myfsio-server/src/handlers/ui.rs       | 10 ------
 crates/myfsio-server/src/lib.rs               |  8 +++--
 .../myfsio-server/src/middleware/session.rs   | 33 +++++++++++++++++--
 3 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/crates/myfsio-server/src/handlers/ui.rs b/crates/myfsio-server/src/handlers/ui.rs
index 3d94114..e08473d 100644
--- a/crates/myfsio-server/src/handlers/ui.rs
+++ b/crates/myfsio-server/src/handlers/ui.rs
@@ -117,16 +117,6 @@ pub async fn logout(Extension(session): Extension<SessionHandle>) -> Response {
     Redirect::to("/login").into_response()
 }
 
-pub async fn csrf_error_page(
-    State(state): State<AppState>,
-    Extension(session): Extension<SessionHandle>,
-) -> Response {
-    let ctx = base_context(&session, None);
-    let mut resp = render(&state, "csrf_error.html", &ctx);
-    *resp.status_mut() = StatusCode::FORBIDDEN;
-    resp
-}
-
 pub async fn root_redirect() -> Response {
     Redirect::to("/ui/buckets").into_response()
 }
diff --git a/crates/myfsio-server/src/lib.rs b/crates/myfsio-server/src/lib.rs
index 217f2b8..93c7b77 100644
--- a/crates/myfsio-server/src/lib.rs
+++ b/crates/myfsio-server/src/lib.rs
@@ -304,8 +304,7 @@ pub fn create_ui_router(state: state::AppState) -> Router {
 
     let public = Router::new()
         .route("/login", get(ui::login_page).post(ui::login_submit))
-        .route("/logout", post(ui::logout).get(ui::logout))
-        .route("/csrf-error", get(ui::csrf_error_page));
+        .route("/logout", post(ui::logout).get(ui::logout));
 
     let session_state = middleware::SessionLayerState {
         store: state.sessions.clone(),
@@ -317,7 +316,10 @@ pub fn create_ui_router(state: state::AppState) -> Router {
     protected
         .merge(public)
         .fallback(ui::not_found_page)
-        .layer(axum::middleware::from_fn(middleware::csrf_layer))
+        .layer(axum::middleware::from_fn_with_state(
+            state.clone(),
+            middleware::csrf_layer,
+        ))
         .layer(axum::middleware::from_fn_with_state(
             session_state,
             middleware::session_layer,
diff --git a/crates/myfsio-server/src/middleware/session.rs b/crates/myfsio-server/src/middleware/session.rs
index 60618bc..68f8207 100644
--- a/crates/myfsio-server/src/middleware/session.rs
+++ b/crates/myfsio-server/src/middleware/session.rs
@@ -90,7 +90,11 @@ pub async fn session_layer(
     resp
 }
 
-pub async fn csrf_layer(req: Request, next: Next) -> Response {
+pub async fn csrf_layer(
+    State(state): State<crate::state::AppState>,
+    req: Request,
+    next: Next,
+) -> Response {
     const CSRF_HEADER_ALIAS: &str = "x-csrftoken";
 
     let method = req.method().clone();
@@ -169,7 +173,32 @@ pub async fn csrf_layer(req: Request, next: Next) -> Response {
         header_present = header_token.is_some(),
         "CSRF token mismatch"
     );
-    (StatusCode::FORBIDDEN, "Invalid CSRF token").into_response()
+
+    let accept = parts
+        .headers
+        .get(header::ACCEPT)
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or("");
+    let is_form_submit = content_type.starts_with("application/x-www-form-urlencoded")
+        || content_type.starts_with("multipart/form-data");
+    let wants_json = accept.contains("application/json")
+        || content_type.starts_with("application/json");
+
+    if is_form_submit && !wants_json {
+        let ctx = crate::handlers::ui::base_context(&handle, None);
+        let mut resp = crate::handlers::ui::render(&state, "csrf_error.html", &ctx);
+        *resp.status_mut() = StatusCode::FORBIDDEN;
+        return resp;
+    }
+
+    let mut resp = (
+        StatusCode::FORBIDDEN,
+        [(header::CONTENT_TYPE, "application/json")],
+        r#"{"error":"Invalid CSRF token"}"#,
+    )
+        .into_response();
+    *resp.status_mut() = StatusCode::FORBIDDEN;
+    resp
 }
 
 fn extract_multipart_token(content_type: &str, body: &[u8]) -> Option<String> {