2
0
Fork 0
mirror of https://git.asonix.dog/asonix/pict-rs synced 2024-12-22 19:31:35 +00:00

Add optional image format coersion

This commit is contained in:
asonix 2020-06-06 20:44:26 -05:00
parent 981e146696
commit 459db42a88
4 changed files with 95 additions and 10 deletions

View file

@ -1,7 +1,7 @@
use std::{net::SocketAddr, path::PathBuf}; use std::{net::SocketAddr, path::PathBuf};
#[derive(structopt::StructOpt)] #[derive(structopt::StructOpt)]
pub struct Config { pub(crate) struct Config {
#[structopt( #[structopt(
short, short,
long, long,
@ -11,6 +11,13 @@ pub struct Config {
#[structopt(short, long, help = "The path to the data directory, e.g. data/")] #[structopt(short, long, help = "The path to the data directory, e.g. data/")]
path: PathBuf, path: PathBuf,
#[structopt(
short,
long,
help = "An image format to convert all uploaded files into, supports 'jpg' and 'png'"
)]
format: Option<Format>,
} }
impl Config { impl Config {
@ -21,4 +28,46 @@ impl Config {
pub(crate) fn data_dir(&self) -> PathBuf { pub(crate) fn data_dir(&self) -> PathBuf {
self.path.clone() self.path.clone()
} }
pub(crate) fn format(&self) -> Option<Format> {
self.format.clone()
}
}
#[derive(Debug, thiserror::Error)]
#[error("Invalid format supplied, {0}")]
pub(crate) struct FormatError(String);
#[derive(Clone, Debug)]
pub(crate) enum Format {
Jpeg,
Png,
}
impl Format {
pub(crate) fn to_image_format(&self) -> image::ImageFormat {
match self {
Format::Jpeg => image::ImageFormat::Jpeg,
Format::Png => image::ImageFormat::Png,
}
}
pub(crate) fn to_mime(&self) -> mime::Mime {
match self {
Format::Jpeg => mime::IMAGE_JPEG,
Format::Png => mime::IMAGE_PNG,
}
}
}
impl std::str::FromStr for Format {
type Err = FormatError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"png" => Ok(Format::Png),
"jpg" => Ok(Format::Jpeg),
other => Err(FormatError(other.to_string())),
}
}
} }

View file

@ -40,6 +40,9 @@ pub enum UploadError {
#[error("Provided token did not match expected token")] #[error("Provided token did not match expected token")]
InvalidToken, InvalidToken,
#[error("Uploaded content could not be validated as an image")]
InvalidImage(image::error::ImageError),
} }
impl From<sled::transaction::TransactionError<UploadError>> for UploadError { impl From<sled::transaction::TransactionError<UploadError>> for UploadError {

View file

@ -239,7 +239,7 @@ async fn main() -> Result<(), anyhow::Error> {
let config = Config::from_args(); let config = Config::from_args();
std::env::set_var("RUST_LOG", "info"); std::env::set_var("RUST_LOG", "info");
env_logger::init(); env_logger::init();
let manager = UploadManager::new(config.data_dir()).await?; let manager = UploadManager::new(config.data_dir(), config.format()).await?;
// Create a new Multipart Form validator // Create a new Multipart Form validator
// //

View file

@ -1,4 +1,4 @@
use crate::{error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES}; use crate::{config::Format, error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES};
use actix_web::web; use actix_web::web;
use futures::stream::{Stream, StreamExt}; use futures::stream::{Stream, StreamExt};
use log::{error, warn}; use log::{error, warn};
@ -11,6 +11,7 @@ pub struct UploadManager {
} }
struct UploadManagerInner { struct UploadManagerInner {
format: Option<Format>,
hasher: sha2::Sha256, hasher: sha2::Sha256,
image_dir: PathBuf, image_dir: PathBuf,
alias_tree: sled::Tree, alias_tree: sled::Tree,
@ -40,7 +41,10 @@ impl UploadManager {
} }
/// Create a new UploadManager /// Create a new UploadManager
pub(crate) async fn new(mut root_dir: PathBuf) -> Result<Self, UploadError> { pub(crate) async fn new(
mut root_dir: PathBuf,
format: Option<Format>,
) -> Result<Self, UploadError> {
let mut sled_dir = root_dir.clone(); let mut sled_dir = root_dir.clone();
sled_dir.push("db"); sled_dir.push("db");
// sled automatically creates it's own directories // sled automatically creates it's own directories
@ -53,6 +57,7 @@ impl UploadManager {
Ok(UploadManager { Ok(UploadManager {
inner: Arc::new(UploadManagerInner { inner: Arc::new(UploadManagerInner {
format,
hasher: sha2::Sha256::new(), hasher: sha2::Sha256::new(),
image_dir: root_dir, image_dir: root_dir,
alias_tree: db.open_tree("alias")?, alias_tree: db.open_tree("alias")?,
@ -181,14 +186,42 @@ impl UploadManager {
return Err(UploadError::ContentType(content_type)); return Err(UploadError::ContentType(content_type));
} }
// -- READ IN BYTES FROM CLIENT -- let (img, format) = {
let mut bytes = bytes::BytesMut::new(); // -- READ IN BYTES FROM CLIENT --
let mut bytes = bytes::BytesMut::new();
while let Some(res) = stream.next().await { while let Some(res) = stream.next().await {
bytes.extend(res?); bytes.extend(res?);
} }
let bytes = bytes.freeze(); let bytes = bytes.freeze();
// -- VALIDATE IMAGE --
let format = image::guess_format(&bytes).map_err(UploadError::InvalidImage)?;
let img = image::load_from_memory(&bytes).map_err(UploadError::InvalidImage)?;
(img, format)
};
let format = self
.inner
.format
.as_ref()
.map(|f| f.to_image_format())
.unwrap_or(format);
let content_type = self
.inner
.format
.as_ref()
.map(|f| f.to_mime())
.unwrap_or(content_type);
let bytes: bytes::Bytes = {
let mut bytes = std::io::Cursor::new(vec![]);
img.write_to(&mut bytes, format)?;
bytes::Bytes::from(bytes.into_inner())
};
// -- DUPLICATE CHECKS -- // -- DUPLICATE CHECKS --