[RFC PATCH proxmox-backup 2/2] fix #7756: api: stream file contents via dedicated thread

[Christian Ebner <c.ebner@proxmox.com>]

Reading the file contents to the streaming body is currently relying
on tokio's async reader implementation, which internally uses the
blocking thread pool.

This is especially problematic during initialization of a new backup
snapshot with pve-qemu [0], which has a 125 second timeout for the
QMP command 'backup', performing download of the previous snapshots
index files and initialization of the new index writers as part
thereof, which can be exceeded for large source disks and heavy load
on the PBS.

By opening and reading the file on the dedicated threads, the
blocking pool is bypassed, only used as fallback when the number of
parallel threads is exceeding the limit of 16, which was arbitrarily
chosen as upper boundary to not allow for denail of service attacks
by authenticated clients requesting a file without consuming the
response. Only one content chunk is queued for sending, thereby
creating an upper boundary for the memory used per file download.
Threads are stopped once the client disappears.

[0] https://git.proxmox.com/?p=pve-qemu.git;a=blob;f=debian/patches/pve/0027-PVE-Backup-Proxmox-backup-patches-for-QEMU.patch;h=67bb0c2f108ee526157174b5ca8722a77013a8db;hb=HEAD#l1425

Fixes: https://bugzilla.proxmox.com/show_bug.cgi?id=7756
Signed-off-by: Christian Ebner <c.ebner@proxmox.com>
---
 src/api2/helpers.rs | 165 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 153 insertions(+), 12 deletions(-)

diff --git a/src/api2/helpers.rs b/src/api2/helpers.rs
index 7d1fb8e35..df5331e5c 100644
--- a/src/api2/helpers.rs
+++ b/src/api2/helpers.rs
@@ -1,30 +1,171 @@
-use std::path::PathBuf;
+use std::fs::File;
+use std::io::Read;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 use anyhow::Error;
+use bytes::Bytes;
 use hyper::{Response, StatusCode, header};
+use tokio::sync::mpsc::Sender;
+use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::io::ReaderStream;
 
 use proxmox_http::Body;
 use proxmox_router::http_bail;
 
-const DOWNLOAD_BUFFER_SIZE: usize = 2 * 1024 * 1024;
+const DOWNLOAD_BUFFER_SIZE: u64 = 2 * 1024 * 1024;
+/// Maximum number of concurrent download from dedicated threads, fallback to
+/// tokio's blocking pool once reached.
+const MAX_DOWNLOAD_THREADS: usize = 16;
+static DOWNLOAD_THREADS: AtomicUsize = AtomicUsize::new(0);
 
 pub async fn create_download_response(path: PathBuf) -> Result<Response<Body>, Error> {
-    let file = match tokio::fs::File::open(path.clone()).await {
-        Ok(file) => file,
-        Err(ref err) if err.kind() == std::io::ErrorKind::NotFound => {
-            http_bail!(NOT_FOUND, "open file {:?} failed - not found", path);
-        }
-        Err(err) => http_bail!(BAD_REQUEST, "open file {:?} failed: {}", path, err),
+    let path_cloned = path.clone();
+    let permit = DOWNLOAD_THREADS
+        .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
+            (count < MAX_DOWNLOAD_THREADS).then_some(count + 1)
+        })
+        .is_ok();
+
+    let receivers = if permit {
+        let (open_tx, open_rx) = tokio::sync::oneshot::channel();
+        let (chunk_tx, chunk_rx) = tokio::sync::mpsc::channel(1);
+
+        std::thread::spawn(move || {
+            match File::open(path_cloned) {
+                Ok(file) => {
+                    if open_tx.send(Ok(())).is_ok() {
+                        send_file_chunked(file, chunk_tx);
+                    }
+                }
+                Err(err) => {
+                    let _ = open_tx.send(Err(err));
+                    DOWNLOAD_THREADS.fetch_sub(1, Ordering::AcqRel);
+                    return;
+                }
+            };
+            DOWNLOAD_THREADS.fetch_sub(1, Ordering::AcqRel);
+        });
+
+        Some((open_rx, chunk_rx))
+    } else {
+        None
     };
 
