diff --git a/aw-server/src/endpoints/util.rs b/aw-server/src/endpoints/util.rs index 8fce4582..a96c583f 100644 --- a/aw-server/src/endpoints/util.rs +++ b/aw-server/src/endpoints/util.rs @@ -1,4 +1,6 @@ -use std::io::{Cursor, Seek, SeekFrom}; +use std::fs::File; +use std::io::{copy, pipe, Cursor, PipeReader, PipeWriter, Seek, SeekFrom}; +use std::thread; use rocket::http::ContentType; use rocket::http::Header; @@ -38,43 +40,121 @@ impl<'r> Responder<'r, 'static> for HttpErrorJson { } pub struct BucketsExportRocket { - file: std::fs::File, + datastore: aw_datastore::Datastore, + bucket_id: Option, filename: String, } +fn export_filename( + datastore: &aw_datastore::Datastore, + bucket_id: Option<&str>, +) -> Result { + let name = match bucket_id { + Some(id) => { + datastore.get_bucket(id)?; + Some(id.to_owned()) + } + None => { + let buckets = datastore.get_buckets()?; + (buckets.len() == 1).then(|| buckets.into_keys().next().unwrap()) + } + }; + Ok(match name { + Some(id) => format!("attachment; filename=aw-bucket-export_{id}.json"), + None => "attachment; filename=aw-buckets-export.json".into(), + }) +} + +#[cfg(not(any(unix, windows)))] +compile_error!("export streaming requires unix or windows anonymous pipes"); + +fn pipe_reader_to_file(reader: PipeReader) -> File { + #[cfg(unix)] + { + File::from(std::os::fd::OwnedFd::from(reader)) + } + #[cfg(windows)] + { + File::from(std::os::windows::io::OwnedHandle::from(reader)) + } +} + +fn pipe_writer_to_file(writer: PipeWriter) -> File { + #[cfg(unix)] + { + File::from(std::os::fd::OwnedFd::from(writer)) + } + #[cfg(windows)] + { + File::from(std::os::windows::io::OwnedHandle::from(writer)) + } +} + +/// Serialize on the datastore worker into a private tempfile, then copy to +/// the client pipe from this thread. The worker stays disk-paced; a slow +/// or dropped download must not stall heartbeats (see `ServerState`). +fn spawn_export_stream( + datastore: aw_datastore::Datastore, + bucket_id: Option, + writer: PipeWriter, +) { + thread::spawn(move || { + let staging = match tempfile::tempfile() { + Ok(file) => file, + Err(err) => { + error!("Failed to create export staging file: {err}"); + return; + } + }; + let mut staging = match datastore.export_to_file(bucket_id.as_deref(), staging) { + Ok((file, _)) => file, + Err(err) => { + error!("Export stream failed: {err:?}"); + return; + } + }; + if let Err(err) = staging.seek(SeekFrom::Start(0)) { + error!("Failed to rewind export staging file: {err}"); + return; + } + let mut writer = pipe_writer_to_file(writer); + if let Err(err) = copy(&mut staging, &mut writer) { + error!("Export stream copy failed: {err}"); + } + }); +} + impl BucketsExportRocket { pub fn new( datastore: &aw_datastore::Datastore, bucket_id: Option<&str>, ) -> Result { - let io_error = |err: std::io::Error| { - error!("Failed to prepare export file: {err}"); - HttpErrorJson::new( - Status::InternalServerError, - "Failed to prepare export file".into(), - ) - }; - // tempfile creates a private file and removes it when the response is - // dropped. Spooling preserves HTTP errors even if serialization or disk - // writes fail, while keeping event buffering bounded. - let file = tempfile::tempfile().map_err(io_error)?; - let (mut file, name) = datastore.export_to_file(bucket_id, file)?; - file.seek(SeekFrom::Start(0)).map_err(io_error)?; - let filename = match name { - Some(id) => format!("attachment; filename=aw-bucket-export_{id}.json"), - None => "attachment; filename=aw-buckets-export.json".into(), - }; - Ok(Self { file, filename }) + // Resolve the download name and 404 missing buckets before the + // response is built. Serialization itself runs after headers so a + // slow export does not look like a hung connection. + let filename = export_filename(datastore, bucket_id)?; + Ok(Self { + datastore: datastore.clone(), + bucket_id: bucket_id.map(str::to_owned), + filename, + }) } } impl<'r> Responder<'r, 'static> for BucketsExportRocket { fn respond_to(self, _: &Request) -> response::Result<'static> { + let (reader, writer) = pipe().map_err(|err| { + error!("Failed to open export pipe: {err}"); + Status::InternalServerError + })?; + spawn_export_stream(self.datastore, self.bucket_id, writer); Response::build() .status(Status::Ok) .header(Header::new("Content-Disposition", self.filename)) .header(ContentType::JSON) - .streamed_body(rocket::tokio::fs::File::from_std(self.file)) + .streamed_body(rocket::tokio::fs::File::from_std(pipe_reader_to_file( + reader, + ))) .ok() } } diff --git a/aw-server/tests/api.rs b/aw-server/tests/api.rs index 99690f00..7a7b7914 100644 --- a/aw-server/tests/api.rs +++ b/aw-server/tests/api.rs @@ -71,6 +71,72 @@ mod api_tests { assert!(body["message"].as_str().unwrap().contains("does not exist")); } + #[test] + fn export_all_empty_uses_plural_filename_and_opens_before_body() { + let server = setup_testserver(); + let client = Client::untracked(server).unwrap(); + let response = client + .get("/api/0/export") + .header(Header::new("Host", "127.0.0.1:5600")) + .dispatch(); + assert_eq!(response.status(), Status::Ok); + assert_eq!(response.content_type(), Some(ContentType::JSON)); + assert_eq!( + response.headers().get_one("Content-Disposition"), + Some("attachment; filename=aw-buckets-export.json") + ); + let body: Value = serde_json::from_str(&response.into_string().unwrap()).unwrap(); + assert_eq!(body["buckets"], json!({})); + } + + #[test] + fn unread_export_body_does_not_block_datastore() { + let server = setup_testserver(); + let datastore = server + .state::() + .unwrap() + .datastore + .clone(); + let bucket: Bucket = serde_json::from_value(json!({ + "id": "big", "type": "test", "client": "test", "hostname": "test" + })) + .unwrap(); + datastore.create_bucket(&bucket).unwrap(); + let mut event = aw_models::Event::default(); + event + .data + .insert("blob".into(), json!("x".repeat(128 * 1024))); + datastore.insert_events("big", &[event]).unwrap(); + + let client = Client::untracked(server).unwrap(); + let response = client + .get("/api/0/buckets/big/export") + .header(Header::new("Host", "127.0.0.1:5600")) + .dispatch(); + assert_eq!(response.status(), Status::Ok); + + // Let Command::Export start. A client-paced pipe would fill here and + // stall the worker; staging to a tempfile must not. + std::thread::sleep(std::time::Duration::from_millis(100)); + let (tx, rx) = std::sync::mpsc::channel(); + let ds = datastore.clone(); + std::thread::spawn(move || { + let _ = tx.send(ds.get_buckets()); + }); + rx.recv_timeout(std::time::Duration::from_secs(3)) + .expect("datastore worker blocked by unread export body") + .unwrap(); + + let body: Value = serde_json::from_str(&response.into_string().unwrap()).unwrap(); + assert_eq!( + body["buckets"]["big"]["events"][0]["data"]["blob"] + .as_str() + .unwrap() + .len(), + 128 * 1024 + ); + } + #[test] fn test_bucket() { let server = setup_testserver();