Merge pull request #24 from KekOnTheWorld/main

Added error handling and format

Author: lbirkert

src/config.rs

@@ -1,3 +1,3 @@
 pub const EXTENSION_MAX_LENGTH: usize = 10;
 pub const STREAM_ID_LENGTH: usize = 64;
-pub const FILE_ID_LENGTH: usize = 7;
+pub const FILE_ID_LENGTH: usize = 7;

src/database.rs

@@ -1,6 +1,6 @@
 use diesel::pg::PgConnection;
 
-use diesel::r2d2::{ Pool, PooledConnection, ConnectionManager, PoolError };
+use diesel::r2d2::{ConnectionManager, Pool, PoolError, PooledConnection};
 
 pub type PgPool = Pool<ConnectionManager<PgConnection>>;
 pub type PgPooledConnection = PooledConnection<ConnectionManager<PgConnection>>;
@@ -12,4 +12,4 @@ fn init_pool(database_url: &str) -> Result<PgPool, PoolError> {
 
 pub fn establish_connection(database_url: String) -> PgPool {
     init_pool(&database_url).expect("Failed to create pool")
-}
+}

src/errors.rs

@@ -10,7 +10,7 @@ use std::fmt::{self, Display, Formatter};
 pub struct JsonError {
     err_type: JsonErrorType,
     field: String,
-    error: String
+    error: String,
 }
 
 #[derive(Debug)]
@@ -35,7 +35,11 @@ impl JsonErrorType {
 
 impl JsonError {
     pub fn new(err_type: JsonErrorType, field: String, error: String) -> Self {
-        Self { err_type, field, error }
+        Self {
+            err_type,
+            field,
+            error,
+        }
     }
 }
 
@@ -96,4 +100,4 @@ json_error_type!(FS_REMOVE, StatusCode::INTERNAL_SERVER_ERROR);
 
 json_error_type!(HASH_MATCH, StatusCode::BAD_REQUEST);
 
-json_error_type!(OVERFLOW, StatusCode::BAD_REQUEST);
+json_error_type!(OVERFLOW, StatusCode::BAD_REQUEST);

src/http.rs

@@ -3,11 +3,11 @@ use std::{io::Error, sync::Arc};
 use std::collections::HashMap;
 
 use actix_cors::Cors;
-use actix_web::{HttpServer, App, web};
+use actix_web::{web, App, HttpServer};
 use sha1::Sha1;
 use tokio::{fs::File, sync::Mutex};
 
-use crate::{database::PgPool, colors, routes};
+use crate::{colors, database::PgPool, routes};
 
 pub struct UploadState {
     pub map: Mutex<HashMap<String, UploadEntry>>,
@@ -22,13 +22,13 @@ pub struct UploadState {
     pub embed_color: String,
     pub download_url: String,
 
-    pub chunk_size: usize
+    pub chunk_size: usize,
 }
 
 pub struct UploadEntry {
     pub file: File,
     pub ext: String,
-    pub hasher: Sha1
+    pub hasher: Sha1,
 }
 
 pub async fn main(state: Arc<UploadState>, address: String, port: u16) -> Result<(), Error> {
@@ -81,4 +81,4 @@ pub async fn main(state: Arc<UploadState>, address: String, port: u16) -> Result
             return Err(error);
         }
     }
-}
+}

src/main.rs

@@ -12,18 +12,18 @@ use std::sync::Arc;
 
 use dotenv::dotenv;
 use http::UploadState;
-use tokio::sync::Mutex;
 use std::env;
+use tokio::sync::Mutex;
 
-pub mod database;
-pub mod schema;
-pub mod routes;
-pub mod util;
+pub mod colors;
 pub mod config;
-pub mod models;
+pub mod database;
 pub mod errors;
 pub mod http;
