fix(web): honor upload limit for table imports

Closes #5957

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::<usize>().ok())
-        .filter(|value| *value > 0)
-        .unwrap_or(DEFAULT_MB);
+    let mb = value.and_then(|value| value.parse::<usize>().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<String, StatusCode> {
+            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 =

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<S, E>(
@@ -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());
     }