427 lines
16 KiB
Rust
427 lines
16 KiB
Rust
use std::time::Duration;
|
|
|
|
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};
|
|
|
|
const UPLOAD_READ_TIMEOUT: Duration = Duration::from_secs(30);
|
|
|
|
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."))?;
|
|
// 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"));
|
|
}
|
|
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, path) VALUES (?, ?)")
|
|
.bind(uploader.0)
|
|
.bind(&key)
|
|
.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::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;
|
|
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);
|
|
}
|
|
|
|
#[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::{
|
|
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, path COLLATE utf8mb4_unicode_ci FROM image_uploads")
|
|
.fetch_one(&pool)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(
|
|
row,
|
|
(
|
|
123456789012345678,
|
|
url.strip_prefix("https://images.example.com/").unwrap().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"));
|
|
}
|
|
}
|