Add support for S3-backed image storage

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

@@ -0,0 +1,319 @@
+use log::warn;
+use rand::{rngs::OsRng, RngCore};
+use rocket::{
+    data::{Data, ToByteUnit},
+    http::{CookieJar, Status},
+    post,
+    request::{FromRequest, Outcome},
+    serde::json::{json, Json, Value},
+    Request, State,
+};
+use sqlx::{MySql, Pool};
+
+use crate::web::image_uploads::{prepare_image, ImageError, ImageUploads, MAX_UPLOAD_BYTES};
+
+pub fn mount(
+    rocket: rocket::Rocket<rocket::Build>,
+    uploads: ImageUploads,
+) -> rocket::Rocket<rocket::Build> {
+    let enabled = uploads.0.is_some();
+    let rocket = rocket.manage(uploads);
+    if enabled {
+        rocket.mount("/dashboard", rocket::routes![upload_image])
+    } else {
+        rocket
+    }
+}
+
+pub struct Uploader(u64);
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for Uploader {
+    type Error = ();
+
+    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
+        match uploader_id(request.cookies()) {
+            Some(id) => Outcome::Success(Self(id)),
+            None => Outcome::Error((Status::Unauthorized, ())),
+        }
+    }
+}
+
+fn uploader_id(cookies: &CookieJar<'_>) -> Option<u64> {
+    cookies.get_private("userid")?.value().parse().ok().filter(|id| *id != 0)
+}
+
+type UploadResult = Result<Json<Value>, (Status, Json<Value>)>;
+
+fn error(status: Status, message: &str) -> (Status, Json<Value>) {
+    (status, Json(json!({ "error": message })))
+}
+
+// Requiring a non-form content type prevents cross-site HTML form submissions.
+// Authorization is a request guard, so it runs before the body is read.
+#[post("/api/images", format = "application/octet-stream", data = "<data>")]
+pub async fn upload_image(
+    uploader: Uploader,
+    uploads: &State<ImageUploads>,
+    pool: &State<Pool<MySql>>,
+    data: Data<'_>,
+) -> UploadResult {
+    let service =
+        uploads.0.as_ref().ok_or_else(|| error(Status::NotFound, "Image uploads disabled"))?;
+    let _permit = service
+        .slots
+        .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"))?;
+    if !bytes.is_complete() {
+        return Err(error(Status::PayloadTooLarge, "Images must be no larger than 8 MiB"));
+    }
+    let service_for_moderation = service.clone();
+    let moderated = tokio::task::spawn_blocking(move || {
+        prepare_image(service_for_moderation.model.as_ref(), &bytes.into_inner())
+    })
+    .await
+    .map_err(|err| {
+        warn!("Image moderation task failed: {}", err);
+        error(Status::ServiceUnavailable, "Image moderation unavailable. Please try again.")
+    })?;
+    let image = moderated.map_err(|err| match err {
+        ImageError::InvalidImage => error(
+            Status::UnprocessableEntity,
+            "Please upload a valid PNG, JPEG, GIF or WebP image (at most 16 million pixels)",
+        ),
+        ImageError::Nsfw => error(
+            Status::UnavailableForLegalReasons,
+            "This image was flagged as probably NSFW and cannot be uploaded.",
+        ),
+        ImageError::ModerationFailed => {
+            error(Status::ServiceUnavailable, "Image moderation unavailable. Please try again.")
+        }
+    })?;
+
+    let mut random = [0u8; 32];
+    OsRng.fill_bytes(&mut random);
+    let key = format!(
+        "uploads/{}.png",
+        random.iter().map(|byte| format!("{:02x}", byte)).collect::<String>()
+    );
+    let url = format!("{}/{}", service.public_url, key);
+    let response =
+        service.bucket.put_object_with_content_type(&key, &image, "image/png").await.map_err(
+            |err| {
+                warn!("Could not upload image to S3: {}", err);
+                error(Status::BadGateway, "Could not store image. Please try again.")
+            },
+        )?;
+    if !(200..300).contains(&response.status_code()) {
+        warn!("S3 image upload returned {}", response.status_code());
+        return Err(error(Status::BadGateway, "Could not store image. Please try again."));
+    }
+    if let Err(err) = sqlx::query("INSERT INTO image_uploads (uploader, url) VALUES (?, ?)")
+        .bind(uploader.0)
+        .bind(&url)
+        .execute(pool.inner())
+        .await
+    {
+        warn!("Could not record image upload: {}", err);
+        match service.bucket.delete_object(&key).await {
+            Ok(response) if (200..300).contains(&response.status_code()) => {}
+            _ => warn!("Could not clean up untracked image object {}", key),
+        }
+        return Err(error(
+            Status::InternalServerError,
+            "Could not record image upload. Please try again.",
+        ));
+    }
+
+    Ok(Json(json!({ "url": url })))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{io::Cursor, sync::Arc};
+
+    use rocket::{http::ContentType, local::asynchronous::Client};
+    use s3::{creds::Credentials, Bucket, Region};
+    use tokio::sync::Semaphore;
+
+    use super::*;
+    use crate::web::image_uploads::ImageUploadService;
+
+    fn service(unsafe_image: bool, endpoint: &str) -> ImageUploads {
+        let model: &[u8] = if unsafe_image {
+            include_bytes!("../../../../../tests/fixtures/image-upload-unsafe.onnx")
+        } else {
+            include_bytes!("../../../../../tests/fixtures/image-upload-safe.onnx")
+        };
+        let bucket = Bucket::new(
+            "images",
+            Region::Custom { region: "eu-west-2".into(), endpoint: endpoint.into() },
+            Credentials::new(Some("test"), Some("test"), None, None, None).unwrap(),
+        )
+        .unwrap()
+        .with_path_style();
+        ImageUploads(Some(Arc::new(ImageUploadService {
+            bucket,
+            public_url: "https://images.example.com".into(),
+            model: Some(nsfw::create_model(model).unwrap()),
+            slots: Arc::new(Semaphore::new(2)),
+        })))
+    }
+
+    fn image_bytes() -> Vec<u8> {
+        let mut bytes = Cursor::new(Vec::new());
+        image::DynamicImage::new_rgb8(2, 2).write_to(&mut bytes, image::ImageFormat::Png).unwrap();
+        bytes.into_inner()
+    }
+
+    async fn client(uploads: ImageUploads) -> Client {
+        let pool = sqlx::mysql::MySqlPoolOptions::new()
+            .connect_lazy("mysql://unused:unused@127.0.0.1:1/unused")
+            .unwrap();
+        Client::tracked(mount(rocket::build().manage(pool), uploads)).await.unwrap()
+    }
+
+    #[rocket::async_test]
+    async fn disabled_upload_route_is_not_available() {
+        let client = client(ImageUploads(None)).await;
+        assert_eq!(
+            client
+                .post("/dashboard/api/images")
+                .header(ContentType::Binary)
+                .body(image_bytes())
+                .dispatch()
+                .await
+                .status(),
+            Status::NotFound
+        );
+    }
+
+    #[rocket::async_test]
+    async fn upload_requires_an_authenticated_private_cookie() {
+        let client = client(service(false, "http://127.0.0.1:1")).await;
+        for cookie in [None, Some("not-an-id"), Some("0")] {
+            let mut request = client
+                .post("/dashboard/api/images")
+                .header(ContentType::Binary)
+                .body(image_bytes());
+            if let Some(cookie) = cookie {
+                request = request.private_cookie(("userid", cookie));
+            }
+            assert_eq!(request.dispatch().await.status(), Status::Unauthorized);
+        }
+    }
+
+    #[rocket::async_test]
+    async fn unsafe_images_return_451_without_contacting_storage_or_database() {
+        let client = client(service(true, "http://127.0.0.1:1")).await;
+        let response = client
+            .post("/dashboard/api/images")
+            .header(ContentType::Binary)
+            .private_cookie(("userid", "123456789012345678"))
+            .body(image_bytes())
+            .dispatch()
+            .await;
+        assert_eq!(response.status(), Status::UnavailableForLegalReasons);
+        assert!(response.into_json::<Value>().await.unwrap()["error"]
+            .as_str()
+            .unwrap()
+            .contains("NSFW"));
+    }
+
+    #[rocket::async_test]
+    async fn rejects_invalid_oversized_and_form_uploads() {
+        let client = client(service(false, "http://127.0.0.1:1")).await;
+        for (body, status) in [
+            (b"this is not an image".to_vec(), Status::UnprocessableEntity),
+            (vec![0; MAX_UPLOAD_BYTES as usize + 1], Status::PayloadTooLarge),
+        ] {
+            let response = client
+                .post("/dashboard/api/images")
+                .header(ContentType::Binary)
+                .private_cookie(("userid", "123456789012345678"))
+                .body(body)
+                .dispatch()
+                .await;
+            assert_eq!(response.status(), status);
+        }
+        let response = client
+            .post("/dashboard/api/images")
+            .header(ContentType::Form)
+            .private_cookie(("userid", "123456789012345678"))
+            .body("file=anything")
+            .dispatch()
+            .await;
+        assert_eq!(response.status(), Status::NotFound);
+    }
+
+    #[sqlx::test]
+    async fn stores_moderated_image_and_tracks_uploader(pool: Pool<MySql>) {
+        use tokio::{
+            io::{AsyncReadExt, AsyncWriteExt},
+            net::TcpListener,
+        };
+
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let endpoint = format!("http://{}", listener.local_addr().unwrap());
+        let storage = tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.unwrap();
+            let mut received = Vec::new();
+            let mut buffer = [0; 4096];
+            let (headers, body) = loop {
+                let count = socket.read(&mut buffer).await.unwrap();
+                assert!(count > 0);
+                received.extend_from_slice(&buffer[..count]);
+                if let Some(end) = received.windows(4).position(|window| window == b"\r\n\r\n") {
+                    let headers = String::from_utf8(received[..end].to_vec()).unwrap();
+                    let length: usize = headers
+                        .lines()
+                        .find_map(|line| {
+                            line.to_lowercase()
+                                .strip_prefix("content-length:")
+                                .map(|length| length.trim().parse().unwrap())
+                        })
+                        .unwrap();
+                    if received.len() >= end + 4 + length {
+                        break (headers, received[end + 4..end + 4 + length].to_vec());
+                    }
+                }
+            };
+            socket
+                .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
+                .await
+                .unwrap();
+            (headers, body)
+        });
+        let client =
+            Client::tracked(mount(rocket::build().manage(pool.clone()), service(false, &endpoint)))
+                .await
+                .unwrap();
+        let response = client
+            .post("/dashboard/api/images")
+            .header(ContentType::Binary)
+            .private_cookie(("userid", "123456789012345678"))
+            .body(image_bytes())
+            .dispatch()
+            .await;
+        assert_eq!(response.status(), Status::Ok);
+        let response = response.into_json::<Value>().await.unwrap();
+        let url = response["url"].as_str().unwrap();
+        let row: (u64, String) = sqlx::query_as("SELECT uploader, url FROM image_uploads")
+            .fetch_one(&pool)
+            .await
+            .unwrap();
+        assert_eq!(row, (123456789012345678, url.to_string()));
+        let (headers, body) = storage.await.unwrap();
+        assert!(headers.starts_with("PUT /images/uploads/"));
+        assert!(headers.to_lowercase().contains("content-type: image/png"));
+        assert_eq!(image::load_from_memory(&body).unwrap().width(), 2);
+        assert!(url.starts_with("https://images.example.com/uploads/"));
+        assert!(url.ends_with(".png"));
+    }
+}

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

