feat: extend the request body size limit and directly forward r2 to response body

This commit is contained in:
qaz741wsd856 2025-12-07 14:13:22 +00:00
parent 73223bfa94
commit 28c791337a
2 changed files with 58 additions and 16 deletions

View file

@ -1,9 +1,8 @@
use std::{collections::HashMap, sync::Arc}; use std::{collections::HashMap, sync::Arc};
use axum::{ use axum::{
body::Body,
extract::{Multipart, Path, Query, State}, extract::{Multipart, Path, Query, State},
http::{header, HeaderValue, StatusCode}, http::{header, StatusCode},
response::Response, response::Response,
Json, Json,
}; };
@ -13,7 +12,10 @@ use log;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use uuid::Uuid; 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::{ use crate::{
auth::Claims, auth::Claims,
@ -415,22 +417,32 @@ pub async fn download_attachment(
.map_err(AppError::Worker)? .map_err(AppError::Worker)?
.ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?; .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() .body()
.ok_or_else(|| AppError::NotFound("Attachment data not found".to_string()))? .ok_or_else(|| AppError::NotFound("Attachment data not found".to_string()))?
.bytes() .response_body()
.await
.map_err(AppError::Worker)?; .map_err(AppError::Worker)?;
let mut builder = Response::builder().status(StatusCode::OK); let headers = WorkerHeaders::new();
builder = builder.header(header::CONTENT_TYPE, "application/octet-stream"); let content_type = http_metadata
if let Ok(value) = HeaderValue::from_str(&body.len().to_string()) { .content_type
builder = builder.header(header::CONTENT_LENGTH, value); .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 let response = WorkerResponse::from_body(response_body)
.body(Body::from(body)) .map_err(AppError::Worker)?
.map_err(|_| AppError::Internal) .with_status(StatusCode::OK.as_u16())
.with_headers(headers);
Ok(response.into())
} }
/// Attach attachment information to Cipher (used by other handlers) /// Attach attachment information to Cipher (used by other handlers)

View file

@ -1,4 +1,4 @@
use axum::Extension; use axum::{extract::DefaultBodyLimit, Extension};
use tower_http::cors::{Any, CorsLayer}; use tower_http::cors::{Any, CorsLayer};
use tower_service::Service; use tower_service::Service;
use worker::*; use worker::*;
@ -39,9 +39,13 @@ pub async fn main(
.allow_headers(Any) .allow_headers(Any)
.allow_origin(Any); .allow_origin(Any);
let body_limit = attachment_body_limit_bytes(&env);
let mut app = router::api_router(env) let mut app = router::api_router(env)
.layer(Extension(BaseUrl(base_url))) .layer(Extension(BaseUrl(base_url)))
.layer(cors); .layer(cors)
// axum 默认 body 限制为 2MiB附件上传需要更大的上限
.layer(DefaultBodyLimit::max(body_limit));
Ok(app.call(req).await?) 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
})
}