Added better error handling

+ Added image deleting

Cargo.lock

@@ -17,17 +17,6 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
-[[package]]
-name = "ahash"
-version = "0.7.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5a824f2aa7e75a0c98c5a504fceb80649e9c35265d44525b5f94de4771a395cd"
-dependencies = [
- "getrandom",
- "once_cell",
- "version_check",
-]
-
 [[package]]
 name = "aho-corasick"
 version = "1.1.2"
@@ -222,9 +211,9 @@ dependencies = [
 
 [[package]]
 name = "base64"
-version = "0.13.1"
+version = "0.21.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
+checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567"
 
 [[package]]
 name = "bitflags"
@@ -232,6 +221,15 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
+[[package]]
+name = "bitflags"
+version = "2.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "block-buffer"
 version = "0.10.4"
@@ -319,7 +317,7 @@ checksum = "a0610544180c38b88101fecf2dd634b174a62eef6946f84dfc6a7127512b381c"
 dependencies = [
  "ansi_term",
  "atty",
- "bitflags",
+ "bitflags 1.3.2",
  "strsim",
  "textwrap",
  "unicode-width",
@@ -334,11 +332,12 @@ checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7"
 
 [[package]]
 name = "config"
-version = "0.13.4"
+version = "0.14.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "23738e11972c7643e4ec947840fc463b6a571afcd3e735bdfce7d03c7a784aca"
+checksum = "7328b20597b53c2454f0b1919720c25c7339051c02b72b7e05409e00b14132be"
 dependencies = [
  "async-trait",
+ "convert_case",
  "json5",
  "lazy_static",
  "nom",
@@ -351,6 +350,35 @@ dependencies = [
  "yaml-rust",
 ]
 
+[[package]]
+name = "const-random"
+version = "0.1.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5aaf16c9c2c612020bcfd042e170f6e32de9b9d75adb5277cdbbd2e2c8c8299a"
+dependencies = [
+ "const-random-macro",
+]
+
+[[package]]
+name = "const-random-macro"
+version = "0.1.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e"
+dependencies = [
+ "getrandom",
+ "once_cell",
+ "tiny-keccak",
+]
+
+[[package]]
+name = "convert_case"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca"
+dependencies = [
+ "unicode-segmentation",
+]
+
 [[package]]
 name = "core-foundation-sys"
 version = "0.8.6"
@@ -390,6 +418,12 @@ version = "0.8.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345"
 
+[[package]]
+name = "crunchy"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
+
 [[package]]
 name = "crypto-common"
 version = "0.1.6"
@@ -412,9 +446,12 @@ dependencies = [
 
 [[package]]
 name = "dlv-list"
-version = "0.3.0"
+version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0688c2a7f92e427f44895cd63841bff7b29f8d7a1648b9e7e07a4a365b2e1257"
+checksum = "442039f5147480ba31067cb00ada1adae6892028e40e45fc5de7b7df6dcc1b5f"
+dependencies = [
+ "const-random",
+]
 
 [[package]]
 name = "encoding_rs"
@@ -575,12 +612,9 @@ dependencies = [
 
 [[package]]
 name = "hashbrown"
-version = "0.12.3"
+version = "0.13.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e"
-dependencies = [
- "ahash",
-]
 
 [[package]]
 name = "hashbrown"
@@ -964,12 +998,12 @@ checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
 
 [[package]]
 name = "ordered-multimap"
-version = "0.4.3"
+version = "0.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ccd746e37177e1711c20dd619a1620f34f5c8b569c53590a72dedd5344d8924a"
+checksum = "4ed8acf08e98e744e5384c8bc63ceb0364e68a6854187221c18df61c4797690e"
 dependencies = [
  "dlv-list",
- "hashbrown 0.12.3",
+ "hashbrown 0.13.2",
 ]
 
 [[package]]
@@ -1124,7 +1158,6 @@ dependencies = [
  "multer",
  "serde",
  "serde_json",
- "sha2",
  "structopt",
  "tokio",
  "tracing-subscriber",
@@ -1226,7 +1259,7 @@ version = "0.2.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a"
 dependencies = [
- "bitflags",
+ "bitflags 1.3.2",
 ]
 
 [[package]]
@@ -1260,20 +1293,21 @@ checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f"
 
 [[package]]
 name = "ron"
-version = "0.7.1"
+version = "0.8.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "88073939a61e5b7680558e6be56b419e208420c2adb92be54921fa6b72283f1a"
+checksum = "b91f7eff05f748767f183df4320a63d6936e9c6107d97c9e6bdd9784f4289c94"
 dependencies = [
  "base64",
- "bitflags",
+ "bitflags 2.4.2",
  "serde",
+ "serde_derive",
 ]
 
 [[package]]
 name = "rust-ini"
-version = "0.18.0"
+version = "0.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f6d5f2436026b4f6e79dc829837d467cc7e9a55ee40e750d716713540715a2df"
+checksum = "7e2a3bcec1f113553ef1c88aae6c020a369d03d55b58de9869a0908930385091"
 dependencies = [
  "cfg-if",
  "ordered-multimap",
@@ -1344,6 +1378,15 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "serde_spanned"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eb3622f419d1296904700073ea6cc23ad690adbd66f13ea683df73298736f0c1"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "serde_urlencoded"
 version = "0.7.1"
@@ -1526,6 +1569,15 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "tiny-keccak"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237"
+dependencies = [
+ "crunchy",
+]
+
 [[package]]
 name = "tinyvec"
 version = "1.6.0"
@@ -1585,11 +1637,36 @@ dependencies = [
 
 [[package]]
 name = "toml"
-version = "0.5.11"
+version = "0.8.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f4f7f0dd8d50a853a531c426359045b1998f04219d88799810762cd4ad314234"
+checksum = "c6a4b9e8023eb94392d3dca65d717c53abc5dad49c07cb65bb8fcd87115fa325"
 dependencies = [
  "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "toml_edit",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3550f4e9685620ac18a50ed434eb3aec30db8ba93b0287467bca5826ea25baf1"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.21.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a8534fd7f78b5405e860340ad6575217ce99f38d4d5c8f2442cb5ecb50090e1"
+dependencies = [
+ "indexmap",
+ "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "winnow",
 ]
 
 [[package]]
@@ -1969,6 +2046,15 @@ version = "0.52.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dff9641d1cd4be8d1a070daf9e3773c5f67e78b4d9d42263020c057706765c04"
 
+[[package]]
+name = "winnow"
+version = "0.5.36"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "818ce546a11a9986bc24f93d0cdf38a8a1a400f1473ea8c82e59f6e0ffab9249"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "yaml-rust"
 version = "0.4.5"

Cargo.toml

@@ -9,12 +9,11 @@ edition = "2021"
 j_db = { version = "0.1.0", registry = "jojo-dev" }
 axum-macros = "0.4.1"
 serde = "1.0.195"
-config = "0.13.4"
+config = "0.14.0"
 tracing-subscriber = "0.3.18"
 chrono = { version = "0.4.31", features = ["serde"] }
 chrono-tz = "0.8.5"
 async-trait = "0.1.77"
-sha2 = "0.10.8"
 hex = "0.4.3"
 url = "2.5.0"
 structopt = "0.3.26"

src/main.rs

@@ -6,14 +6,15 @@ mod storage_manager;
 
 use crate::config::PicOxConfig;
 use crate::model::album::Album;
-use crate::model::image::ImageData;
+use crate::model::image::{Image, ImageData};
 use crate::state::Context;
-use crate::storage_manager::StorageManager;
+use crate::storage_manager::{StorageManager, StoreError};
 use axum::extract::{Multipart, Path, Query, State};
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::routing::{get, post};
 use axum::{Json, Router};
+use axum_macros::FromRequest;
 use j_db::database::Database;
 use j_db::model::JdbModel;
 use log::info;
@@ -43,41 +44,100 @@ struct AlbumQuery {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct AddImage {
     pub album: AlbumQuery,
-    pub image_name: String,
+    pub tags: Vec<String>,
+}
+
+#[derive(FromRequest)]
+#[from_request(via(axum::Json), rejection(PicOxError))]
+struct Response<T>(T);
+
+impl<T> IntoResponse for Response<T>
+where
+    Json<T>: IntoResponse,
+{
+    fn into_response(self) -> axum::response::Response {
+        Json(self.0).into_response()
+    }
+}
+
+enum PicOxError {
+    StoreError(StoreError),
+    DbError(j_db::error::JDbError),
+}
+
+impl From<StoreError> for PicOxError {
+    fn from(value: StoreError) -> Self {
+        Self::StoreError(value)
+    }
+}
+
+impl From<j_db::error::JDbError> for PicOxError {
+    fn from(value: j_db::error::JDbError) -> Self {
+        Self::DbError(value)
+    }
+}
+
+#[derive(Serialize)]
+struct ErrorResponse {
+    pub message: String,
+}
+
+impl IntoResponse for PicOxError {
+    fn into_response(self) -> axum::response::Response {
+        let (status, message) = match self {
+            PicOxError::StoreError(err) => match err {
+                StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()),
+                StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()),
+                StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()),
+                StoreError::IOError(_) => (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    "IO Error Has Occurred!".to_string(),
+                ),
+            },
+            PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
+        };
+
+        (status, Response(ErrorResponse { message })).into_response()
+    }
 }
 
 async fn create_album(
     State(context): State<PicContext>,
     Json(album): Json<CreateAlbum>,
-) -> impl IntoResponse {
+) -> Result<Response<Album>, PicOxError> {
     info!("Creating new album '{}'", album.album_name);
 
     let new_album = Album::new(&album.album_name, Vec::new(), 0);
 
-    let new_album = context.db.insert(new_album).unwrap();
+    let new_album = context.db.insert(new_album)?;
 
-    (StatusCode::OK, Json(new_album))
+    Ok(Response(new_album))
 }
 
-async fn get_album(album_id: Path<u64>, State(context): State<PicContext>) -> impl IntoResponse {
+async fn get_album(
-    let album = context.db.get::<Album>(*album_id).unwrap();
+    album_id: Path<u64>,
+    State(context): State<PicContext>,
+) -> Result<Response<Album>, PicOxError> {
+    let album = context.db.get::<Album>(*album_id)?;
 
-    (StatusCode::OK, Json(album))
+    Ok(Response(album))
 }
 
 async fn add_image(
     State(context): State<PicContext>,
     mut img_data: Multipart,
-) -> impl IntoResponse {
+) -> Result<Response<Image>, PicOxError> {
     let mut data: Vec<u8> = Vec::new();
     let mut metadata: Option<AddImage> = None;
+    let mut file_name = None;
     while let Some(field) = img_data.next_field().await.unwrap() {
-        let field_name = field.name().clone();
+        let field_name = field.name();
         if let Some(field_name) = field_name {
             if field_name == "metadata" {
                 let metadata_json = field.text().await.unwrap();
                 metadata = Some(serde_json::from_str(&metadata_json).unwrap())
             } else if field_name == "img_data" {
+                file_name = Some(field.file_name().unwrap().to_string());
                 data.extend_from_slice(field.bytes().await.unwrap().as_ref());
             }
         }
@@ -92,25 +152,24 @@ async fn add_image(
             &context.db,
             None,
             ImageData::Bytes(data),
-            &metadata.unwrap().image_name,
+            &file_name.unwrap(),
             0,
         )
-        .await
+        .await?;
-        .unwrap();
 
     album.images.push(img.id().unwrap());
 
-    context.db.insert::<Album>(album).unwrap();
+    context.db.insert::<Album>(album)?;
 
-    StatusCode::OK
+    Ok(Response(img))
 }
 
 async fn query_album(
     album_query: Query<AlbumQuery>,
     State(context): State<PicContext>,
-) -> impl IntoResponse {
+) -> Result<Response<Option<Album>>, PicOxError> {
     let resp = Album::find_album_by_query(&context.db, album_query.0);
-    (StatusCode::OK, Json(resp))
+    Ok(Response(resp))
 }
 
 #[tokio::main]

src/model/image.rs

@@ -60,3 +60,12 @@ pub enum ImageData {
     Bytes(Vec<u8>),
     Link(String),
 }
+
+impl ImageData {
+    pub fn size(&self) -> usize {
+        match self {
+            ImageData::Bytes(data) => data.len(),
+            ImageData::Link(_) => 0,
+        }
+    }
+}

src/storage_manager/file_store.rs

@@ -3,7 +3,10 @@ use crate::storage_manager::{Store, StoreError};
 use async_trait::async_trait;
 use j_db::database::Database;
 use serde::{Deserialize, Serialize};
-use sha2::Digest;
+use std::collections::hash_map::DefaultHasher;
+use std::fs;
+use std::hash::Hasher;
+use std::os::unix::fs::MetadataExt;
 use std::path::PathBuf;
 use url::Url;
 
@@ -38,8 +41,9 @@ impl Store for FileStore {
             ImageData::Bytes(b) => b,
             ImageData::Link(_) => unimplemented!("No link support"),
         };
-        let hash = sha2::Sha256::digest(&img_data);
+        let mut hasher = DefaultHasher::new();
-        let disk_file_name = hex::encode(hash);
+        hasher.write(img_data.as_slice());
+        let disk_file_name = hex::encode(hasher.finish().to_be_bytes());
 
         let file = PathBuf::from(file_name);
         let ext = file.extension().unwrap().to_str().unwrap();
@@ -60,8 +64,15 @@ impl Store for FileStore {
         Ok((img_link, storage_location))
     }
 
-    async fn delete_img(&mut self, img: Image) -> StoreError {
+    async fn delete_img(&mut self, img: Image) -> Result<(), StoreError> {
-        todo!()
+        match img.storage_location {
+            StorageLocation::FileStore { path } => {
+                tokio::fs::remove_file(path).await?;
+            }
+            StorageLocation::Link => unimplemented!("No link support"),
+        }
+
+        Ok(())
     }
 
     fn max_image_size(&self) -> usize {
@@ -72,7 +83,21 @@ impl Store for FileStore {
         self.config.max_total_storage
     }
 
-    fn current_store_size(&self, _db: &Database) -> usize {
+    async fn current_store_size(&self, db: &Database) -> Result<usize, StoreError> {
-        0
+        let files: Vec<StorageLocation> = db
+            .filter(|_, _file: &Image| true)
+            .unwrap()
+            .map(|file: Image| file.storage_location.clone())
+            .collect();
+
+        let mut total_size = 0;
+        for file in files {
+            if let StorageLocation::FileStore { path } = file {
+                let md = fs::metadata(path)?;
+
+                total_size += md.size();
+            }
+        }
+        Ok(total_size as usize)
     }
 }

src/storage_manager/mod.rs

@@ -4,6 +4,7 @@ use axum::async_trait;
 use j_db::database::Database;
 use serde::{Deserialize, Serialize};
 use std::fmt::{Display, Formatter};
+use std::path::PathBuf;
 use url::Url;
 
 pub mod file_store;
@@ -11,7 +12,6 @@ pub mod file_store;
 #[derive(Debug)]
 pub enum StoreError {
     InvalidFile,
-    ImageNotFound,
     OutOfStorage,
     ImageTooBig,
     IOError(tokio::io::Error),
@@ -20,8 +20,7 @@ pub enum StoreError {
 impl Display for StoreError {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
-            StoreError::InvalidFile => write!(f, "Invalid file"),
+            StoreError::InvalidFile => write!(f, "Invalid file type"),
-            StoreError::ImageNotFound => write!(f, "Image not found"),
             StoreError::IOError(err) => write!(f, "IO Error: {}", err),
             StoreError::OutOfStorage => write!(f, "Underlying store full"),
             StoreError::ImageTooBig => write!(f, "Image too big for store"),
@@ -62,13 +61,13 @@ pub trait Store: Send {
         ))
     }
 
-    async fn delete_img(&mut self, img: Image) -> StoreError;
+    async fn delete_img(&mut self, img: Image) -> Result<(), StoreError>;
 
     fn max_image_size(&self) -> usize;
 
     fn max_total_storage(&self) -> usize;
 
-    fn current_store_size(&self, db: &Database) -> usize;
+    async fn current_store_size(&self, db: &Database) -> Result<usize, StoreError>;
 }
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize)]
@@ -80,6 +79,7 @@ pub enum StorageTypes {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct StorageManagerConfig {
     pub default_storage_method: StorageTypes,
+    pub allowed_types: Vec<String>,
     pub file_store_config: Option<FileStoreConfig>,
 }
 
@@ -92,26 +92,23 @@ pub struct StorageManager {
 
 impl StorageManager {
     pub fn new(storage_manager_config: StorageManagerConfig) -> Self {
-        let file_store =
+        let file_store = storage_manager_config
-            if let Some(file_store_config) = storage_manager_config.file_store_config.clone() {
+            .file_store_config
-                Some(FileStore::new(file_store_config))
+            .clone()
-            } else {
+            .map(FileStore::new);
-                None
-            };
-
         Self {
             config: storage_manager_config,
             file_store,
         }
     }
 
-    fn get_default_storage_manager(&mut self) -> Box<&mut dyn Store> {
+    fn get_default_storage_manager(&mut self) -> &mut dyn Store {
         self.get_storage_manager(self.config.default_storage_method)
     }
 
-    fn get_storage_manager(&mut self, storage_type: StorageTypes) -> Box<&mut dyn Store> {
+    fn get_storage_manager(&mut self, storage_type: StorageTypes) -> &mut dyn Store {
         match storage_type {
-            StorageTypes::FileStore => Box::new(self.file_store.as_mut().unwrap()),
+            StorageTypes::FileStore => self.file_store.as_mut().unwrap(),
             StorageTypes::LinkStore => {
                 unimplemented!()
             }
@@ -126,12 +123,33 @@ impl StorageManager {
         file_name: &str,
         created_by: u64,
     ) -> Result<Image, StoreError> {
+        let file_name_path = PathBuf::from(file_name);
+
+        if let Some(ext) = file_name_path.extension() {
+            let ext = ext.to_str().unwrap();
+            if !self.config.allowed_types.contains(&ext.to_string()) {
+                return Err(StoreError::InvalidFile);
+            }
+        } else {
+            return Err(StoreError::InvalidFile);
+        }
+
         let store_type = if let Some(store_type) = store {
             self.get_storage_manager(store_type)
         } else {
             self.get_default_storage_manager()
         };
 
+        if img_data.size() > store_type.max_image_size() {
+            return Err(StoreError::ImageTooBig);
+        }
+
+        let new_size = store_type.current_store_size(db).await? + img_data.size();
+
+        if new_size > store_type.max_total_storage() {
+            return Err(StoreError::OutOfStorage);
+        }
+
         let img = store_type
             .create_img(img_data, file_name, created_by)
             .await?;