Add support for S3-backed image storage
This commit is contained in:
jude committed
1 parent
7fcb2eb747
commit
8763a8e74a
55 files changed
+2611
-111
No files matched your search
@@ -0,0 +1,284 @@
|
||||
use std::{env, io::Cursor, sync::Arc};
|
||||
|
||||
use image::{ImageFormat, ImageReader, Limits};
|
||||
use log::{info, warn};
|
||||
use nsfw::model::{Classification, Metric};
|
||||
use s3::{creds::Credentials, Bucket, Region};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
pub const MAX_UPLOAD_BYTES: u64 = 8 * 1024 * 1024;
|
||||
const MAX_IMAGE_PIXELS: u64 = 16_000_000;
|
||||
const NSFW_THRESHOLD: f32 = 0.5;
|
||||
const MODEL_BYTES: &[u8] =
|
||||
include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/nsfw/model.onnx"));
|
||||
|
||||
#[cfg(test)]
|
||||
mod fixture_tests;
|
||||
|
||||
// An empty configuration keeps the dashboard's existing URL-only behavior.
|
||||
pub struct ImageUploads(pub Option<Arc<ImageUploadService>>);
|
||||
|
||||
pub struct ImageUploadService {
|
||||
pub bucket: Box<Bucket>,
|
||||
pub public_url: String,
|
||||
pub model: Option<nsfw::Model>,
|
||||
pub slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl ImageUploads {
|
||||
pub fn from_env() -> Self {
|
||||
let names =
|
||||
["S3_BUCKET", "S3_REGION", "S3_ACCESS_KEY_ID", "S3_SECRET_ACCESS_KEY", "S3_PUBLIC_URL"];
|
||||
let values = names.map(|name| env::var(name).unwrap_or_default());
|
||||
Self::from_settings(&values)
|
||||
}
|
||||
|
||||
fn from_settings(values: &[String; 5]) -> Self {
|
||||
if values.iter().all(|value| value.trim().is_empty()) {
|
||||
return Self(None);
|
||||
}
|
||||
if values.iter().any(|value| value.trim().is_empty()) {
|
||||
warn!("Image uploads disabled: all S3 upload settings must be supplied");
|
||||
return Self(None);
|
||||
}
|
||||
|
||||
let check_nsfw = nsfw_check_enabled(env::var("NSFW_CHECK_ENABLED").ok().as_deref());
|
||||
match ImageUploadService::new(values, check_nsfw) {
|
||||
Ok(service) => {
|
||||
info!("Dashboard image uploads enabled");
|
||||
Self(Some(Arc::new(service)))
|
||||
}
|
||||
Err(error) => {
|
||||
warn!("Image uploads disabled: {}", error);
|
||||
Self(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ImageUploadService {
|
||||
fn new(values: &[String; 5], check_nsfw: bool) -> Result<Self, Box<dyn std::error::Error>> {
|
||||
let url = reqwest::Url::parse(&values[4])?;
|
||||
let public_url = url.as_str().trim_end_matches('/').to_string();
|
||||
if !matches!(url.scheme(), "https" | "http")
|
||||
|| url.host_str().is_none()
|
||||
|| !url.username().is_empty()
|
||||
|| url.password().is_some()
|
||||
|| url.query().is_some()
|
||||
|| url.fragment().is_some()
|
||||
|| public_url.len() > 400
|
||||
{
|
||||
return Err("S3_PUBLIC_URL must be an HTTP(S) base URL of at most 400 bytes".into());
|
||||
}
|
||||
let endpoint = env::var("S3_ENDPOINT").ok().filter(|value| !value.trim().is_empty());
|
||||
let region = match &endpoint {
|
||||
Some(endpoint) => {
|
||||
Region::Custom { region: values[1].clone(), endpoint: endpoint.clone() }
|
||||
}
|
||||
None => values[1].parse()?,
|
||||
};
|
||||
let credentials = Credentials::new(Some(&values[2]), Some(&values[3]), None, None, None)?;
|
||||
let mut bucket = Bucket::new(&values[0], region, credentials)?;
|
||||
if endpoint.is_some() {
|
||||
bucket = bucket.with_path_style();
|
||||
}
|
||||
bucket.set_request_timeout(Some(std::time::Duration::from_secs(30)));
|
||||
let model = if check_nsfw {
|
||||
Some(nsfw::create_model(MODEL_BYTES)?)
|
||||
} else {
|
||||
info!("NSFW checks disabled for dashboard image uploads");
|
||||
None
|
||||
};
|
||||
|
||||
Ok(Self { bucket, public_url, model, slots: Arc::new(Semaphore::new(2)) })
|
||||
}
|
||||
}
|
||||
|
||||
fn nsfw_check_enabled(value: Option<&str>) -> bool {
|
||||
match value.map(str::trim).map(str::to_ascii_lowercase).as_deref() {
|
||||
None | Some("" | "true" | "1" | "on") => true,
|
||||
Some("false" | "0" | "off") => false,
|
||||
Some(_) => {
|
||||
warn!("Unrecognized NSFW_CHECK_ENABLED value; keeping NSFW checks enabled");
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub enum ImageError {
|
||||
InvalidImage,
|
||||
Nsfw,
|
||||
ModerationFailed,
|
||||
}
|
||||
|
||||
fn probable_nsfw(predictions: &[Classification]) -> Result<bool, ImageError> {
|
||||
if predictions.len() != 5
|
||||
|| predictions.iter().any(|prediction| {
|
||||
!prediction.score.is_finite() || !(0.0..=1.0).contains(&prediction.score)
|
||||
})
|
||||
{
|
||||
return Err(ImageError::ModerationFailed);
|
||||
}
|
||||
let score: f32 = predictions
|
||||
.iter()
|
||||
.filter(|prediction| {
|
||||
matches!(prediction.metric, Metric::Porn | Metric::Hentai | Metric::Sexy)
|
||||
})
|
||||
.map(|prediction| prediction.score)
|
||||
.sum();
|
||||
Ok(score >= NSFW_THRESHOLD)
|
||||
}
|
||||
|
||||
// Publish only the decoded pixels, which are moderated when NSFW checks are enabled.
|
||||
// Static PNGs strip metadata and prevent unchecked animation frames from being published.
|
||||
fn decode_image(bytes: &[u8]) -> Result<image::RgbaImage, ImageError> {
|
||||
let mut reader = ImageReader::new(Cursor::new(bytes))
|
||||
.with_guessed_format()
|
||||
.map_err(|_| ImageError::InvalidImage)?;
|
||||
if !matches!(
|
||||
reader.format(),
|
||||
Some(ImageFormat::Jpeg | ImageFormat::Png | ImageFormat::Gif | ImageFormat::WebP)
|
||||
) {
|
||||
return Err(ImageError::InvalidImage);
|
||||
}
|
||||
let mut limits = Limits::default();
|
||||
limits.max_image_width = Some(8192);
|
||||
limits.max_image_height = Some(8192);
|
||||
limits.max_alloc = Some(128 * 1024 * 1024);
|
||||
reader.limits(limits);
|
||||
let decoder = reader.into_decoder().map_err(|_| ImageError::InvalidImage)?;
|
||||
use image::ImageDecoder;
|
||||
let (width, height) = decoder.dimensions();
|
||||
if width == 0 || height == 0 || u64::from(width) * u64::from(height) > MAX_IMAGE_PIXELS {
|
||||
return Err(ImageError::InvalidImage);
|
||||
}
|
||||
image::DynamicImage::from_decoder(decoder)
|
||||
.map(|image| image.to_rgba8())
|
||||
.map_err(|_| ImageError::InvalidImage)
|
||||
}
|
||||
|
||||
pub fn prepare_image(model: Option<&nsfw::Model>, bytes: &[u8]) -> Result<Vec<u8>, ImageError> {
|
||||
let image = decode_image(bytes)?;
|
||||
// Flatten transparency so the classifier sees exactly the published image.
|
||||
let mut visible = image::RgbImage::new(image.width(), image.height());
|
||||
for (target, source) in visible.pixels_mut().zip(image.pixels()) {
|
||||
for channel in 0..3 {
|
||||
let alpha = u32::from(source[3]);
|
||||
target[channel] =
|
||||
((u32::from(source[channel]) * alpha + 255 * (255 - alpha)) / 255) as u8;
|
||||
}
|
||||
}
|
||||
let visible = image::DynamicImage::ImageRgb8(visible);
|
||||
if let Some(model) = model {
|
||||
let predictions = nsfw::examine(model, &visible.to_rgba8()).map_err(|error| {
|
||||
warn!("Image moderation failed: {}", error);
|
||||
ImageError::ModerationFailed
|
||||
})?;
|
||||
if probable_nsfw(&predictions)? {
|
||||
return Err(ImageError::Nsfw);
|
||||
}
|
||||
}
|
||||
let mut output = Cursor::new(Vec::new());
|
||||
visible.write_to(&mut output, ImageFormat::Png).map_err(|_| ImageError::InvalidImage)?;
|
||||
Ok(output.into_inner())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn predictions(scores: [f32; 5]) -> Vec<Classification> {
|
||||
[Metric::Drawings, Metric::Hentai, Metric::Neutral, Metric::Porn, Metric::Sexy]
|
||||
.into_iter()
|
||||
.zip(scores)
|
||||
.map(|(metric, score)| Classification { metric, score })
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn combines_unsafe_categories_and_rejects_invalid_scores() {
|
||||
assert_eq!(probable_nsfw(&predictions([0.1, 0.2, 0.4, 0.2, 0.1])), Ok(true));
|
||||
assert_eq!(probable_nsfw(&predictions([0.1, 0.1, 0.7, 0.05, 0.05])), Ok(false));
|
||||
assert_eq!(
|
||||
probable_nsfw(&predictions([0.0, 0.0, f32::NAN, 0.0, 0.0])),
|
||||
Err(ImageError::ModerationFailed)
|
||||
);
|
||||
assert_eq!(probable_nsfw(&[]), Err(ImageError::ModerationFailed));
|
||||
assert_eq!(
|
||||
probable_nsfw(&predictions([0.0, -0.5, 1.0, 0.0, 0.0])),
|
||||
Err(ImageError::ModerationFailed)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_or_partial_configuration_disables_uploads() {
|
||||
let mut settings = std::array::from_fn(|_| String::new());
|
||||
assert!(ImageUploads::from_settings(&settings).0.is_none());
|
||||
settings[0] = "images".into();
|
||||
assert!(ImageUploads::from_settings(&settings).0.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nsfw_check_defaults_to_enabled_and_accepts_explicit_toggles() {
|
||||
for value in [None, Some(""), Some("true"), Some("1"), Some("ON"), Some("invalid")] {
|
||||
assert!(nsfw_check_enabled(value));
|
||||
}
|
||||
for value in [Some("false"), Some("0"), Some("OFF"), Some(" false ")] {
|
||||
assert!(!nsfw_check_enabled(value));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabling_nsfw_checks_skips_model_initialization() {
|
||||
let settings = ["images", "eu-west-2", "test", "test", "https://images.example.com"]
|
||||
.map(str::to_string);
|
||||
assert!(ImageUploadService::new(&settings, false).unwrap().model.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_moderation_allows_flagged_images_but_still_validates_files() {
|
||||
let model = nsfw::create_model(
|
||||
&include_bytes!("../../tests/fixtures/image-upload-unsafe.onnx")[..],
|
||||
)
|
||||
.unwrap();
|
||||
let mut bytes = Cursor::new(Vec::new());
|
||||
image::DynamicImage::new_rgb8(2, 2).write_to(&mut bytes, ImageFormat::Png).unwrap();
|
||||
assert_eq!(prepare_image(Some(&model), bytes.get_ref()), Err(ImageError::Nsfw));
|
||||
let published = prepare_image(None, bytes.get_ref()).unwrap();
|
||||
assert_eq!(image::load_from_memory(&published).unwrap().width(), 2);
|
||||
assert_eq!(prepare_image(None, b"not an image"), Err(ImageError::InvalidImage));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bundled_model_loads_and_classifies_an_image() {
|
||||
let model = nsfw::create_model(MODEL_BYTES).unwrap();
|
||||
let image = image::RgbaImage::from_pixel(224, 224, image::Rgba([255, 255, 255, 255]));
|
||||
let predictions = nsfw::examine(&model, &image).unwrap();
|
||||
assert!(probable_nsfw(&predictions).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_images_and_truncated_images() {
|
||||
assert_eq!(decode_image(b"<svg></svg>"), Err(ImageError::InvalidImage));
|
||||
assert_eq!(decode_image(b"\x89PNG\r\n\x1a\n"), Err(ImageError::InvalidImage));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decodes_real_images_without_trusting_the_filename() {
|
||||
let image = image::DynamicImage::new_rgb8(2, 2);
|
||||
for format in [ImageFormat::Png, ImageFormat::Jpeg, ImageFormat::Gif, ImageFormat::WebP] {
|
||||
let mut bytes = Cursor::new(Vec::new());
|
||||
image.write_to(&mut bytes, format).unwrap();
|
||||
assert_eq!(decode_image(bytes.get_ref()).unwrap().dimensions(), (2, 2));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_excessive_image_dimensions() {
|
||||
let image = image::DynamicImage::new_rgb8(8193, 1);
|
||||
let mut bytes = Cursor::new(Vec::new());
|
||||
image.write_to(&mut bytes, ImageFormat::Png).unwrap();
|
||||
assert_eq!(decode_image(bytes.get_ref()), Err(ImageError::InvalidImage));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
fs,
|
||||
io::Cursor,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use image::{ImageFormat, RgbaImage};
|
||||
|
||||
use super::{prepare_image, ImageError, MODEL_BYTES};
|
||||
|
||||
const NSFW_IMAGES: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/nsfw");
|
||||
const NON_NSFW_IMAGES: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/non_nsfw");
|
||||
|
||||
fn fixture_pairs(directory: &Path) -> Vec<(PathBuf, PathBuf)> {
|
||||
let mut indices = BTreeSet::new();
|
||||
for entry in fs::read_dir(directory).expect("Could not read image fixture directory") {
|
||||
let entry = entry.expect("Could not read image fixture entry");
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|extension| extension.to_str()) != Some("png") {
|
||||
continue;
|
||||
}
|
||||
let name = path.file_name().unwrap().to_str().expect("Fixture name is not valid UTF-8");
|
||||
let index = name
|
||||
.strip_prefix("image_")
|
||||
.and_then(|name| name.strip_suffix("a.png").or_else(|| name.strip_suffix("b.png")))
|
||||
.filter(|index| index.parse::<u32>().is_ok_and(|index| index > 0))
|
||||
.unwrap_or_else(|| {
|
||||
panic!("Invalid fixture name {name}; expected image_1a.png/image_1b.png")
|
||||
});
|
||||
indices.insert(index.to_string());
|
||||
}
|
||||
|
||||
assert!(
|
||||
!indices.is_empty(),
|
||||
"No paired fixtures in {}. Add images with scripts/split_nsfw_test_image.py.",
|
||||
directory.display()
|
||||
);
|
||||
indices
|
||||
.into_iter()
|
||||
.map(|index| {
|
||||
let first = directory.join(format!("image_{index}a.png"));
|
||||
let second = directory.join(format!("image_{index}b.png"));
|
||||
assert!(
|
||||
first.is_file() && second.is_file(),
|
||||
"Both fixture parts must exist: {} and {}",
|
||||
first.display(),
|
||||
second.display()
|
||||
);
|
||||
(first, second)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn recombine_parts(first: &RgbaImage, second: &RgbaImage) -> Result<RgbaImage, &'static str> {
|
||||
if first.dimensions() != second.dimensions() {
|
||||
return Err("Fixture parts must have identical dimensions");
|
||||
}
|
||||
let pixels = first.as_raw().iter().zip(second.as_raw()).map(|(a, b)| a ^ b).collect();
|
||||
RgbaImage::from_raw(first.width(), first.height(), pixels).ok_or("Invalid fixture pixel data")
|
||||
}
|
||||
|
||||
fn check_directory(directory: &str, expected_nsfw: bool) {
|
||||
let pairs = fixture_pairs(Path::new(directory));
|
||||
let model = nsfw::create_model(MODEL_BYTES).expect("Could not load bundled NSFW model");
|
||||
for (first, second) in pairs {
|
||||
let a = image::open(&first).expect("Could not decode first fixture part").to_rgba8();
|
||||
let b = image::open(&second).expect("Could not decode second fixture part").to_rgba8();
|
||||
let reconstructed = recombine_parts(&a, &b).expect("Could not recombine fixture parts");
|
||||
let mut bytes = Cursor::new(Vec::new());
|
||||
reconstructed
|
||||
.write_to(&mut bytes, ImageFormat::Png)
|
||||
.expect("Could not encode reconstructed fixture");
|
||||
|
||||
// Use the production path, including flattening transparency and the same
|
||||
// combined category threshold. Never write the reconstructed image to disk.
|
||||
let result = prepare_image(Some(&model), bytes.get_ref());
|
||||
if expected_nsfw {
|
||||
assert_eq!(
|
||||
result.as_ref().err(),
|
||||
Some(&ImageError::Nsfw),
|
||||
"Expected {} to be blocked as NSFW, got {:?}",
|
||||
first.display(),
|
||||
result.as_ref().err()
|
||||
);
|
||||
} else {
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"Expected {} to be accepted, got {:?}",
|
||||
first.display(),
|
||||
result.as_ref().err()
|
||||
);
|
||||
}
|
||||
// The toggle must allow either category, while retaining image validation.
|
||||
assert!(
|
||||
prepare_image(None, bytes.get_ref()).is_ok(),
|
||||
"Expected {} to be accepted with moderation disabled",
|
||||
first.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Add local split fixtures to tests/fixtures/nsfw, then run with --ignored"]
|
||||
fn nsfw_directory_images_are_blocked() {
|
||||
check_directory(NSFW_IMAGES, true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Add local split fixtures to tests/fixtures/non_nsfw, then run with --ignored"]
|
||||
fn non_nsfw_directory_images_are_allowed() {
|
||||
check_directory(NON_NSFW_IMAGES, false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recombining_parts_restores_every_pixel_and_alpha_channel() {
|
||||
let original = [0, 17, 254, 255, 45, 99, 128, 0];
|
||||
let random = vec![18, 201, 50, 34, 166, 76, 11, 220];
|
||||
let masked = original.iter().zip(&random).map(|(pixel, mask)| pixel ^ mask).collect();
|
||||
let first = RgbaImage::from_raw(2, 1, random).unwrap();
|
||||
let second = RgbaImage::from_raw(2, 1, masked).unwrap();
|
||||
assert_eq!(recombine_parts(&first, &second).unwrap().as_raw(), &original);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_parts_with_different_dimensions() {
|
||||
assert_eq!(
|
||||
recombine_parts(&RgbaImage::new(1, 1), &RgbaImage::new(2, 1)),
|
||||
Err("Fixture parts must have identical dimensions")
|
||||
);
|
||||
}
|
||||
+7
-4
@@ -4,6 +4,7 @@ mod macros;
|
||||
mod catchers;
|
||||
mod fairings;
|
||||
mod guards;
|
||||
mod image_uploads;
|
||||
mod ip_blocking;
|
||||
mod routes;
|
||||
|
||||
@@ -124,7 +125,8 @@ pub async fn initialize(
|
||||
if Path::new("static").exists() { "static" } else { "/lib/reminder-rs/static" };
|
||||
info!("Using static path: {}", static_path);
|
||||
|
||||
rocket::build()
|
||||
let image_uploads = image_uploads::ImageUploads::from_env();
|
||||
let rocket = rocket::build()
|
||||
.attach(MetricProducer)
|
||||
.attach(Template::fairing())
|
||||
.register(
|
||||
@@ -214,9 +216,10 @@ pub async fn initialize(
|
||||
routes::dashboard::export::todos::export,
|
||||
routes::dashboard::export::todos::import,
|
||||
],
|
||||
)
|
||||
.launch()
|
||||
.await?;
|
||||
);
|
||||
|
||||
let rocket = routes::dashboard::api::images::mount(rocket, image_uploads);
|
||||
rocket.launch().await?;
|
||||
|
||||
warn!("Exiting rocket runtime");
|
||||
// distribute kill signal
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod guild;
|
||||
pub mod images;
|
||||
pub mod user;
|
||||
|
||||
use log::warn;
|
||||
|
||||
@@ -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()),
|
||||
|
||||
Reference in new issue
Block a user