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 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)

View file

@ -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
})
}