feat: extend the request body size limit and directly forward r2 to response body
This commit is contained in:
parent
73223bfa94
commit
28c791337a
2 changed files with 58 additions and 16 deletions
|
|
@ -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)
|
||||
|
|
|
|||
34
src/lib.rs
34
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
|
||||
})
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue