A Simple duplicate-resolving image host

.gitignore

@@ -0,0 +1,2 @@
+/target
+/data

Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "pict-rs"
+version = "0.1.0"
+authors = ["asonix <asonix@asonix.dog>"]
+edition = "2018"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+actix-form-data = { git = "https://git.asonix.dog/Aardwolf/actix-form-data" }
+actix-fs = { git = "https://git.asonix.dog/asonix/actix-fs" }
+actix-rt = "1.1.1"
+actix-web = "3.0.0-alpha.2"
+anyhow = "1.0"
+bytes = "0.5"
+env_logger = "0.7"
+futures = "0.3.4"
+image = "0.23.4"
+log = "0.4"
+mime = "0.3.1"
+rand = "0.7.3"
+serde_json = "1.0"
+sha2 = "0.8.2"
+sled = "0.32.0-rc1"
+thiserror = "1.0"

client.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+
+# This script could be used for actix-web multipart example test
+# just start server and run client.py
+
+import asyncio
+import aiofiles
+import aiohttp
+
+file_name = 'test.png'
+url = 'http://localhost:8080/image'
+
+async def file_sender(file_name=None):
+    async with aiofiles.open(file_name, 'rb') as f:
+        chunk = await f.read(64*1024)
+        while chunk:
+            yield chunk
+            chunk = await f.read(64*1024)
+
+
+async def req():
+    async with aiohttp.ClientSession() as session:
+        data = aiohttp.FormData(quote_fields=False)
+        data.add_field("images[]", file_sender(file_name=file_name), filename="image1.png", content_type="image/png")
+        data.add_field("images[]", file_sender(file_name=file_name), filename="image2.png", content_type="image/png")
+        data.add_field("images[]", file_sender(file_name=file_name), filename="image3.png", content_type="image/png")
+
+        async with session.post(url, data=data) as resp:
+            text = await resp.text()
+            print(text)
+            assert 201 == resp.status
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(req())

src/main.rs

