Adjust migration

src/web/routes/dashboard/api/images.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use log::warn;
 use rand::{rngs::OsRng, RngCore};
 use rocket::{
@@ -12,6 +14,8 @@ use sqlx::{MySql, Pool};
 
 use crate::web::image_uploads::{prepare_image, ImageError, ImageUploads, MAX_UPLOAD_BYTES};
 
+const UPLOAD_READ_TIMEOUT: Duration = Duration::from_secs(30);
+
 pub fn mount(
     rocket: rocket::Rocket<rocket::Build>,
     uploads: ImageUploads,
@@ -65,11 +69,14 @@ pub async fn upload_image(
         .clone()
         .try_acquire_owned()
         .map_err(|_| error(Status::TooManyRequests, "Image uploads busy. Please try again."))?;
-    let bytes = data
-        .open(MAX_UPLOAD_BYTES.bytes())
-        .into_bytes()
-        .await
-        .map_err(|_| error(Status::BadRequest, "Could not read image"))?;
+    // Bound the entire body read so stalled or trickling uploads release their slot.
+    let bytes =
+        tokio::time::timeout(UPLOAD_READ_TIMEOUT, data.open(MAX_UPLOAD_BYTES.bytes()).into_bytes())
+            .await
+            .map_err(|_| {
+                error(Status::RequestTimeout, "Image upload timed out. Please try again.")
+            })?
+            .map_err(|_| error(Status::BadRequest, "Could not read image"))?;
     if !bytes.is_complete() {
         return Err(error(Status::PayloadTooLarge, "Images must be no larger than 8 MiB"));
     }
@@ -179,6 +186,11 @@ mod tests {
         Client::tracked(mount(rocket::build().manage(pool), uploads)).await.unwrap()
     }
 
+    #[rocket::get("/test-login")]
+    fn login(cookies: &CookieJar<'_>) {
+        cookies.add_private(("userid", "123"));
+    }
+
     #[rocket::async_test]
     async fn disabled_upload_route_is_not_available() {
         let client = client(ImageUploads(None)).await;
@@ -252,6 +264,94 @@ mod tests {
         assert_eq!(response.status(), Status::NotFound);
     }
 
+    #[rocket::async_test]
+    async fn stalled_uploads_time_out_and_release_global_slots() {
+        use rocket::{config::LogLevel, fairing::AdHoc, Config};
+        use tokio::{
+            io::{AsyncReadExt, AsyncWriteExt},
+            net::TcpStream,
+            sync::oneshot,
+            time::{sleep, timeout},
+        };
+
+        let uploads = service(false, "http://127.0.0.1:1");
+        let slots = uploads.0.as_ref().unwrap().slots.clone();
+        let pool = sqlx::mysql::MySqlPoolOptions::new()
+            .connect_lazy("mysql://unused:unused@127.0.0.1:1/unused")
+            .unwrap();
+        let config = Config {
+            address: "127.0.0.1".parse().unwrap(),
+            port: 0,
+            log_level: LogLevel::Off,
+            ..Config::debug_default()
+        };
+        let (send_port, receive_port) = oneshot::channel();
+        let rocket =
+            mount(rocket::custom(config).manage(pool).mount("/", rocket::routes![login]), uploads)
+                .attach(AdHoc::on_liftoff("test address", |rocket| {
+                    Box::pin(async move { send_port.send(rocket.config().port).unwrap() })
+                }))
+                .ignite()
+                .await
+                .unwrap();
+        let shutdown = rocket.shutdown();
+        let server = tokio::spawn(rocket.launch());
+        let port = timeout(Duration::from_secs(5), receive_port).await.unwrap().unwrap();
+        let base = format!("http://127.0.0.1:{port}");
+        let http =
+            reqwest::Client::builder().no_proxy().timeout(Duration::from_secs(5)).build().unwrap();
+        let response = http.get(format!("{base}/test-login")).send().await.unwrap();
+        let cookie = response.headers()["set-cookie"]
+            .to_str()
+            .unwrap()
+            .split(';')
+            .next()
+            .unwrap()
+            .to_string();
+        let mut stalled = Vec::new();
+        for _ in 0..2 {
+            let mut socket = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
+            socket.write_all(format!(
+                "POST /dashboard/api/images HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/octet-stream\r\nContent-Length: 1024\r\nCookie: {cookie}\r\nConnection: close\r\n\r\nincomplete image"
+            ).as_bytes()).await.unwrap();
+            stalled.push(socket);
+        }
+        timeout(Duration::from_secs(5), async {
+            while slots.available_permits() != 0 {
+                sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .unwrap();
+
+        let upload = || {
+            http.post(format!("{base}/dashboard/api/images"))
+                .header("cookie", &cookie)
+                .header("content-type", "application/octet-stream")
+                .body("not an image")
+                .send()
+        };
+        assert_eq!(upload().await.unwrap().status(), Status::TooManyRequests.code);
+        for mut socket in stalled {
+            let mut response = String::new();
+            timeout(
+                UPLOAD_READ_TIMEOUT + Duration::from_secs(5),
+                socket.read_to_string(&mut response),
+            )
+            .await
+            .unwrap()
+            .unwrap();
+            assert!(response.starts_with("HTTP/1.1 408"), "{response}");
+            assert!(response.contains("Image upload timed out"), "{response}");
+        }
+        assert_eq!(slots.available_permits(), 2);
+        assert_eq!(upload().await.unwrap().status(), Status::UnprocessableEntity.code);
+
+        drop(http);
+        shutdown.notify();
+        timeout(Duration::from_secs(5), server).await.unwrap().unwrap().unwrap();
+    }
+
     #[sqlx::test]
     async fn stores_moderated_image_and_tracks_uploader(pool: Pool<MySql>) {
         use tokio::{