diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 18383314ce..9434b02344 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -67,12 +67,13 @@ fn web_compression_predicate() -> impl Predicate { } fn web_body_limit_bytes() -> usize { + let value = std::env::var("DBX_MAX_UPLOAD_MB").ok(); + web_body_limit_bytes_from_value(value.as_deref()) +} + +fn web_body_limit_bytes_from_value(value: Option<&str>) -> usize { const DEFAULT_MB: usize = 1024; - let mb = std::env::var("DBX_MAX_UPLOAD_MB") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(DEFAULT_MB); + let mb = value.and_then(|value| value.parse::().ok()).filter(|value| *value > 0).unwrap_or(DEFAULT_MB); mb.saturating_mul(1024 * 1024) } @@ -877,7 +878,12 @@ async fn main() { .route("/sql-file/progress/{executionId}", get(routes::sql_file::sql_file_progress)) .route("/sql-file/cancel", post(routes::sql_file::cancel_sql_file)) // Table import - .route("/import/preview", post(routes::table_import::preview_import)) + .route( + "/import/preview", + post(routes::table_import::preview_import).layer(DefaultBodyLimit::max( + routes::table_import::import_request_body_limit_for_upload(web_body_limit_bytes()), + )), + ) .route("/import/preview-source", post(routes::table_import::preview_uploaded_import)) .route("/import/source/release", post(routes::table_import::release_import_source)) .route("/import/execute", post(routes::table_import::execute_import)) @@ -982,13 +988,15 @@ async fn main() { #[cfg(test)] mod tests { use super::{ - mount_public_base_path, normalize_public_base_path, web_agent_dir_from_env, web_compression_predicate, - XLSX_CONTENT_TYPE, + mount_public_base_path, normalize_public_base_path, web_agent_dir_from_env, web_body_limit_bytes_from_value, + web_compression_predicate, XLSX_CONTENT_TYPE, }; + use crate::routes::table_import; use axum::body::Body; + use axum::extract::{DefaultBodyLimit, Multipart}; use axum::http::header::CONTENT_TYPE; - use axum::http::Response; - use axum::routing::get; + use axum::http::{Response, StatusCode}; + use axum::routing::{get, post}; use axum::Router; use tower_http::compression::predicate::Predicate; @@ -1040,6 +1048,70 @@ mod tests { ); } + #[test] + fn web_upload_limit_parses_valid_values_and_preserves_safe_fallbacks() { + const MIB: usize = 1024 * 1024; + + assert_eq!(web_body_limit_bytes_from_value(None), 1024 * MIB); + assert_eq!(web_body_limit_bytes_from_value(Some("")), 1024 * MIB); + assert_eq!(web_body_limit_bytes_from_value(Some("0")), 1024 * MIB); + assert_eq!(web_body_limit_bytes_from_value(Some("invalid")), 1024 * MIB); + assert_eq!(web_body_limit_bytes_from_value(Some("4096")), 4096usize.saturating_mul(MIB)); + assert_eq!(web_body_limit_bytes_from_value(Some(&usize::MAX.to_string())), usize::MAX); + } + + #[tokio::test] + async fn import_route_reserves_multipart_framing_above_the_file_limit() { + const FILE_LIMIT: usize = 8; + + async fn uploaded_file_size(mut multipart: Multipart) -> Result { + let field = + multipart.next_field().await.map_err(|_| StatusCode::BAD_REQUEST)?.ok_or(StatusCode::BAD_REQUEST)?; + let bytes = field.bytes().await.map_err(|_| StatusCode::BAD_REQUEST)?; + Ok(bytes.len().to_string()) + } + + let router = Router::new() + .route("/general", post(uploaded_file_size)) + .route( + "/import", + post(uploaded_file_size) + .layer(DefaultBodyLimit::max(table_import::import_request_body_limit_for_upload(FILE_LIMIT))), + ) + .layer(DefaultBodyLimit::max(FILE_LIMIT)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind test listener"); + let address = listener.local_addr().expect("test listener address"); + let server = tokio::spawn(async move { + axum::serve(listener, router).await.expect("serve test router"); + }); + let boundary = "dbx-import-boundary"; + let body = format!( + "--{boundary}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"data.csv\"\r\n\r\n12345678\r\n--{boundary}--\r\n" + ); + let client = reqwest::Client::new(); + + let general_response = client + .post(format!("http://{address}/general")) + .header(CONTENT_TYPE, format!("multipart/form-data; boundary={boundary}")) + .body(body.clone()) + .send() + .await + .expect("send request through general limit"); + assert_eq!(general_response.status(), reqwest::StatusCode::BAD_REQUEST); + + let import_response = client + .post(format!("http://{address}/import")) + .header(CONTENT_TYPE, format!("multipart/form-data; boundary={boundary}")) + .body(body) + .send() + .await + .expect("send request through import limit"); + assert_eq!(import_response.status(), reqwest::StatusCode::OK); + assert_eq!(import_response.text().await.expect("read import response"), FILE_LIMIT.to_string()); + + server.abort(); + } + #[tokio::test] async fn public_base_path_routes_preserve_redirect_query_static_files_and_api() { let client = diff --git a/crates/dbx-web/src/routes/table_import.rs b/crates/dbx-web/src/routes/table_import.rs index 867cee1195..1000193ed4 100644 --- a/crates/dbx-web/src/routes/table_import.rs +++ b/crates/dbx-web/src/routes/table_import.rs @@ -19,9 +19,13 @@ use tokio::io::AsyncWriteExt; use crate::error::AppError; use crate::state::WebState; -const MAX_IMPORT_UPLOAD_BYTES: usize = 100 * 1024 * 1024; +pub(crate) const IMPORT_MULTIPART_OVERHEAD_BYTES: usize = 1024 * 1024; const TABLE_IMPORT_PROGRESS_TTL: Duration = Duration::from_secs(30); +pub(crate) fn import_request_body_limit_for_upload(max_upload_bytes: usize) -> usize { + max_upload_bytes.saturating_add(IMPORT_MULTIPART_OVERHEAD_BYTES) +} + fn initial_import_progress(import_id: &str, started_at: Instant) -> TableImportProgress { TableImportProgress { import_id: import_id.to_string(), @@ -227,7 +231,7 @@ pub async fn release_import_source( } async fn write_import_upload(field: axum::extract::multipart::Field<'_>, file_path: &StdPath) -> Result<(), AppError> { - write_import_upload_stream(field, file_path, MAX_IMPORT_UPLOAD_BYTES).await + write_import_upload_stream(field, file_path, crate::web_body_limit_bytes()).await } async fn write_import_upload_stream( @@ -461,6 +465,25 @@ mod tests { cleanup_uploaded_import_path(&file_path).await; } + #[test] + fn multipart_request_limit_adds_framing_without_reducing_file_limit() { + assert_eq!(import_request_body_limit_for_upload(8), IMPORT_MULTIPART_OVERHEAD_BYTES + 8); + assert_eq!(import_request_body_limit_for_upload(usize::MAX), usize::MAX); + } + + #[tokio::test] + async fn streams_import_upload_past_historical_limit_when_configured() { + const MIB: usize = 1024 * 1024; + const CONFIGURED_LIMIT: usize = 101 * MIB; + let file_path = test_upload_path(); + let chunk = Bytes::from(vec![b'x'; MIB]); + let chunks = stream::iter((0..101).map(|_| Ok::<_, String>(chunk.clone()))); + + assert!(write_import_upload_stream(chunks, &file_path, CONFIGURED_LIMIT).await.is_ok()); + assert_eq!(tokio::fs::metadata(&file_path).await.unwrap().len(), CONFIGURED_LIMIT as u64); + cleanup_uploaded_import_path(&file_path).await; + } + #[tokio::test] async fn removes_partial_upload_when_size_limit_is_exceeded() { let file_path = test_upload_path(); @@ -468,7 +491,7 @@ mod tests { let error = write_import_upload_stream(chunks, &file_path, 4).await.unwrap_err(); - assert!(error.message.contains("File too large")); + assert_eq!(error.message, "File too large: 5 bytes received (max 4 bytes)"); assert!(!file_path.exists()); } diff --git a/docs/content/docs/web-api.cn.mdx b/docs/content/docs/web-api.cn.mdx index ba38a915d5..8c71f660ac 100644 --- a/docs/content/docs/web-api.cn.mdx +++ b/docs/content/docs/web-api.cn.mdx @@ -89,7 +89,7 @@ Cookie: dbx_session=... | `DBX_PORT` | 修改监听端口，默认 `4224` | | `DBX_DATA_DIR` | 存放 `dbx.db` 的数据目录 | | `DBX_PUBLIC_BASE_PATH` | 在子路径下提供服务，例如 `/dbx` | -| `DBX_MAX_UPLOAD_MB` | 调整通用请求/上传大小上限，默认 1024 MB | +| `DBX_MAX_UPLOAD_MB` | 调整通用请求及上传文件大小上限，默认 1024 MB；multipart 路由会额外预留固定的协议开销（如需上传 4 GiB 文件，请设为至少 `4096`） | | `DBX_AGENT_DIR` | 覆盖 Web 端 Agent/驱动目录，默认位于数据目录下的 `agents` | | `DBX_STATIC_DIR` | 覆盖 Web 静态文件目录 | diff --git a/docs/content/docs/web-api.mdx b/docs/content/docs/web-api.mdx index 73f6882da2..33621301ee 100644 --- a/docs/content/docs/web-api.mdx +++ b/docs/content/docs/web-api.mdx @@ -89,7 +89,7 @@ Cookie: dbx_session=... | `DBX_PORT` | Change the listen port (default `4224`) | | `DBX_DATA_DIR` | Data directory containing `dbx.db` | | `DBX_PUBLIC_BASE_PATH` | Serve DBX under a subpath such as `/dbx` | -| `DBX_MAX_UPLOAD_MB` | Override the general request/upload limit; default 1024 MB | +| `DBX_MAX_UPLOAD_MB` | Override the general request and uploaded-file limit; default 1024 MB. Multipart routes reserve fixed protocol overhead (set at least `4096` for a 4 GiB file) | | `DBX_AGENT_DIR` | Override the Web Agent/driver directory; defaults to `agents` under the data directory | | `DBX_STATIC_DIR` | Override the Web static asset directory |