@@ -1,4 +1,5 @@
 pub mod guild;
+pub mod images;
 pub mod user;
 
 use log::warn;

src/web/routes/dashboard/api/user/mod.rs

@@ -23,6 +23,7 @@ use serenity::{
 use sqlx::{MySql, Pool};
 
 use crate::web::guards::transaction::Transaction;
+use crate::web::image_uploads::ImageUploads;
 
 #[derive(Serialize)]
 struct UserPreferences {
@@ -36,6 +37,7 @@ struct UserPreferences {
 struct UserInfo {
     name: String,
     patreon: bool,
+    image_uploads: bool,
     preferences: UserPreferences,
 }
 
@@ -52,10 +54,12 @@ pub async fn get_user_info(
     cookies: &CookieJar<'_>,
     ctx: &State<Context>,
     pool: &State<Pool<MySql>>,
+    uploads: &State<ImageUploads>,
 ) -> JsonValue {
     offline!(json!(UserInfo {
         name: "Discord".to_string(),
         patreon: true,
+        image_uploads: uploads.0.is_some(),
         preferences: UserPreferences {
             timezone: "UTC".to_string(),
             use_browser_timezone: true,
@@ -101,6 +105,7 @@ pub async fn get_user_info(
         );
 
         let user_info = UserInfo {
+            image_uploads: uploads.0.is_some(),
             name: cookies
                 .get_private("username")
                 .map_or("Discord User".to_string(), |c| c.value().to_string()),