-pub mod colors;
+pub mod models;
+pub mod routes;
+pub mod schema;
+pub mod util;
 
 fn clean_tmp<'a>(tmp: &'a str) {
     fs::remove_dir_all(tmp).expect("Failed to remove temp directory!");
@@ -33,42 +33,38 @@ fn clean_tmp<'a>(tmp: &'a str) {
 #[tokio::main]
 async fn main() {
     dotenv().ok();
-    
-    let tmp_dir = env::var("tmp_dir")
-        .unwrap_or("tmp/".to_owned());
-        
-    let upload_dir = env::var("upload_dir")
-        .unwrap_or("upload/".to_owned());
 
-    let web_dir = env::var("web_dir")
-        .unwrap_or("web/".to_owned());
+    let tmp_dir = env::var("tmp_dir").unwrap_or("tmp/".to_owned());
+
+    let upload_dir = env::var("upload_dir").unwrap_or("upload/".to_owned());
 
-    let embed_description = env::var("embed_description")
-        .unwrap_or("No description availlable".to_owned());
+    let web_dir = env::var("web_dir").unwrap_or("web/".to_owned());
 
-    let embed_color = env::var("embed_color")
-        .unwrap_or("#ffffff".to_owned());
+    let embed_description =
+        env::var("embed_description").unwrap_or("No description availlable".to_owned());
 
-    let download_url = env::var("download_url")
-        .unwrap_or("http://example.com/".to_owned());
+    let embed_color = env::var("embed_color").unwrap_or("#ffffff".to_owned());
+
+    let download_url = env::var("download_url").unwrap_or("http://example.com/".to_owned());
 
     let port = env::var("port")
         .unwrap_or("6942".to_owned())
         .parse()
         .unwrap_or(6942);
 
-    let address = env::var("address")
-        .unwrap_or("0.0.0.0".to_owned());
+    let address = env::var("address").unwrap_or("0.0.0.0".to_owned());
 
     let chunk_size: usize = env::var("chunksize")
         .unwrap_or("2048".to_owned())
         .parse()
-        .unwrap_or(2048) * 1024;
+        .unwrap_or(2048)
+        * 1024;
 
     // Clean temp directory
     clean_tmp(tmp_dir.as_str());
 
-    let pool = database::establish_connection(env::var("DATABASE_URL").expect("Database url not set!"));
+    let pool =
+        database::establish_connection(env::var("DATABASE_URL").expect("Database url not set!"));
 
     let state = UploadState {
         map: Mutex::new(HashMap::new()),
@@ -79,8 +75,8 @@ async fn main() {
         embed_description,
         embed_color,
         download_url,
-        chunk_size
+        chunk_size,
     };
 
     let _result = http::main(Arc::new(state), address, port).await;
-}
+}

src/models/file.rs

@@ -1,16 +1,16 @@
-use diesel::{PgConnection, QueryResult, QueryDsl};
 use crate::diesel::RunQueryDsl;
+use diesel::{PgConnection, QueryDsl, QueryResult};
 
 use crate::schema::files;
 
 use crate::diesel::ExpressionMethods;
 
 #[derive(Queryable, Insertable)]
-#[table_name="files"]
+#[table_name = "files"]
 pub struct File {
     pub id: String,
     pub ext: String,
-    pub hash: String
+    pub hash: String,
 }
 
 impl File {
@@ -20,11 +20,10 @@ impl File {
             .execute(connection)
     }
 
-    pub fn find(id: String, connection: &PgConnection) -> Vec<File> {
+    pub fn find(id: String, connection: &PgConnection) -> QueryResult<Vec<File>> {
         files::table
             .filter(files::dsl::id.eq(id))
             .select((files::dsl::id, files::dsl::ext, files::dsl::hash))
             .load::<File>(connection)
-            .expect("Error while executing query!")
     }
-}
+}

src/routes/api/create.rs

@@ -1,25 +1,35 @@
 use std::sync::Arc;
 
-use actix_web::{web, Responder, post, Result};
-use sha1::{Sha1, Digest};
+use actix_web::{post, web, Responder, Result};
+use sha1::{Digest, Sha1};
 use tokio::fs::File;
 
-use crate::{http::{UploadState, UploadEntry}, util::{checker, random}, config};
-
+use crate::{
+    config,
+    http::{UploadEntry, UploadState},
+    util::{checker, random},
+};
 
 #[post("/api/c/{ext}")]
 pub async fn create(
     path: web::Path<(String,)>,
     state: web::Data<Arc<UploadState>>,
 ) -> Result<impl Responder> {
     let ext = path.into_inner().0;
-    checker::in_bounds("Length of extension ", ext.len(), 0, config::EXTENSION_MAX_LENGTH)?;
+    checker::in_bounds(
+        "Length of extension ",
+        ext.len(),
+        0,
+        config::EXTENSION_MAX_LENGTH,
+    )?;
 
     let map = &mut state.map.lock().await;
 
     let stream = random::random_b64(config::STREAM_ID_LENGTH);
-    
-    let file = File::create(state.tmp_dir.clone() + &stream).await.map_err(|e| crate::error!(FS_CREATE, FILE, "Error while creating file: {}", e))?;
+
+    let file = File::create(state.tmp_dir.clone() + &stream)
+        .await
+        .map_err(|e| crate::error!(FS_CREATE, FILE, "Error while creating file: {}", e))?;
 
     let hasher = Sha1::new();
 
@@ -29,7 +39,5 @@ pub async fn create(
 
     map.insert(stream.clone(), entry);
 
-    Ok(web::Json(json!({
-        "stream": stream
-    })))
+    Ok(web::Json(json!({ "stream": stream })))
 }

src/routes/api/download.rs

@@ -1,9 +1,20 @@
-use std::{sync::Arc, path::Path};
+use std::{path::Path, sync::Arc};
 
 use actix_files::NamedFile;
-use actix_web::{web, get, Result, Responder, http::header::{ContentDisposition, DispositionType, DispositionParam}};
-
-use crate::{http::UploadState, util::{checker, files}, models::file};
+use actix_web::{
+    get,
+    http::header::{ContentDisposition, DispositionParam, DispositionType},
+    web, Responder, Result,
+};
+
+use crate::{
+    http::UploadState,
+    models::file,
+    util::{
+        checker::{self, map_qres},
+        files,
+    },
+};
 
 #[get("/api/d/{id}/")]
 pub async fn download(
@@ -14,21 +25,25 @@ pub async fn download(
 
     let db_connection = &checker::get_con(&state.pool)?;
 
-    if let Some(entry) = file::File::find(id, &db_connection).into_iter().next() {
+    if let Some(entry) = map_qres(
+        file::File::find(id, &db_connection),
+        "Error while selecting files",
+    )?
+    .into_iter()
+    .next()
+    {
         let filename = files::get_filename(entry.hash.clone(), entry.ext);
 
         let named_file = NamedFile::open(Path::new(state.upload_dir.as_str()).join(entry.hash))
             .map_err(|e| crate::error!(FS_OPEN, FILE, "Error while opening file: {}", e))?;
 
         let content_disposition = ContentDisposition {
             disposition: DispositionType::Attachment,
-            parameters: vec![
-                DispositionParam::Filename(filename)
-            ],
+            parameters: vec![DispositionParam::Filename(filename)],
         };
 
         Ok(named_file.set_content_disposition(content_disposition))
     } else {
         Err(crate::error!(NOT_FOUND, ID, "File with id not found").into())
     }
-}
+}

src/routes/api/finish.rs

@@ -1,11 +1,18 @@
 use std::sync::Arc;
 
-use actix_web::{web, Responder, post, Result};
+use actix_web::{post, web, Responder, Result};
 use sha1::Digest;
 use tokio::fs;
 
-use crate::{http::UploadState, util::{checker::{self, map_qres}, random}, config, models::file};
-
+use crate::{
+    config,
+    http::UploadState,
+    models::file,
+    util::{
+        checker::{self, map_qres},
+        random,
+    },
+};
 
 #[post("/api/f/{stream}/{hash}")]
 pub async fn finish(
@@ -21,27 +28,27 @@ pub async fn finish(
         let file_path = state.tmp_dir.clone() + &stream;
 
         if file_hash.eq(&hash) {
-            fs::rename(file_path, state.upload_dir.clone() + &file_hash).await.map_err(|e| crate::error!(FS_RENAME, FILE, "Error while renaming file: {}", e))?;
+            fs::rename(file_path, state.upload_dir.clone() + &file_hash)
+                .await
+                .map_err(|e| crate::error!(FS_RENAME, FILE, "Error while renaming file: {}", e))?;
 
             let id = random::random_b64(config::FILE_ID_LENGTH);
 
             let db_connection = &checker::get_con(&state.pool)?;
 
             let new_file = file::File {
-                id: id.clone(), 
-                ext: entry.ext.clone(), 
-                hash: file_hash
+                id: id.clone(),
+                ext: entry.ext.clone(),
+                hash: file_hash,
             };
-            
+
             map_qres(new_file.create(db_connection), "Error while inserting file")?;
 
-            Ok(web::Json(json!({
-                "id": id
-            })))
+            Ok(web::Json(json!({ "id": id })))
         } else {
             Err(crate::error!(HASH_MATCH, HASH, "Hash doesn't match").into())
         }
     } else {
         Err(crate::error!(NOT_FOUND, STREAM, "Stream not found").into())
     }
-}
+}

src/routes/api/mod.rs

@@ -1,5 +1,5 @@
 pub mod create;
-pub mod upload;
-pub mod remove;
+pub mod download;
 pub mod finish;
-pub mod download;
+pub mod remove;
+pub mod upload;