diff --git a/Cargo.toml b/Cargo.toml index 8e283d4..5591125 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "pict-rs" description = "A simple image hosting service" -version = "0.1.2" +version = "0.1.3" authors = ["asonix "] license = "AGPL-3.0" readme = "README.md" diff --git a/README.md b/README.md index 862e493..aaa3d14 100644 --- a/README.md +++ b/README.md @@ -4,21 +4,26 @@ _a simple image hosting service_ ## Usage ### Running ``` -pict-rs 0.1.0 +pict-rs 0.1.3 USAGE: - pict-rs [OPTIONS] --addr --path + pict-rs [FLAGS] [OPTIONS] --path FLAGS: - -h, --help Prints help information - -V, --version Prints version information + -h, --help Prints help information + -s, --skip-validate-imports Whether to skip validating images uploaded via the internal import API + -V, --version Prints version information OPTIONS: - -a, --addr The address and port the server binds to, e.g. 127.0.0.1:80 - -f, --format An image format to convert all uploaded files into, supports 'jpg' and 'png' - -p, --path The path to the data directory, e.g. data/ - -w, --whitelist ... An optional list of filters to whitelist, supports 'identity', 'thumbnail', and - 'blur' + -a, --addr The address and port the server binds to. Default: 0.0.0.0:8080 [env: + PICTRS_ADDR=] [default: 0.0.0.0:8080] + -f, --format An optional image format to convert all uploaded files into, supports 'jpg' + and 'png' [env: PICTRS_FORMAT=] + -m, --max-file-size Specify the maximum allowed uploaded file size (in Megabytes) [env: + PICTRS_MAX_FILE_SIZE=] [default: 40] + -p, --path The path to the data directory, e.g. data/ [env: PICTRS_PATH=] + -w, --whitelist ... An optional list of filters to whitelist, supports 'identity', 'thumbnail', + and 'blur' [env: PICTRS_FILTER_WHITELIST=] ``` #### Example: @@ -80,6 +85,9 @@ pict-rs offers four endpoints: "msg": "ok" } ``` +- `POST /import` for uploading an image while preserving the filename. This should not be exposed to + the public internet, as it can cause naming conflicts with saved files. The upload format and + response format are the same as the `POST /image` endpoint. - `GET /image/download?url=...` Download an image from a remote server, returning the same JSON payload as the `POST` endpoint - `GET /image/{file}` for getting a full-resolution image. `file` here is the `file` key from the diff --git a/src/config.rs b/src/config.rs index 40a56f9..30434ac 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2,6 +2,13 @@ use std::{collections::HashSet, net::SocketAddr, path::PathBuf}; #[derive(Clone, Debug, structopt::StructOpt)] pub(crate) struct Config { + #[structopt( + short, + long, + help = "Whether to skip validating images uploaded via the internal import API" + )] + skip_validate_imports: bool, + #[structopt( short, long, @@ -12,10 +19,11 @@ pub(crate) struct Config { addr: SocketAddr, #[structopt( - short, - long, - env = "PICTRS_PATH", - help = "The path to the data directory, e.g. data/")] + short, + long, + env = "PICTRS_PATH", + help = "The path to the data directory, e.g. data/" + )] path: PathBuf, #[structopt( @@ -33,6 +41,15 @@ pub(crate) struct Config { help = "An optional list of filters to whitelist, supports 'identity', 'thumbnail', and 'blur'" )] whitelist: Option>, + + #[structopt( + short, + long, + env = "PICTRS_MAX_FILE_SIZE", + help = "Specify the maximum allowed uploaded file size (in Megabytes)", + default_value = "40" + )] + max_file_size: usize, } impl Config { @@ -53,6 +70,14 @@ impl Config { .as_ref() .map(|wl| wl.iter().cloned().collect()) } + + pub(crate) fn validate_imports(&self) -> bool { + !self.skip_validate_imports + } + + pub(crate) fn max_file_size(&self) -> usize { + self.max_file_size + } } #[derive(Debug, thiserror::Error)] diff --git a/src/error.rs b/src/error.rs index 406995f..d55c38a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -61,6 +61,9 @@ pub enum UploadError { #[error("Error converting Path to String")] Path, + + #[error("Tried to save an image with an already-taken name")] + DuplicateAlias, } impl From for UploadError { @@ -99,9 +102,10 @@ where impl ResponseError for UploadError { fn status_code(&self) -> StatusCode { match self { - UploadError::NoFiles | UploadError::ContentType(_) | UploadError::Upload(_) => { - StatusCode::BAD_REQUEST - } + UploadError::DuplicateAlias + | UploadError::NoFiles + | UploadError::ContentType(_) + | UploadError::Upload(_) => StatusCode::BAD_REQUEST, UploadError::MissingAlias | UploadError::MissingFilename => StatusCode::NOT_FOUND, UploadError::InvalidToken => StatusCode::FORBIDDEN, _ => StatusCode::INTERNAL_SERVER_ERROR, diff --git a/src/main.rs b/src/main.rs index 11aaeb3..a30c791 100644 --- a/src/main.rs +++ b/src/main.rs @@ -270,7 +270,7 @@ async fn main() -> Result<(), anyhow::Error> { let manager2 = manager.clone(); let form = Form::new() .max_files(10) - .max_file_size(40 * MEGABYTES) + .max_file_size(config.max_file_size() * MEGABYTES) .field( "images", Field::array(Field::file(move |_, _, stream| { @@ -286,6 +286,32 @@ async fn main() -> Result<(), anyhow::Error> { })), ); + // Create a new Multipart Form validator for internal imports + // + // This form is expecting a single array field, 'images' with at most 10 files in it + let validate_imports = config.validate_imports(); + let manager2 = manager.clone(); + let import_form = Form::new() + .max_files(10) + .max_file_size(config.max_file_size() * MEGABYTES) + .field( + "images", + Field::array(Field::file(move |filename, content_type, stream| { + let manager = manager2.clone(); + + async move { + manager + .import(filename, content_type, validate_imports, stream) + .await + .map(|alias| { + let mut path = PathBuf::new(); + path.push(alias); + Some(path) + }) + } + })), + ); + let config2 = config.clone(); HttpServer::new(move || { let client = Client::build() @@ -316,6 +342,11 @@ async fn main() -> Result<(), anyhow::Error> { ) .service(web::resource("/{tail:.*}").route(web::get().to(serve))), ) + .service( + web::resource("/import") + .wrap(import_form.clone()) + .route(web::post().to(upload)), + ) }) .bind(config.bind_address())? .run() diff --git a/src/upload_manager.rs b/src/upload_manager.rs index b471614..34defcf 100644 --- a/src/upload_manager.rs +++ b/src/upload_manager.rs @@ -68,6 +68,7 @@ impl UploadManager { }) } + /// Store the path to a generated image variant so we can easily clean it up later pub(crate) async fn store_variant(&self, path: PathBuf) -> Result<(), UploadError> { let filename = path .file_name() @@ -88,6 +89,7 @@ impl UploadManager { Ok(()) } + /// Delete the alias, and the file & variants if no more aliases exist pub(crate) async fn delete(&self, alias: String, token: String) -> Result<(), UploadError> { use sled::Transactional; let db = self.inner.db.clone(); @@ -159,6 +161,111 @@ impl UploadManager { Ok(()) } + /// Generate a delete token for an alias + pub(crate) async fn delete_token(&self, alias: String) -> Result { + use rand::distributions::{Alphanumeric, Distribution}; + let rng = rand::thread_rng(); + let s: String = Alphanumeric.sample_iter(rng).take(10).collect(); + let delete_token = s.clone(); + + let alias_tree = self.inner.alias_tree.clone(); + let key = delete_key(&alias); + let res = web::block(move || { + alias_tree.compare_and_swap( + key.as_bytes(), + None as Option, + Some(s.as_bytes()), + ) + }) + .await?; + + if let Err(sled::CompareAndSwapError { + current: Some(ivec), + .. + }) = res + { + let s = String::from_utf8(ivec.to_vec())?; + + return Ok(s); + } + + Ok(delete_token) + } + + /// Upload the file while preserving the filename, optionally validating the uploaded image + pub(crate) async fn import( + &self, + alias: String, + content_type: mime::Mime, + validate: bool, + stream: UploadStream, + ) -> Result + where + UploadError: From, + { + let bytes = read_stream(stream).await?; + + let (bytes, content_type) = if validate { + self.validate_image(bytes).await? + } else { + (bytes, content_type) + }; + + // -- DUPLICATE CHECKS -- + + // Cloning bytes is fine because it's actually a pointer + let hash = self.hash(bytes.clone()).await?; + + self.add_existing_alias(&hash, &alias).await?; + + self.save_upload(bytes, hash, content_type).await?; + + // Return alias to file + Ok(alias) + } + + /// Upload the file, discarding bytes if it's already present, or saving if it's new + pub(crate) async fn upload(&self, stream: UploadStream) -> Result + where + UploadError: From, + { + // -- READ IN BYTES FROM CLIENT -- + let bytes = read_stream(stream).await?; + + // -- VALIDATE IMAGE -- + let (bytes, content_type) = self.validate_image(bytes).await?; + + // -- DUPLICATE CHECKS -- + + // Cloning bytes is fine because it's actually a pointer + let hash = self.hash(bytes.clone()).await?; + + let alias = self.add_alias(&hash, content_type.clone()).await?; + + self.save_upload(bytes, hash, content_type).await?; + + // Return alias to file + Ok(alias) + } + + /// Fetch the real on-disk filename given an alias + pub(crate) async fn from_alias(&self, alias: String) -> Result { + let tree = self.inner.alias_tree.clone(); + let hash = web::block(move || tree.get(alias.as_bytes())) + .await? + .ok_or(UploadError::MissingAlias)?; + + let db = self.inner.db.clone(); + let filename = web::block(move || db.get(hash)) + .await? + .ok_or(UploadError::MissingFile)?; + + let filename = String::from_utf8(filename.to_vec())?; + + Ok(filename) + } + + // Find image variants and remove them from the DB and the disk async fn cleanup_files(&self, filename: sled::IVec) -> Result<(), UploadError> { let mut path = self.image_dir(); let fname = String::from_utf8(filename.to_vec())?; @@ -201,61 +308,41 @@ impl UploadManager { Ok(()) } - /// Generate a delete token for an alias - pub(crate) async fn delete_token(&self, alias: String) -> Result { - use rand::distributions::{Alphanumeric, Distribution}; - let rng = rand::thread_rng(); - let s: String = Alphanumeric.sample_iter(rng).take(10).collect(); - let delete_token = s.clone(); + // check duplicates & store image if new + async fn save_upload( + &self, + bytes: bytes::Bytes, + hash: Vec, + content_type: mime::Mime, + ) -> Result<(), UploadError> { + let (dup, name) = self.check_duplicate(hash, content_type).await?; - let alias_tree = self.inner.alias_tree.clone(); - let key = delete_key(&alias); - let res = web::block(move || { - alias_tree.compare_and_swap( - key.as_bytes(), - None as Option, - Some(s.as_bytes()), - ) - }) - .await?; - - if let Err(sled::CompareAndSwapError { - current: Some(ivec), - .. - }) = res - { - let s = String::from_utf8(ivec.to_vec())?; - - return Ok(s); + // bail early with alias to existing file if this is a duplicate + if dup.exists() { + return Ok(()); } - Ok(delete_token) + // -- WRITE NEW FILE -- + let mut real_path = self.image_dir(); + real_path.push(name); + + safe_save_file(real_path, bytes).await?; + + Ok(()) } - /// Upload the file, discarding bytes if it's already present, or saving if it's new - pub(crate) async fn upload(&self, mut stream: UploadStream) -> Result - where - UploadError: From, - { - let (img, format) = { - // -- READ IN BYTES FROM CLIENT -- - let mut bytes = bytes::BytesMut::new(); + // import & export image using the image crate + async fn validate_image( + &self, + bytes: bytes::Bytes, + ) -> Result<(bytes::Bytes, mime::Mime), UploadError> { + let (img, format) = web::block(move || { + let format = image::guess_format(&bytes).map_err(UploadError::InvalidImage)?; + let img = image::load_from_memory(&bytes).map_err(UploadError::InvalidImage)?; - while let Some(res) = stream.next().await { - bytes.extend(res?); - } - - let bytes = bytes.freeze(); - - // -- VALIDATE IMAGE -- - web::block(move || { - let format = image::guess_format(&bytes).map_err(UploadError::InvalidImage)?; - let img = image::load_from_memory(&bytes).map_err(UploadError::InvalidImage)?; - - Ok((img, format)) as Result<(image::DynamicImage, image::ImageFormat), UploadError> - }) - .await? - }; + Ok((img, format)) as Result<(image::DynamicImage, image::ImageFormat), UploadError> + }) + .await?; let (format, content_type) = self .inner @@ -275,43 +362,7 @@ impl UploadManager { }) .await?; - // -- DUPLICATE CHECKS -- - - // Cloning bytes is fine because it's actually a pointer - let hash = self.hash(bytes.clone()).await?; - - let alias = self.add_alias(&hash, content_type.clone()).await?; - let (dup, name) = self.check_duplicate(hash, content_type).await?; - - // bail early with alias to existing file if this is a duplicate - if dup.exists() { - return Ok(alias); - } - - // -- WRITE NEW FILE -- - let mut real_path = self.image_dir(); - real_path.push(name); - - safe_save_file(real_path, bytes).await?; - - // Return alias to file - Ok(alias) - } - - pub(crate) async fn from_alias(&self, alias: String) -> Result { - let tree = self.inner.alias_tree.clone(); - let hash = web::block(move || tree.get(alias.as_bytes())) - .await? - .ok_or(UploadError::MissingAlias)?; - - let db = self.inner.db.clone(); - let filename = web::block(move || db.get(hash)) - .await? - .ok_or(UploadError::MissingFile)?; - - let filename = String::from_utf8(filename.to_vec())?; - - Ok(filename) + Ok((bytes, content_type)) } // produce a sh256sum of the uploaded file @@ -387,6 +438,14 @@ impl UploadManager { } } + async fn add_existing_alias(&self, hash: &[u8], alias: &str) -> Result<(), UploadError> { + self.save_alias(hash, alias).await??; + + self.store_alias(hash, alias).await?; + + Ok(()) + } + // Add an alias to an existing file // // This will help if multiple 'users' upload the same file, and one of them wants to delete it @@ -397,6 +456,16 @@ impl UploadManager { ) -> Result { let alias = self.next_alias(hash, content_type).await?; + self.store_alias(hash, &alias).await?; + + Ok(alias) + } + + // Add a pre-defined alias to an existin file + // + // DANGER: this can cause BAD BAD BAD conflicts if the same alias is used for multiple files + async fn store_alias(&self, hash: &[u8], alias: &str) -> Result<(), UploadError> { + let alias = alias.to_string(); loop { let db = self.inner.db.clone(); let id = web::block(move || db.generate_id()).await?.to_string(); @@ -418,7 +487,7 @@ impl UploadManager { } } - Ok(alias) + Ok(()) } // Generate an alias to the file @@ -430,26 +499,54 @@ impl UploadManager { use rand::distributions::{Alphanumeric, Distribution}; let mut limit: usize = 10; let rng = rand::thread_rng(); - let hvec = hash.to_vec(); loop { let s: String = Alphanumeric.sample_iter(rng).take(limit).collect(); - let filename = file_name(s, content_type.clone()); + let alias = file_name(s, content_type.clone()); - let tree = self.inner.alias_tree.clone(); - let vec = hvec.clone(); - let filename2 = filename.clone(); - let res = web::block(move || { - tree.compare_and_swap(filename2.as_bytes(), None as Option, Some(vec)) - }) - .await?; + let res = self.save_alias(hash, &alias).await?; if res.is_ok() { - return Ok(filename); + return Ok(alias); } limit += 1; } } + + // Save an alias to the database + async fn save_alias( + &self, + hash: &[u8], + alias: &str, + ) -> Result, UploadError> { + let tree = self.inner.alias_tree.clone(); + let vec = hash.to_vec(); + let alias = alias.to_string(); + + let res = web::block(move || { + tree.compare_and_swap(alias.as_bytes(), None as Option, Some(vec)) + }) + .await?; + + if res.is_err() { + return Ok(Err(UploadError::DuplicateAlias)); + } + + return Ok(Ok(())); + } +} + +async fn read_stream(mut stream: UploadStream) -> Result +where + UploadError: From, +{ + let mut bytes = bytes::BytesMut::new(); + + while let Some(res) = stream.next().await { + bytes.extend(res?); + } + + Ok(bytes.freeze()) } async fn remove_path(path: sled::IVec) -> Result<(), UploadError> {