From b818a051bc93deea12bf04a2a658e929aba4e17e Mon Sep 17 00:00:00 2001 From: jude Date: Wed, 28 Oct 2026 19:56:12 +0000 Subject: [PATCH] Adjust migration --- migrations/20261011120000_image_uploads.sql | 48 ++++++--- src/image_upload/tests.rs | 15 ++- src/web/routes/dashboard/api/images.rs | 110 +++++++++++++++++++- 3 files changed, 151 insertions(+), 22 deletions(-) diff --git a/migrations/20261011120000_image_uploads.sql b/migrations/20261011120000_image_uploads.sql index 793b1f7..34356b8 100644 --- a/migrations/20261011120000_image_uploads.sql +++ b/migrations/20261011120000_image_uploads.sql @@ -10,36 +10,54 @@ CREATE TABLE image_uploads ( INDEX image_uploads_orphaned_at (orphaned_at) ); +-- Keep column additions separate so existing rows do not need to be rebuilt. ALTER TABLE reminders ADD COLUMN avatar_id BIGINT UNSIGNED DEFAULT NULL, - ADD INDEX reminders_avatar_id (avatar_id), - ADD FOREIGN KEY (avatar_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, ADD COLUMN embed_image_id BIGINT UNSIGNED DEFAULT NULL, - ADD INDEX reminders_embed_image_id (embed_image_id), - ADD FOREIGN KEY (embed_image_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, ADD COLUMN embed_thumbnail_id BIGINT UNSIGNED DEFAULT NULL, - ADD INDEX reminders_embed_thumbnail_id (embed_thumbnail_id), - ADD FOREIGN KEY (embed_thumbnail_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, ADD COLUMN embed_author_id BIGINT UNSIGNED DEFAULT NULL, - ADD INDEX reminders_embed_author_id (embed_author_id), - ADD FOREIGN KEY (embed_author_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, ADD COLUMN embed_footer_id BIGINT UNSIGNED DEFAULT NULL, - ADD INDEX reminders_embed_footer_id (embed_footer_id), - ADD FOREIGN KEY (embed_footer_id) REFERENCES image_uploads (id) ON DELETE RESTRICT; + ALGORITHM=INSTANT; ALTER TABLE reminder_template ADD COLUMN avatar_id BIGINT UNSIGNED DEFAULT NULL, + ADD COLUMN embed_image_id BIGINT UNSIGNED DEFAULT NULL, + ADD COLUMN embed_thumbnail_id BIGINT UNSIGNED DEFAULT NULL, + ADD COLUMN embed_author_id BIGINT UNSIGNED DEFAULT NULL, + ADD COLUMN embed_footer_id BIGINT UNSIGNED DEFAULT NULL, + ALGORITHM=INSTANT; + +-- Run before the application populates the new columns: existing values are NULL. +-- Disabling checks on this connection permits online foreign key creation. +SET @image_uploads_foreign_key_checks = @@SESSION.foreign_key_checks; +SET SESSION foreign_key_checks = 0; + +ALTER TABLE reminders + ADD INDEX reminders_avatar_id (avatar_id), + ADD FOREIGN KEY (avatar_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, + ADD INDEX reminders_embed_image_id (embed_image_id), + ADD FOREIGN KEY (embed_image_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, + ADD INDEX reminders_embed_thumbnail_id (embed_thumbnail_id), + ADD FOREIGN KEY (embed_thumbnail_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, + ADD INDEX reminders_embed_author_id (embed_author_id), + ADD FOREIGN KEY (embed_author_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, + ADD INDEX reminders_embed_footer_id (embed_footer_id), + ADD FOREIGN KEY (embed_footer_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, + ALGORITHM=INPLACE, + LOCK=NONE; + +ALTER TABLE reminder_template ADD INDEX reminder_template_avatar_id (avatar_id), ADD FOREIGN KEY (avatar_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, - ADD COLUMN embed_image_id BIGINT UNSIGNED DEFAULT NULL, ADD INDEX reminder_template_embed_image_id (embed_image_id), ADD FOREIGN KEY (embed_image_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, - ADD COLUMN embed_thumbnail_id BIGINT UNSIGNED DEFAULT NULL, ADD INDEX reminder_template_embed_thumbnail_id (embed_thumbnail_id), ADD FOREIGN KEY (embed_thumbnail_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, - ADD COLUMN embed_author_id BIGINT UNSIGNED DEFAULT NULL, ADD INDEX reminder_template_embed_author_id (embed_author_id), ADD FOREIGN KEY (embed_author_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, - ADD COLUMN embed_footer_id BIGINT UNSIGNED DEFAULT NULL, ADD INDEX reminder_template_embed_footer_id (embed_footer_id), - ADD FOREIGN KEY (embed_footer_id) REFERENCES image_uploads (id) ON DELETE RESTRICT; + ADD FOREIGN KEY (embed_footer_id) REFERENCES image_uploads (id) ON DELETE RESTRICT, + ALGORITHM=INPLACE, + LOCK=NONE; + +SET SESSION foreign_key_checks = @image_uploads_foreign_key_checks; diff --git a/src/image_upload/tests.rs b/src/image_upload/tests.rs index 8d9f689..2bd4d3c 100644 --- a/src/image_upload/tests.rs +++ b/src/image_upload/tests.rs @@ -72,10 +72,21 @@ async fn image_migration_preserves_external_urls_and_adds_indexed_references(poo sqlx::query(&sql).bind(external).execute(&pool).await.unwrap(); } } - sqlx::raw_sql(include_str!("../../migrations/20261011120000_image_uploads.sql")) - .execute(&pool) + let mut connection = pool.acquire().await.unwrap(); + let foreign_key_checks: i64 = sqlx::query_scalar("SELECT @@SESSION.foreign_key_checks") + .fetch_one(&mut *connection) .await .unwrap(); + sqlx::raw_sql(include_str!("../../migrations/20261011120000_image_uploads.sql")) + .execute(&mut *connection) + .await + .unwrap(); + let restored_checks: i64 = sqlx::query_scalar("SELECT @@SESSION.foreign_key_checks") + .fetch_one(&mut *connection) + .await + .unwrap(); + assert_eq!(restored_checks, foreign_key_checks); + drop(connection); let path = format!("uploads/{:064x}.png", 1); let upload = sqlx::query("INSERT INTO image_uploads (uploader, path) VALUES (123, ?)") .bind(&path) diff --git a/src/web/routes/dashboard/api/images.rs b/src/web/routes/dashboard/api/images.rs index 487b1e6..18a03b9 100644 --- a/src/web/routes/dashboard/api/images.rs +++ b/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, 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) { use tokio::{