-    let payload = ReaderStream::with_capacity(file, DOWNLOAD_BUFFER_SIZE);
-    let body = Body::wrap_stream(payload);
+    let body = match receivers {
+        Some((open_rx, chunk_rx)) => {
+            // dedicated thread used for opening and reading file
+            let open_result = match open_rx.await {
+                Ok(res) => res,
+                Err(err) => http_bail!(INTERNAL_SERVER_ERROR, "download thread failed: {err}"),
+            };
+            check_open_result(&path, open_result)?;
+            Body::wrap_stream(ReceiverStream::new(chunk_rx))
+        }
+        None => {
+            // fallback to tokio's blocking pool to open and read file
+            let file = check_open_result(&path, tokio::fs::File::open(&path).await)?;
+            let stream = ReaderStream::with_capacity(file, DOWNLOAD_BUFFER_SIZE as usize);
+            Body::wrap_stream(stream)
+        }
+    };
 
     // fixme: set other headers ?
     Ok(Response::builder()
         .status(StatusCode::OK)
         .header(header::CONTENT_TYPE, "application/octet-stream")
-        .body(body)
-        .unwrap())
+        .body(body)?)
+}
+
+fn check_open_result<T>(path: &Path, result: std::io::Result<T>) -> Result<T, Error> {
+    match result {
+        Ok(value) => Ok(value),
+        Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
+            http_bail!(NOT_FOUND, "open file {path:?} failed - not found");
+        }
+        Err(err) => http_bail!(BAD_REQUEST, "open file {path:?} failed: {err}"),
+    }
+}
+
+/// Send the file contents via the sender in chunks of up to `DOWNLOAD_BUFFER_SIZE`.
+/// Stop early if the receiver was dropped because the client went away.
+fn send_file_chunked(file: File, chunk_tx: Sender<std::io::Result<Bytes>>) {
+    // optimize buffer size based on expected file size
+    let mut remaining = file
+        .metadata()
+        .map(|metadata| metadata.len())
+        .unwrap_or(DOWNLOAD_BUFFER_SIZE);
+
+    loop {
+        let capacity = remaining.clamp(1, DOWNLOAD_BUFFER_SIZE) as usize;
+        let mut chunk = Vec::with_capacity(capacity);
+
+        match (&file).take(DOWNLOAD_BUFFER_SIZE).read_to_end(&mut chunk) {
+            Ok(0) => return,
+            Ok(bytes_read) => {
+                remaining = remaining.saturating_sub(bytes_read as u64);
+                if chunk_tx.blocking_send(Ok(Bytes::from(chunk))).is_err() {
+                    return;
+                }
+            }
+            Err(err) => {
+                let _res = chunk_tx.blocking_send(Err(err));
+                return;
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::path::PathBuf;
+
+    use http_body_util::BodyExt;
+    use hyper::StatusCode;
+
+    use super::create_download_response;
+
+    fn test_file(name: &str, size: usize) -> PathBuf {
+        let path = std::env::temp_dir().join(name);
+        let data: Vec<u8> = (0..size).map(|i| (i % 256) as u8).collect();
+        std::fs::write(&path, data).expect("failed to write test file");
+        path
+    }
+
+    fn download(path: PathBuf) -> Result<Vec<u8>, anyhow::Error> {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()?;
+        runtime.block_on(async move {
+            let response = create_download_response(path).await?;
+            let body = response.into_body().collect().await?;
+            Ok(body.to_bytes().to_vec())
+        })
+    }
+
+    #[test]
+    fn download_file_contents() {
+        for (name, size) in [
+            ("empty", 0),
+            ("small", 100),
+            ("exact", 2 * 1024 * 1024),
+            ("multi", 5 * 1024 * 1024 + 3),
+        ] {
+            let path = test_file(name, size);
+            let expected = std::fs::read(&path).expect("failed to read test file");
+            let result = download(path.clone());
+            let _ = std::fs::remove_file(&path);
+            assert_eq!(result.expect("download failed"), expected, "file '{name}'");
+        }
+    }
+
+    #[test]
+    fn download_missing_file() {
+        let path = std::env::temp_dir().join("missing-file");
+        let err = download(path).expect_err("download of missing file succeeded");
+        let err = err
+            .downcast_ref::<proxmox_router::HttpError>()
+            .expect("expected HTTP error");
+        assert_eq!(err.code, StatusCode::NOT_FOUND);
+    }
 }
-- 
2.47.3