diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index dd0db43e92..2756284ca3 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -2162,7 +2162,7 @@ export async function startTableExport(request: TableExportRequest, onProgress: finish(() => reject(new Error(progress.errorMessage || "Export failed"))); } else if (progress.status === "Done") { // Trigger browser download - downloadTableExportFile(exportId, request.format); + downloadTableExportFile(exportId); finish(() => resolve(progress)); } else { finish(() => resolve(progress)); @@ -2176,11 +2176,9 @@ export async function startTableExport(request: TableExportRequest, onProgress: }); } -function downloadTableExportFile(exportId: string, format: string): void { - const ext = format === "markdown" || format === "md" ? "md" : format; +function downloadTableExportFile(exportId: string): void { const a = document.createElement("a"); a.href = apiUrl(`/api/export/table/download/${exportId}`); - a.download = `table_export_${exportId}.${ext}`; a.click(); } @@ -2218,7 +2216,7 @@ export async function startQueryResultExport(request: QueryResultExportRequest, if (progress.status === "Error") { finish(() => reject(new Error(progress.errorMessage || "Export failed"))); } else if (progress.status === "Done") { - downloadQueryResultExportFile(exportId, request.format); + downloadQueryResultExportFile(exportId); finish(() => resolve(progress)); } else { finish(() => resolve(progress)); @@ -2232,10 +2230,9 @@ export async function startQueryResultExport(request: QueryResultExportRequest, }); } -function downloadQueryResultExportFile(exportId: string, format: string): void { +function downloadQueryResultExportFile(exportId: string): void { const a = document.createElement("a"); a.href = apiUrl(`/api/export/query-result/download/${exportId}`); - a.download = `query_result_export_${exportId}.${format}`; a.click(); } diff --git a/crates/dbx-web/src/routes/database_export.rs b/crates/dbx-web/src/routes/database_export.rs index 523f6c50f0..aba89f79f1 100644 --- a/crates/dbx-web/src/routes/database_export.rs +++ b/crates/dbx-web/src/routes/database_export.rs @@ -14,7 +14,8 @@ use futures::stream::Stream; use serde::Deserialize; use crate::error::AppError; -use crate::state::WebState; +use crate::routes::export_download::attachment_content_disposition; +use crate::state::{WebExportFile, WebState}; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] @@ -48,7 +49,10 @@ pub async fn start_database_export( let download_filename = format!("{}.sql", sanitize_export_filename(&req.database)); // Store export file mapping for download - state.export_files.write().await.insert(export_id.clone(), (tmp_file_path.clone(), download_filename)); + state.export_files.write().await.insert( + export_id.clone(), + WebExportFile { file_path: tmp_file_path.clone(), download_filename, format: "sql".to_string() }, + ); let (tx, _) = tokio::sync::broadcast::channel::(256); state.sse_channels.write().await.insert(export_id.clone(), tx.clone()); @@ -147,21 +151,21 @@ pub async fn database_export_download( State(state): State>, Path(export_id): Path, ) -> Result { - let (file_path, download_filename) = state + let export_file = state .export_files .write() .await .remove(&export_id) .ok_or_else(|| AppError::from("Export file not found".to_string()))?; - let data = tokio::fs::read(&file_path).await.map_err(|e| AppError::from(e.to_string()))?; + let data = tokio::fs::read(&export_file.file_path).await.map_err(|e| AppError::from(e.to_string()))?; // Clean up temp file - let _ = tokio::fs::remove_file(&file_path).await; + let _ = tokio::fs::remove_file(&export_file.file_path).await; Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, "application/sql; charset=utf-8") - .header(header::CONTENT_DISPOSITION, format!("attachment; filename=\"{download_filename}\"")) + .header(header::CONTENT_DISPOSITION, attachment_content_disposition(&export_file.download_filename)) .body(Body::from(data)) .map_err(|e| AppError::from(e.to_string())) } diff --git a/crates/dbx-web/src/routes/export_download.rs b/crates/dbx-web/src/routes/export_download.rs new file mode 100644 index 0000000000..cfa55f2202 --- /dev/null +++ b/crates/dbx-web/src/routes/export_download.rs @@ -0,0 +1,83 @@ +use std::fmt::Write; + +pub(crate) fn export_download_filename(requested_path: &str, fallback_base: &str, extension: &str) -> String { + let requested_name = requested_path.rsplit(['/', '\\']).next().unwrap_or_default().trim(); + let fallback_name = format!("{fallback_base}.{extension}"); + let is_web_placeholder = requested_name.starts_with("__web_export_") + && requested_name.to_ascii_lowercase().ends_with(&format!(".{extension}").to_ascii_lowercase()); + let source = if requested_name.is_empty() || is_web_placeholder { fallback_name.as_str() } else { requested_name }; + let mut sanitized = source + .chars() + .map(|character| { + if character.is_control() || matches!(character, '"' | ';' | '/' | '\\' | '<' | '>' | ':' | '|' | '?' | '*') + { + '_' + } else { + character + } + }) + .collect::(); + sanitized = sanitized.trim_matches([' ', '.']).to_string(); + if sanitized.is_empty() { + sanitized = fallback_name; + } + + let expected_suffix = format!(".{extension}"); + if !sanitized.to_ascii_lowercase().ends_with(&expected_suffix.to_ascii_lowercase()) { + sanitized.push_str(&expected_suffix); + } + sanitized +} + +pub(crate) fn attachment_content_disposition(file_name: &str) -> String { + let fallback = file_name + .chars() + .map(|character| { + if character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.' | ' ' | '(' | ')') { + character + } else { + '_' + } + }) + .collect::(); + format!("attachment; filename=\"{fallback}\"; filename*=UTF-8''{}", encode_rfc5987_value(file_name)) +} + +fn encode_rfc5987_value(value: &str) -> String { + let mut encoded = String::with_capacity(value.len()); + for byte in value.bytes() { + if byte.is_ascii_alphanumeric() + || matches!(byte, b'!' | b'#' | b'$' | b'&' | b'+' | b'-' | b'.' | b'^' | b'_' | b'`' | b'|' | b'~') + { + encoded.push(char::from(byte)); + } else { + write!(&mut encoded, "%{byte:02X}").expect("writing to a String cannot fail"); + } + } + encoded +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn export_download_names_keep_the_requested_basename_and_expected_extension() { + assert_eq!( + export_download_filename(r"C:\\exports\\custom_question_260812161843.xlsx", "agents", "xlsx"), + "custom_question_260812161843.xlsx" + ); + assert_eq!(export_download_filename("__web_export_123.xlsx", "agents", "xlsx"), "agents.xlsx"); + assert_eq!(export_download_filename("../../unsafe.csv", "agents", "xlsx"), "unsafe.csv.xlsx"); + assert_eq!(export_download_filename("\r\n\".xlsx", "agents", "xlsx"), "_.xlsx"); + assert_eq!(export_download_filename("", "agents", "xlsx"), "agents.xlsx"); + } + + #[test] + fn content_disposition_supports_unicode_without_unsafe_header_characters() { + assert_eq!( + attachment_content_disposition("智能体列表.xlsx"), + "attachment; filename=\"_____.xlsx\"; filename*=UTF-8''%E6%99%BA%E8%83%BD%E4%BD%93%E5%88%97%E8%A1%A8.xlsx" + ); + } +} diff --git a/crates/dbx-web/src/routes/mod.rs b/crates/dbx-web/src/routes/mod.rs index 47d3fa747b..3586461803 100644 --- a/crates/dbx-web/src/routes/mod.rs +++ b/crates/dbx-web/src/routes/mod.rs @@ -10,6 +10,7 @@ pub mod dialect; pub mod docs; pub mod document_store; pub mod etcd; +pub(crate) mod export_download; pub mod hbase; pub mod history; pub mod jdbc; diff --git a/crates/dbx-web/src/routes/query_result_export.rs b/crates/dbx-web/src/routes/query_result_export.rs index f74b72449e..d1d196055b 100644 --- a/crates/dbx-web/src/routes/query_result_export.rs +++ b/crates/dbx-web/src/routes/query_result_export.rs @@ -18,7 +18,8 @@ use futures::stream::Stream; use serde::Deserialize; use crate::error::AppError; -use crate::state::WebState; +use crate::routes::export_download::{attachment_content_disposition, export_download_filename}; +use crate::state::{WebExportFile, WebState}; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] @@ -49,9 +50,14 @@ pub async fn start_query_result_export( std::fs::create_dir_all(&tmp_dir).map_err(|e| AppError::from(e.to_string()))?; let tmp_file = tmp_dir.join(format!("query_result_export_{export_id}.{ext}")); let file_path = tmp_file.to_string_lossy().to_string(); + let download_filename = export_download_filename(&req.file_path, "query-result", ext); req.file_path = file_path.clone(); - state.export_files.write().await.insert(export_id.clone(), (file_path, req.format.clone())); + state + .export_files + .write() + .await + .insert(export_id.clone(), WebExportFile { file_path, download_filename, format: req.format.clone() }); let tx = { let mut channels = state.sse_channels.write().await; @@ -141,28 +147,67 @@ pub async fn query_result_export_download( State(state): State>, Path(export_id): Path, ) -> Result { - let (file_path, format) = state + let export_file = state .export_files .write() .await .remove(&export_id) .ok_or_else(|| AppError::from("Export file not found".to_string()))?; - let data = tokio::fs::read(&file_path).await.map_err(|e| AppError::from(e.to_string()))?; - let _ = tokio::fs::remove_file(&file_path).await; + let data = tokio::fs::read(&export_file.file_path).await.map_err(|e| AppError::from(e.to_string()))?; + let _ = tokio::fs::remove_file(&export_file.file_path).await; - let (content_type, file_ext) = match format.as_str() { - "csv" => ("text/csv; charset=utf-8", "csv"), - "xlsx" => ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "xlsx"), - _ => return Err(AppError::from(format!("Unknown format: {format}"))), + let content_type = match export_file.format.as_str() { + "csv" => "text/csv; charset=utf-8", + "xlsx" => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + format => return Err(AppError::from(format!("Unknown format: {format}"))), }; - let filename = format!("query_result_export_{export_id}.{file_ext}"); - Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, content_type) - .header(header::CONTENT_DISPOSITION, format!("attachment; filename=\"{filename}\"")) + .header(header::CONTENT_DISPOSITION, attachment_content_disposition(&export_file.download_filename)) .body(Body::from(data)) .map_err(|e| AppError::from(e.to_string())) } + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::to_bytes; + use dbx_core::connection::AppState; + use dbx_core::storage::Storage; + + use crate::state::WebExportFile; + + #[tokio::test] + async fn query_result_export_download_uses_the_requested_web_filename() { + let dir = std::env::temp_dir().join(format!("dbx-web-query-export-download-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let app = Arc::new(AppState::new_with_plugin_dir(storage, dir.join("plugins"))); + let state = Arc::new(WebState::for_tests(app, dir.clone())); + let file_path = dir.join("query-export.csv"); + tokio::fs::write(&file_path, b"id,name\n1,test").await.unwrap(); + state.export_files.write().await.insert( + "query-export-id".to_string(), + WebExportFile { + file_path: file_path.to_string_lossy().to_string(), + download_filename: "agents_260812161843.csv".to_string(), + format: "csv".to_string(), + }, + ); + + let response = + query_result_export_download(State(state.clone()), Path("query-export-id".to_string())).await.unwrap(); + + assert_eq!( + response.headers().get(header::CONTENT_DISPOSITION).unwrap().to_str().unwrap(), + "attachment; filename=\"agents_260812161843.csv\"; filename*=UTF-8''agents_260812161843.csv" + ); + assert_eq!(to_bytes(response.into_body(), usize::MAX).await.unwrap().as_ref(), b"id,name\n1,test"); + assert!(!file_path.exists()); + assert!(!state.export_files.read().await.contains_key("query-export-id")); + std::fs::remove_dir_all(dir).ok(); + } +} diff --git a/crates/dbx-web/src/routes/table_export.rs b/crates/dbx-web/src/routes/table_export.rs index d19291a46d..797ea95b29 100644 --- a/crates/dbx-web/src/routes/table_export.rs +++ b/crates/dbx-web/src/routes/table_export.rs @@ -16,7 +16,8 @@ use futures::stream::Stream; use serde::Deserialize; use crate::error::AppError; -use crate::state::WebState; +use crate::routes::export_download::{attachment_content_disposition, export_download_filename}; +use crate::state::{WebExportFile, WebState}; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] @@ -50,10 +51,15 @@ pub async fn start_table_export( }; let tmp_file = tmp_dir.join(format!("table_export_{export_id}.{ext}")); let file_path = tmp_file.to_string_lossy().to_string(); + let download_filename = export_download_filename(&req.file_path, &req.table_name, ext); req.file_path = file_path.clone(); // Store export file mapping for download - state.export_files.write().await.insert(export_id.clone(), (file_path, req.format.clone())); + state + .export_files + .write() + .await + .insert(export_id.clone(), WebExportFile { file_path, download_filename, format: req.format.clone() }); let tx = { let mut channels = state.sse_channels.write().await; @@ -127,32 +133,70 @@ pub async fn table_export_download( State(state): State>, Path(export_id): Path, ) -> Result { - let (file_path, format) = state + let export_file = state .export_files .write() .await .remove(&export_id) .ok_or_else(|| AppError::from("Export file not found".to_string()))?; - let data = tokio::fs::read(&file_path).await.map_err(|e| AppError::from(e.to_string()))?; + let data = tokio::fs::read(&export_file.file_path).await.map_err(|e| AppError::from(e.to_string()))?; // Clean up temp file - let _ = tokio::fs::remove_file(&file_path).await; + let _ = tokio::fs::remove_file(&export_file.file_path).await; - let (content_type, file_ext) = match format.as_str() { - "csv" => ("text/csv; charset=utf-8", "csv"), - "xlsx" => ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "xlsx"), - "json" => ("application/json; charset=utf-8", "json"), - "markdown" | "md" => ("text/markdown; charset=utf-8", "md"), - "sql" => ("application/sql; charset=utf-8", "sql"), - _ => return Err(AppError::from(format!("Unknown format: {format}"))), + let content_type = match export_file.format.as_str() { + "csv" => "text/csv; charset=utf-8", + "xlsx" => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "json" => "application/json; charset=utf-8", + "markdown" | "md" => "text/markdown; charset=utf-8", + "sql" => "application/sql; charset=utf-8", + format => return Err(AppError::from(format!("Unknown format: {format}"))), }; - let filename = format!("table_export_{export_id}.{file_ext}"); - Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, content_type) - .header(header::CONTENT_DISPOSITION, format!("attachment; filename=\"{filename}\"")) + .header(header::CONTENT_DISPOSITION, attachment_content_disposition(&export_file.download_filename)) .body(Body::from(data)) .map_err(|e| AppError::from(e.to_string())) } + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::to_bytes; + use dbx_core::connection::AppState; + use dbx_core::storage::Storage; + + use crate::state::WebExportFile; + + #[tokio::test] + async fn table_export_download_uses_the_requested_web_filename() { + let dir = std::env::temp_dir().join(format!("dbx-web-table-export-download-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let app = Arc::new(AppState::new_with_plugin_dir(storage, dir.join("plugins"))); + let state = Arc::new(WebState::for_tests(app, dir.clone())); + let file_path = dir.join("table-export.xlsx"); + tokio::fs::write(&file_path, b"xlsx-data").await.unwrap(); + state.export_files.write().await.insert( + "table-export-id".to_string(), + WebExportFile { + file_path: file_path.to_string_lossy().to_string(), + download_filename: "智能体列表.xlsx".to_string(), + format: "xlsx".to_string(), + }, + ); + + let response = table_export_download(State(state.clone()), Path("table-export-id".to_string())).await.unwrap(); + + assert_eq!( + response.headers().get(header::CONTENT_DISPOSITION).unwrap().to_str().unwrap(), + "attachment; filename=\"_____.xlsx\"; filename*=UTF-8''%E6%99%BA%E8%83%BD%E4%BD%93%E5%88%97%E8%A1%A8.xlsx" + ); + assert_eq!(to_bytes(response.into_body(), usize::MAX).await.unwrap().as_ref(), b"xlsx-data"); + assert!(!file_path.exists()); + assert!(!state.export_files.read().await.contains_key("table-export-id")); + std::fs::remove_dir_all(dir).ok(); + } +} diff --git a/crates/dbx-web/src/state.rs b/crates/dbx-web/src/state.rs index deb5d93d3b..3083632db0 100644 --- a/crates/dbx-web/src/state.rs +++ b/crates/dbx-web/src/state.rs @@ -12,6 +12,13 @@ pub struct LoginRateLimit { pub locked_until: Option, } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WebExportFile { + pub file_path: String, + pub download_filename: String, + pub format: String, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct NacosImportContext { pub owner_session: Option, @@ -33,8 +40,8 @@ pub struct WebState { pub sql_file_executions: RwLock>, pub nacos_imports: RwLock>, pub login_rate_limit: Mutex, - /// Table export temp files: export_id -> (file_path, format) - pub export_files: RwLock>, + /// Completed Web export temp files waiting for the browser download. + pub export_files: RwLock>, pub ssh_prompts: Arc, }