@@ -0,0 +1,437 @@
+use actix_form_data::{Field, Form, Value};
+use actix_web::{
+    guard,
+    http::{
+        header::{CacheControl, CacheDirective},
+        StatusCode,
+    },
+    middleware::Logger,
+    web, App, HttpResponse, HttpServer, ResponseError,
+};
+use futures::stream::{Stream, StreamExt, TryStreamExt};
+use log::{error, info};
+use sha2::Digest;
+use std::{path::PathBuf, pin::Pin, sync::Arc};
+
+const ACCEPTED_MIMES: &[mime::Mime] = &[
+    mime::IMAGE_BMP,
+    mime::IMAGE_GIF,
+    mime::IMAGE_JPEG,
+    mime::IMAGE_PNG,
+];
+
+const MEGABYTES: usize = 1024 * 1024;
+const HOURS: u32 = 60 * 60;
+
+#[derive(Clone)]
+struct UploadManager {
+    inner: Arc<UploadManagerInner>,
+}
+
+struct UploadManagerInner {
+    hasher: sha2::Sha256,
+    base_dir: PathBuf,
+    db: sled::Db,
+}
+
+#[derive(Debug, thiserror::Error)]
+enum UploadError {
+    #[error("Invalid content type provided, {0}")]
+    ContentType(mime::Mime),
+
+    #[error("Couln't upload file, {0}")]
+    Upload(String),
+
+    #[error("Couldn't save file, {0}")]
+    Save(#[from] actix_fs::Error),
+
+    #[error("Error in DB, {0}")]
+    Db(#[from] sled::Error),
+
+    #[error("Error parsing string, {0}")]
+    ParseString(#[from] std::string::FromUtf8Error),
+
+    #[error("Error processing image, {0}")]
+    Image(#[from] image::error::ImageError),
+
+    #[error("Panic in blocking operation")]
+    Canceled,
+
+    #[error("No files present in upload")]
+    NoFiles,
+
+    #[error("Uploaded image could not be served, extension is missing")]
+    MissingExtension,
+}
+
+impl From<actix_form_data::Error> for UploadError {
+    fn from(e: actix_form_data::Error) -> Self {
+        UploadError::Upload(e.to_string())
+    }
+}
+
+impl<T> From<actix_web::error::BlockingError<T>> for UploadError
+where
+    T: Into<UploadError> + std::fmt::Debug,
+{
+    fn from(e: actix_web::error::BlockingError<T>) -> Self {
+        match e {
+            actix_web::error::BlockingError::Error(e) => e.into(),
+            _ => UploadError::Canceled,
+        }
+    }
+}
+
+impl ResponseError for UploadError {
+    fn status_code(&self) -> StatusCode {
+        match self {
+            UploadError::ContentType(_) | UploadError::Upload(_) => StatusCode::BAD_REQUEST,
+            _ => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+
+    fn error_response(&self) -> HttpResponse {
+        HttpResponse::build(self.status_code()).json(serde_json::json!({ "msg": self.to_string() }))
+    }
+}
+
+type UploadStream = Pin<Box<dyn Stream<Item = Result<bytes::Bytes, actix_form_data::Error>>>>;
+
+enum Dup {
+    Exists,
+    New,
+}
+
+impl Dup {
+    fn exists(&self) -> bool {
+        match self {
+            Dup::Exists => true,
+            _ => false,
+        }
+    }
+}
+
+impl UploadManager {
+    async fn new(mut root_dir: PathBuf) -> Result<Self, UploadError> {
+        let mut sled_dir = root_dir.clone();
+        sled_dir.push("db");
+        let db = sled::open(sled_dir)?;
+
+        root_dir.push("files");
+
+        actix_fs::create_dir_all(root_dir.clone()).await?;
+
+        Ok(UploadManager {
+            inner: Arc::new(UploadManagerInner {
+                hasher: sha2::Sha256::new(),
+                base_dir: root_dir,
+                db,
+            }),
+        })
+    }
+
+    async fn upload(
+        &self,
+        _filename: String,
+        content_type: mime::Mime,
+        mut stream: UploadStream,
+    ) -> Result<Option<PathBuf>, UploadError> {
+        if ACCEPTED_MIMES.iter().all(|valid| *valid != content_type) {
+            return Err(UploadError::ContentType(content_type));
+        }
+
+        let mut bytes = bytes::BytesMut::new();
+
+        while let Some(res) = stream.next().await {
+            bytes.extend(res?);
+        }
+
+        let bytes = bytes.freeze();
+
+        let hash = self.hash(bytes.clone()).await?;
+
+        let (dup, path) = self.check_duplicate(hash, content_type).await?;
+
+        if dup.exists() {
+            return Ok(Some(path));
+        }
+
+        let file = actix_fs::file::create(path.clone()).await?;
+
+        if let Err(e) = actix_fs::file::write(file, bytes).await {
+            error!("Error saving file, {}", e);
+            actix_fs::remove_file(path).await?;
+            return Err(e.into());
+        }
+
+        Ok(Some(path))
+    }
+
+    async fn hash(&self, bytes: bytes::Bytes) -> Result<Vec<u8>, UploadError> {
+        let mut hasher = self.inner.hasher.clone();
+        let hash = web::block(move || {
+            hasher.input(&bytes);
+            Ok(hasher.result().to_vec()) as Result<_, UploadError>
+        })
+        .await?;
+
+        Ok(hash)
+    }
+
+    async fn check_duplicate(
+        &self,
+        hash: Vec<u8>,
+        content_type: mime::Mime,
+    ) -> Result<(Dup, PathBuf), UploadError> {
+        let mut path = self.inner.base_dir.clone();
+        let db = self.inner.db.clone();
+
+        let filename = self.next_file(content_type).await?;
+
+        let filename2 = filename.clone();
+        let res = web::block(move || {
+            db.compare_and_swap(hash, None as Option<sled::IVec>, Some(filename2.as_bytes()))
+        })
+        .await?;
+
+        if let Err(sled::CompareAndSwapError {
+            current: Some(ivec),
+            ..
+        }) = res
+        {
+            let name = String::from_utf8(ivec.to_vec())?;
+            path.push(name);
+
+            return Ok((Dup::Exists, path));
+        }
+
+        path.push(filename);
+
+        Ok((Dup::New, path))
+    }
+
+    async fn next_file(&self, content_type: mime::Mime) -> Result<String, UploadError> {
+        let base_dir = self.inner.base_dir.clone();
+        use rand::distributions::{Alphanumeric, Distribution};
+        let mut limit: usize = 10;
+        let rng = rand::thread_rng();
+        loop {
+            let mut path = base_dir.clone();
+            let s: String = Alphanumeric.sample_iter(rng).take(limit).collect();
+
+            let filename = format!("{}{}", s, to_ext(content_type.clone()));
+
+            path.push(filename.clone());
+
+            if let Err(e) = actix_fs::metadata(path).await {
+                if e.kind() == Some(std::io::ErrorKind::NotFound) {
+                    return Ok(filename);
+                }
+                return Err(e.into());
+            }
+
+            limit += 1;
+        }
+    }
+}
+
+fn to_ext(mime: mime::Mime) -> &'static str {
+    if mime == mime::IMAGE_PNG {
+        ".png"
+    } else if mime == mime::IMAGE_JPEG {
+        ".jpg"
+    } else if mime == mime::IMAGE_GIF {
+        ".gif"
+    } else {
+        ".bmp"
+    }
+}
+
+fn from_ext(ext: std::ffi::OsString) -> mime::Mime {
+    match ext.to_str() {
+        Some("png") => mime::IMAGE_PNG,
+        Some("jpg") => mime::IMAGE_JPEG,
+        Some("gif") => mime::IMAGE_GIF,
+        _ => mime::IMAGE_BMP,
+    }
+}
+
+async fn upload(value: Value) -> Result<HttpResponse, UploadError> {
+    let images = value
+        .map()
+        .and_then(|mut m| m.remove("images"))
+        .and_then(|images| images.array())
+        .ok_or(UploadError::NoFiles)?;
+
+    let mut files = Vec::new();
+    for image in images.into_iter().filter_map(|i| i.file()) {
+        if let Some(saved_as) = image
+            .saved_as
+            .as_ref()
+            .and_then(|s| s.file_name())
+            .and_then(|s| s.to_str())
+        {
+            info!("Uploaded {} as {:?}", image.filename, saved_as);
+            files.push(serde_json::json!({ "file": saved_as }));
+        }
+    }
+
+    Ok(HttpResponse::Created().json(serde_json::json!({ "msg": "ok", "files": files })))
+}
+
+async fn serve(
+    manager: web::Data<UploadManager>,
+    filename: web::Path<String>,
+) -> Result<HttpResponse, UploadError> {
+    let mut path = manager.inner.base_dir.clone();
+    path.push(filename.into_inner());
+    let ext = path
+        .extension()
+        .ok_or(UploadError::MissingExtension)?
+        .to_owned();
+    let ext = from_ext(ext);
+
+    let stream = actix_fs::read_to_stream(path).await?;
+
+    Ok(srv_response(stream, ext))
+}
+
+async fn serve_resized(
+    manager: web::Data<UploadManager>,
+    filename: web::Path<(u32, String)>,
+) -> Result<HttpResponse, UploadError> {
+    use image::GenericImageView;
+
+    let mut path = manager.inner.base_dir.clone();
+
+    let (size, name) = filename.into_inner();
+    path.push(size.to_string());
+    path.push(name.clone());
+    let ext = path
+        .extension()
+        .ok_or(UploadError::MissingExtension)?
+        .to_owned();
+    let ext = from_ext(ext);
+
+    if let Err(e) = actix_fs::metadata(path.clone()).await {
+        if e.kind() != Some(std::io::ErrorKind::NotFound) {
+            error!("Error looking up thumbnail, {}", e);
+            return Err(e.into());
+        }
+
+        let mut original_path = manager.inner.base_dir.clone();
+        original_path.push(name.clone());
+
+        let (img, format) = {
+            let bytes = actix_fs::read(original_path.clone()).await?;
+            let format = image::guess_format(&bytes)?;
+            let img = image::load_from_memory(&bytes)?;
+
+            (img, format)
+        };
+
+        if !img.in_bounds(size, size) {
+            // return original image if resize target is larger
+            drop(img);
+            let stream = actix_fs::read_to_stream(original_path).await?;
+            return Ok(srv_response(stream, ext));
+        }
+
+        let img_bytes: bytes::Bytes = web::block(move || {
+            let mut bytes = std::io::Cursor::new(vec![]);
+            img.thumbnail(size, size).write_to(&mut bytes, format)?;
+            Ok(bytes::Bytes::from(bytes.into_inner())) as Result<_, image::error::ImageError>
+        })
+        .await?;
+
+        let path2 = path.clone();
+        let img_bytes2 = img_bytes.clone();
+
+        actix_rt::spawn(async move {
+            if let Some(path) = path2.parent() {
+                if let Err(e) = actix_fs::create_dir_all(path.to_owned()).await {
+                    error!("Couldn't create directory for thumbnail, {}", e);
+                }
+            }
+
+            if let Err(e) = actix_fs::metadata(path2.clone()).await {
+                if e.kind() == Some(std::io::ErrorKind::NotFound) {
+                    if let Err(e) = actix_fs::write(path2, img_bytes2).await {
+                        error!("Error saving image, {}", e);
+                    }
+                } else {
+                    error!("Error checking image, {}", e);
+                }
+            }
+        });
+
+        return Ok(srv_response(
+            Box::pin(futures::stream::once(async {
+                Ok(img_bytes) as Result<_, UploadError>
+            })),
+            ext,
+        ));
+    }
+
+    let stream = actix_fs::read_to_stream(path).await?;
+
+    Ok(srv_response(stream, ext))
+}
+
+fn srv_response<S, E>(stream: S, ext: mime::Mime) -> HttpResponse
+where
+    S: Stream<Item = Result<bytes::Bytes, E>> + Unpin + 'static,
+    E: Into<UploadError>,
+{
+    HttpResponse::Ok()
+        .set(CacheControl(vec![
+            CacheDirective::Public,
+            CacheDirective::MaxAge(24 * HOURS),
+            CacheDirective::Extension("immutable".to_owned(), None),
+        ]))
+        .content_type(ext.to_string())
+        .streaming(stream.err_into())
+}
+
+#[actix_rt::main]
+async fn main() -> Result<(), anyhow::Error> {
+    env_logger::init();
+    let manager = UploadManager::new("data/".to_string().into()).await?;
+
+    let manager2 = manager.clone();
+    let form = Form::new()
+        .max_files(10)
+        .max_file_size(40 * MEGABYTES)
+        .field(
+            "images",
+            Field::array(Field::file(move |filename, content_type, stream| {
+                let manager = manager2.clone();
+
+                async move { manager.upload(filename, content_type, stream).await }
+            })),
+        );
+
+    HttpServer::new(move || {
+        App::new()
+            .wrap(Logger::default())
+            .data(manager.clone())
+            .service(
+                web::scope("/image")
+                    .service(
+                        web::resource("")
+                            .guard(guard::Post())
+                            .wrap(form.clone())
+                            .route(web::post().to(upload)),
+                    )
+                    .service(web::resource("/{filename}").route(web::get().to(serve)))
+                    .service(
+                        web::resource("/{size}/{filename}").route(web::get().to(serve_resized)),
+                    ),
+            )
+    })
+    .bind("127.0.0.1:8080")?
+    .run()
+    .await?;
+
+    Ok(())
+}