From 28c791337a889d6c627f16c595bb0ad389ac7de4 Mon Sep 17 00:00:00 2001
From: qaz741wsd856 <hysth@163.com>
Date: Sun, 7 Dec 2025 14:13:22 +0000
Subject: [PATCH] feat: extend the request body size limit and directly forward
 r2 to response body

---
 src/handlers/attachments.rs | 40 ++++++++++++++++++++++++-------------
 src/lib.rs                  | 34 +++++++++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/src/handlers/attachments.rs b/src/handlers/attachments.rs
index 62a4521..3175f7d 100644
--- a/src/handlers/attachments.rs
+++ b/src/handlers/attachments.rs
@@ -1,9 +1,8 @@
 use std::{collections::HashMap, sync::Arc};
 
 use axum::{
-    body::Body,
     extract::{Multipart, Path, Query, State},
-    http::{header, HeaderValue, StatusCode},
+    http::{header, StatusCode},
     response::Response,
     Json,
 };
@@ -13,7 +12,10 @@ use log;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use uuid::Uuid;
-use worker::{query, wasm_bindgen::JsValue, Bucket, D1Database, Env, HttpMetadata};
+use worker::{
+    query, wasm_bindgen::JsValue, Bucket, D1Database, Env, Headers as WorkerHeaders, HttpMetadata,
+    Response as WorkerResponse,
+};
 
 use crate::{
     auth::Claims,
@@ -415,22 +417,32 @@ pub async fn download_attachment(
         .map_err(AppError::Worker)?
         .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?;
 
-    let body = object
+    let content_length = object.size();
+    let http_metadata = object.http_metadata();
+
+    let response_body = object
         .body()
         .ok_or_else(|| AppError::NotFound("Attachment data not found".to_string()))?
-        .bytes()
-        .await
+        .response_body()
         .map_err(AppError::Worker)?;
 
-    let mut builder = Response::builder().status(StatusCode::OK);
-    builder = builder.header(header::CONTENT_TYPE, "application/octet-stream");
-    if let Ok(value) = HeaderValue::from_str(&body.len().to_string()) {
-        builder = builder.header(header::CONTENT_LENGTH, value);
-    }
+    let headers = WorkerHeaders::new();
+    let content_type = http_metadata
+        .content_type
+        .unwrap_or_else(|| "application/octet-stream".to_string());
+    headers
+        .set(header::CONTENT_TYPE.as_str(), &content_type)
+        .map_err(AppError::Worker)?;
+    headers
+        .set(header::CONTENT_LENGTH.as_str(), &content_length.to_string())
+        .map_err(AppError::Worker)?;
 
-    builder
-        .body(Body::from(body))
-        .map_err(|_| AppError::Internal)
+    let response = WorkerResponse::from_body(response_body)
+        .map_err(AppError::Worker)?
+        .with_status(StatusCode::OK.as_u16())
+        .with_headers(headers);
+
+    Ok(response.into())
 }
 
 /// Attach attachment information to Cipher (used by other handlers)
diff --git a/src/lib.rs b/src/lib.rs
index 9db7306..ff9b683 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,4 @@
-use axum::Extension;
+use axum::{extract::DefaultBodyLimit, Extension};
 use tower_http::cors::{Any, CorsLayer};
 use tower_service::Service;
 use worker::*;
@@ -39,9 +39,13 @@ pub async fn main(
         .allow_headers(Any)
         .allow_origin(Any);
 
+    let body_limit = attachment_body_limit_bytes(&env);
+
     let mut app = router::api_router(env)
         .layer(Extension(BaseUrl(base_url)))
-        .layer(cors);
+        .layer(cors)
+        // axum 默认 body 限制为 2MiB，附件上传需要更大的上限
+        .layer(DefaultBodyLimit::max(body_limit));
 
     Ok(app.call(req).await?)
 }
@@ -68,3 +72,29 @@ pub async fn scheduled(_event: ScheduledEvent, env: Env, _ctx: ScheduleContext)
         }
     }
 }
+
+/// Resolve a permissive body size limit, prioritizing ATTACHMENT_MAX_BYTES and falling back to 64MiB.
+fn attachment_body_limit_bytes(env: &Env) -> usize {
+    const DEFAULT_LIMIT: usize = 64 * 1024 * 1024;
+
+    let max_bytes = env.var("ATTACHMENT_MAX_BYTES").ok().and_then(|v| {
+        let raw = v.to_string();
+        match raw.parse::<u64>() {
+            Ok(val) => Some(val),
+            Err(err) => {
+                log::error!("Invalid ATTACHMENT_MAX_BYTES '{}': {}", raw, err);
+                None
+            }
+        }
+    });
+
+    let limit_u64 = max_bytes.unwrap_or(DEFAULT_LIMIT as u64);
+
+    limit_u64.try_into().unwrap_or_else(|_| {
+        log::error!(
+            "Attachment body limit {} did not fit into usize, falling back to default",
+            limit_u64
+        );
+        DEFAULT_LIMIT
+    })
+}