From 3bed4de1273a638e39c372436ecd2e27b406a865 Mon Sep 17 00:00:00 2001 From: neri Date: Sun, 4 Apr 2021 14:30:31 +0200 Subject: [PATCH] extract more modules, fix FileKind enum names --- src/config.rs | 27 ++++ src/db.rs | 46 +++++++ src/download.rs | 102 +++++++++++++++ src/file_kind.rs | 12 +- src/main.rs | 322 +++-------------------------------------------- src/multipart.rs | 4 +- src/upload.rs | 136 ++++++++++++++++++++ 7 files changed, 333 insertions(+), 316 deletions(-) create mode 100644 src/config.rs create mode 100644 src/db.rs create mode 100644 src/download.rs create mode 100644 src/upload.rs diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..b8597af --- /dev/null +++ b/src/config.rs @@ -0,0 +1,27 @@ +use std::env; + +use async_std::{fs, path::PathBuf}; + +#[derive(Clone)] +pub struct Config { + pub files_dir: PathBuf, + pub max_file_size: Option, +} + +pub async fn get_config() -> Config { + let max_file_size = env::var("UPLOAD_MAX_BYTES") + .ok() + .and_then(|variable| variable.parse().ok()) + .unwrap_or(8 * 1024 * 1024); + let max_file_size = (max_file_size != 0).then(|| max_file_size); + + let files_dir = PathBuf::from(env::var("FILES_DIR").unwrap_or_else(|_| "./files".to_owned())); + fs::create_dir_all(&files_dir) + .await + .expect("could not create directory for storing files"); + + Config { + files_dir, + max_file_size, + } +} diff --git a/src/db.rs b/src/db.rs new file mode 100644 index 0000000..18e8eb1 --- /dev/null +++ b/src/db.rs @@ -0,0 +1,46 @@ +use sqlx::postgres::{PgPool, PgPoolOptions}; +use std::env; + +pub async fn setup_db() -> PgPool { + let conn_url = &get_db_url(); + log::info!("Using Connection string {}", conn_url); + + let pool = PgPoolOptions::new() + .max_connections(5) + .connect_timeout(std::time::Duration::from_secs(5)) + .connect(conn_url) + .await + .expect("could not create db pool"); + + for query in include_str!("../init-db.sql").split_inclusive(";") { + sqlx::query(query) + .execute(&pool) + .await + .expect("could not initialize database schema"); + } + + pool +} + +fn get_db_url() -> String { + if let Ok(database_url) = env::var("DATABASE_URL") { + return database_url; + } + + let auth = if let Ok(user) = env::var("DATABASE_USER") { + if let Ok(pass) = env::var("DATABASE_PASS") { + format!("{}:{}@", user, pass) + } else { + format!("{}@", user) + } + } else { + String::new() + }; + + format!( + "postgresql://{auth}{host}/{name}", + auth = auth, + host = env::var("DATABASE_HOST").unwrap_or_else(|_| "localhost".to_string()), + name = env::var("DATABASE_NAME").unwrap_or_else(|_| "datatrash".to_string()) + ) +} diff --git a/src/download.rs b/src/download.rs new file mode 100644 index 0000000..0812669 --- /dev/null +++ b/src/download.rs @@ -0,0 +1,102 @@ +use actix_files::NamedFile; +use actix_web::{ + error, + http::header::{Charset, ContentDisposition, DispositionParam, DispositionType, ExtendedValue}, + web, Error, HttpRequest, HttpResponse, +}; +use async_std::{fs, path::Path}; +use futures::TryStreamExt; +use mime::Mime; +use sqlx::{ + postgres::{PgPool, PgRow}, + Row, +}; + +use crate::config::Config; +use crate::deleter; + +const VIEW_HTML: &str = include_str!("../template/view.html"); + +pub async fn download( + req: HttpRequest, + db: web::Data, + config: web::Data, +) -> Result { + let id = req.match_info().query("id"); + let mut rows = + sqlx::query("SELECT file_id, file_name, delete_on_download from files WHERE file_id = $1") + .bind(id) + .fetch(db.as_ref()); + let row: PgRow = rows + .try_next() + .await + .map_err(|_| error::ErrorInternalServerError("could not run select statement"))? + .ok_or_else(|| error::ErrorNotFound("file does not exist or has expired"))?; + + let file_id: String = row.get("file_id"); + let file_name: String = row.get("file_name"); + let delete_on_download: bool = row.get("delete_on_download"); + let mut path = config.files_dir.clone(); + path.push(&file_id); + + let download = req.query_string().contains("dl"); + let (content_type, mut content_disposition) = get_content_types(&path, &file_name); + let response = if content_type.type_() == mime::TEXT && !download { + let content = fs::read_to_string(path).await.map_err(|_| { + error::ErrorInternalServerError("this file should be here but could not be found") + })?; + let encoded = htmlescape::encode_minimal(&content); + let view_html = VIEW_HTML.replace("{text}", &encoded); + let response = HttpResponse::Ok().content_type("text/html").body(view_html); + Ok(response) + } else { + if download { + content_disposition.disposition = DispositionType::Attachment; + } + let file = NamedFile::open(path) + .map_err(|_| { + error::ErrorInternalServerError("this file should be here but could not be found") + })? + .set_content_type(content_type) + .set_content_disposition(content_disposition); + file.into_response(&req) + }; + if delete_on_download { + deleter::delete_by_id(&db, &file_id, &config.files_dir) + .await + .map_err(|_| error::ErrorInternalServerError("could not delete file"))?; + } + response +} + +fn get_content_types(path: &Path, filename: &str) -> (Mime, ContentDisposition) { + let std_path = std::path::Path::new(path.as_os_str()); + let ct = tree_magic_mini::from_filepath(std_path) + .unwrap_or("application/octet-stream") + .parse::() + .expect("tree_magic_mini should not produce invalid mime"); + + let disposition = match ct.type_() { + mime::IMAGE | mime::TEXT | mime::AUDIO | mime::VIDEO => DispositionType::Inline, + _ => DispositionType::Attachment, + }; + + let cd = ContentDisposition { + disposition, + parameters: get_disposition_params(filename), + }; + + (ct, cd) +} + +fn get_disposition_params(filename: &str) -> Vec { + let mut parameters = vec![DispositionParam::Filename(filename.to_owned())]; + if !filename.is_ascii() { + parameters.push(DispositionParam::FilenameExt(ExtendedValue { + charset: Charset::Ext(String::from("UTF-8")), + language_tag: None, + value: filename.to_owned().into_bytes(), + })) + } + parameters +} diff --git a/src/file_kind.rs b/src/file_kind.rs index d3ce841..4b3b4f2 100644 --- a/src/file_kind.rs +++ b/src/file_kind.rs @@ -2,15 +2,15 @@ use std::{fmt::Display, str::FromStr}; #[derive(Debug, PartialEq)] pub(crate) enum FileKind { - TEXT, - BINARY, + Text, + Binary, } impl Display for FileKind { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - FileKind::TEXT => write!(f, "text"), - FileKind::BINARY => write!(f, "binary"), + FileKind::Text => write!(f, "text"), + FileKind::Binary => write!(f, "binary"), } } } @@ -19,8 +19,8 @@ impl FromStr for FileKind { type Err = String; fn from_str(s: &str) -> Result { match s.to_lowercase().as_str() { - "text" => Ok(FileKind::TEXT), - "binary" => Ok(FileKind::BINARY), + "text" => Ok(FileKind::Text), + "binary" => Ok(FileKind::Binary), _ => Err(format!("unknown kind {}", s)), } } diff --git a/src/main.rs b/src/main.rs index 2421d85..23ab8ce 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,324 +1,30 @@ +mod config; +mod db; mod deleter; +mod download; mod file_kind; mod multipart; +mod upload; -use actix_files::{Files, NamedFile}; -use actix_multipart::Multipart; -use actix_web::{ - error, - http::header::{Charset, ContentDisposition, DispositionParam, DispositionType, ExtendedValue}, - middleware, web, App, Error, HttpRequest, HttpResponse, HttpServer, -}; -use async_std::{ - channel::{self, Sender}, - fs, - path::{Path, PathBuf}, - task, -}; +use actix_files::Files; +use actix_web::{middleware, web, App, Error, HttpResponse, HttpServer}; +use async_std::{channel, task}; use env_logger::Env; -use file_kind::FileKind; -use futures::TryStreamExt; -use mime::Mime; -use multipart::UploadConfig; -use rand::prelude::SliceRandom; -use sqlx::{ - postgres::{PgPool, PgPoolOptions, PgRow}, - Row, -}; +use sqlx::postgres::PgPool; use std::env; -const INDEX_HTML: &str = include_str!("../template/index.html"); -const UPLOAD_HTML: &str = include_str!("../template/upload.html"); -const VIEW_HTML: &str = include_str!("../template/view.html"); - -const ID_CHARS: &[char] = &[ - 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'j', 'k', 'm', 'n', 'p', 'q', 'r', 's', 't', 'u', 'v', - 'w', 'x', 'y', 'z', '1', '2', '3', '4', '5', '6', '7', '8', '9', -]; - -async fn index(req: web::HttpRequest) -> Result { - let upload_url = format!("{}/upload", get_host_url(&req)); - let index_html = INDEX_HTML.replace("{upload_url}", upload_url.as_str()); - Ok(HttpResponse::Ok() - .content_type("text/html") - .body(index_html)) -} - -// multipart data -// required: either 'file' or 'text' -// optional: 'keep_for' default to 30 minutes -async fn upload( - req: web::HttpRequest, - payload: Multipart, - db: web::Data, - expiry_watch_sender: web::Data>, - config: web::Data, -) -> Result { - let file_id = gen_file_id(); - let mut filename = config.files_dir.clone(); - filename.push(&file_id); - - let parsed_multipart = - multipart::parse_multipart(payload, &file_id, &filename, config.max_file_size).await; - let UploadConfig { - original_name, - valid_till, - kind, - delete_on_download, - } = match parsed_multipart { - Ok(data) => data, - Err(err) => { - if filename.exists().await { - fs::remove_file(filename).await.map_err(|_| { - error::ErrorInternalServerError( - "could not parse multipart; could not remove file", - ) - })?; - } - return Err(err); - } - }; - - let db_insert = sqlx::query( - "INSERT INTO Files (file_id, file_name, valid_till, kind, delete_on_download) \ - VALUES ($1, $2, $3, $4, $5)", - ) - .bind(&file_id) - .bind(&original_name) - .bind(valid_till.naive_local()) - .bind(kind.to_string()) - .bind(delete_on_download) - .execute(db.as_ref()) - .await; - if db_insert.is_err() { - fs::remove_file(filename).await.map_err(|_| { - error::ErrorInternalServerError( - "could not insert file into database; could not remove file", - ) - })?; - return Err(error::ErrorInternalServerError( - "could not insert file into database", - )); - } - - log::info!( - "{} create new file {} (valid_till: {}, kind: {}, delete_on_download: {})", - req.connection_info().realip_remote_addr().unwrap_or("-"), - file_id, - valid_till, - kind, - delete_on_download - ); - - expiry_watch_sender.send(()).await.unwrap(); - - let redirect = if kind == FileKind::BINARY { - let encoded_name = urlencoding::encode(&original_name); - format!("/upload/{}/{}", file_id, encoded_name) - } else { - format!("/upload/{}", file_id) - }; - - let url = get_file_url(&req, &file_id, Some(&original_name)); - Ok(HttpResponse::SeeOther() - .header("location", redirect) - .body(format!("{}\n", url))) -} - -fn gen_file_id() -> String { - let mut rng = rand::thread_rng(); - let mut id = String::with_capacity(5); - for _ in 0..5 { - id.push(*ID_CHARS.choose(&mut rng).expect("ID_CHARS is not empty")); - } - id -} - -fn get_host_url(req: &web::HttpRequest) -> String { - let conn = req.connection_info(); - format!("{}://{}", conn.scheme(), conn.host()) -} - -fn get_file_url(req: &web::HttpRequest, id: &str, name: Option<&str>) -> String { - if let Some(name) = name { - let encoded_name = urlencoding::encode(name); - format!("{}/{}/{}", get_host_url(req), id, encoded_name) - } else { - format!("{}/{}", get_host_url(req), id) - } -} - -async fn uploaded(req: web::HttpRequest) -> Result { - let id = req.match_info().query("id"); - let name = req.match_info().get("name"); - let url = get_file_url(&req, id, name); - let upload_html = UPLOAD_HTML.replace("{url}", url.as_str()); - Ok(HttpResponse::Ok() - .content_type("text/html") - .body(upload_html)) -} - -async fn download( - req: HttpRequest, - db: web::Data, - config: web::Data, -) -> Result { - let id = req.match_info().query("id"); - let mut rows = - sqlx::query("SELECT file_id, file_name, delete_on_download from files WHERE file_id = $1") - .bind(id) - .fetch(db.as_ref()); - let row: PgRow = rows - .try_next() - .await - .map_err(|_| error::ErrorInternalServerError("could not run select statement"))? - .ok_or_else(|| error::ErrorNotFound("file does not exist or has expired"))?; - - let file_id: String = row.get("file_id"); - let file_name: String = row.get("file_name"); - let delete_on_download: bool = row.get("delete_on_download"); - let mut path = config.files_dir.clone(); - path.push(&file_id); - - let download = req.query_string().contains("dl"); - let (content_type, mut content_disposition) = get_content_types(&path, &file_name); - let response = if content_type.type_() == mime::TEXT && !download { - let content = fs::read_to_string(path).await.map_err(|_| { - error::ErrorInternalServerError("this file should be here but could not be found") - })?; - let encoded = htmlescape::encode_minimal(&content); - let view_html = VIEW_HTML.replace("{text}", &encoded); - let response = HttpResponse::Ok().content_type("text/html").body(view_html); - Ok(response) - } else { - if download { - content_disposition.disposition = DispositionType::Attachment; - } - let file = NamedFile::open(path) - .map_err(|_| { - error::ErrorInternalServerError("this file should be here but could not be found") - })? - .set_content_type(content_type) - .set_content_disposition(content_disposition); - file.into_response(&req) - }; - if delete_on_download { - deleter::delete_by_id(&db, &file_id, &config.files_dir) - .await - .map_err(|_| error::ErrorInternalServerError("could not delete file"))?; - } - response -} - -fn get_content_types(path: &Path, filename: &str) -> (Mime, ContentDisposition) { - let std_path = std::path::Path::new(path.as_os_str()); - let ct = tree_magic_mini::from_filepath(std_path) - .unwrap_or("application/octet-stream") - .parse::() - .expect("tree_magic_mini should not produce invalid mime"); - - let disposition = match ct.type_() { - mime::IMAGE | mime::TEXT | mime::AUDIO | mime::VIDEO => DispositionType::Inline, - _ => DispositionType::Attachment, - }; - - let cd = ContentDisposition { - disposition, - parameters: get_disposition_params(filename), - }; - - (ct, cd) -} - -fn get_disposition_params(filename: &str) -> Vec { - let mut parameters = vec![DispositionParam::Filename(filename.to_owned())]; - if !filename.is_ascii() { - parameters.push(DispositionParam::FilenameExt(ExtendedValue { - charset: Charset::Ext(String::from("UTF-8")), - language_tag: None, - value: filename.to_owned().into_bytes(), - })) - } - parameters -} - async fn not_found() -> Result { Ok(HttpResponse::NotFound() .content_type("text/plain") .body("not found")) } -fn get_db_url() -> String { - if let Ok(database_url) = env::var("DATABASE_URL") { - return database_url; - } - - let auth = if let Ok(user) = env::var("DATABASE_USER") { - if let Ok(pass) = env::var("DATABASE_PASS") { - format!("{}:{}@", user, pass) - } else { - format!("{}@", user) - } - } else { - String::new() - }; - - format!( - "postgresql://{auth}{host}/{name}", - auth = auth, - host = env::var("DATABASE_HOST").unwrap_or_else(|_| "localhost".to_string()), - name = env::var("DATABASE_NAME").unwrap_or_else(|_| "datatrash".to_string()) - ) -} - -async fn setup_db() -> PgPool { - let conn_url = &get_db_url(); - log::info!("Using Connection string {}", conn_url); - - let pool = PgPoolOptions::new() - .max_connections(5) - .connect_timeout(std::time::Duration::from_secs(5)) - .connect(conn_url) - .await - .expect("could not create db pool"); - - for query in include_str!("../init-db.sql").split_inclusive(";") { - sqlx::query(query) - .execute(&pool) - .await - .expect("could not initialize database schema"); - } - - pool -} - -#[derive(Clone)] -struct Config { - files_dir: PathBuf, - max_file_size: Option, -} - #[actix_web::main] async fn main() -> std::io::Result<()> { env_logger::Builder::from_env(Env::default().default_filter_or("info,sqlx=warn")).init(); - let pool: PgPool = setup_db().await; - let max_file_size = env::var("UPLOAD_MAX_BYTES") - .ok() - .and_then(|variable| variable.parse().ok()) - .unwrap_or(8 * 1024 * 1024); - let max_file_size = if max_file_size == 0 { - None - } else { - Some(max_file_size) - }; - let config = Config { - files_dir: PathBuf::from(env::var("FILES_DIR").unwrap_or_else(|_| "./files".to_owned())), - max_file_size, - }; - fs::create_dir_all(&config.files_dir) - .await - .expect("could not create directory for storing files"); + let pool: PgPool = db::setup_db().await; + let config = config::get_config().await; let (sender, receiver) = channel::bounded(8); log::info!("omnomnom"); @@ -340,16 +46,16 @@ async fn main() -> std::io::Result<()> { .app_data(db.clone()) .app_data(expiry_watch_sender.clone()) .data(config.clone()) - .service(web::resource("/").route(web::get().to(index))) - .service(web::resource("/upload").route(web::post().to(upload))) + .service(web::resource("/").route(web::get().to(upload::index))) + .service(web::resource("/upload").route(web::post().to(upload::upload))) .service( web::resource(["/upload/{id}", "/upload/{id}/{name}"]) - .route(web::get().to(uploaded)), + .route(web::get().to(upload::uploaded)), ) .service(Files::new("/static", "static").disable_content_disposition()) .service( web::resource(["/{id:[a-z0-9]{5}}", "/{id:[a-z0-9]{5}}/{name}"]) - .route(web::get().to(download)), + .route(web::get().to(download::download)), ) .default_service(web::route().to(not_found)) } diff --git a/src/multipart.rs b/src/multipart.rs index 9d2f09f..f354e2b 100644 --- a/src/multipart.rs +++ b/src/multipart.rs @@ -36,7 +36,7 @@ pub(crate) async fn parse_multipart( continue; } original_name = file_original_name; - kind = Some(FileKind::BINARY); + kind = Some(FileKind::Binary); let mut file = fs::File::create(&filename) .await .map_err(|_| error::ErrorInternalServerError("could not create file"))?; @@ -47,7 +47,7 @@ pub(crate) async fn parse_multipart( continue; } original_name = Some(format!("{}.txt", file_id)); - kind = Some(FileKind::TEXT); + kind = Some(FileKind::Text); let mut file = fs::File::create(&filename) .await .map_err(|_| error::ErrorInternalServerError("could not create file"))?; diff --git a/src/upload.rs b/src/upload.rs new file mode 100644 index 0000000..203a301 --- /dev/null +++ b/src/upload.rs @@ -0,0 +1,136 @@ +use crate::config::Config; +use crate::file_kind::FileKind; +use crate::multipart; +use crate::multipart::UploadConfig; +use actix_multipart::Multipart; +use actix_web::{error, web, Error, HttpResponse}; +use async_std::{channel::Sender, fs}; +use rand::prelude::SliceRandom; +use sqlx::postgres::PgPool; + +const INDEX_HTML: &str = include_str!("../template/index.html"); +const UPLOAD_HTML: &str = include_str!("../template/upload.html"); + +const ID_CHARS: &[char] = &[ + 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'j', 'k', 'm', 'n', 'p', 'q', 'r', 's', 't', 'u', 'v', + 'w', 'x', 'y', 'z', '1', '2', '3', '4', '5', '6', '7', '8', '9', +]; + +pub async fn index(req: web::HttpRequest) -> Result { + let upload_url = format!("{}/upload", get_host_url(&req)); + let index_html = INDEX_HTML.replace("{upload_url}", upload_url.as_str()); + Ok(HttpResponse::Ok() + .content_type("text/html") + .body(index_html)) +} + +pub async fn upload( + req: web::HttpRequest, + payload: Multipart, + db: web::Data, + expiry_watch_sender: web::Data>, + config: web::Data, +) -> Result { + let file_id = gen_file_id(); + let mut filename = config.files_dir.clone(); + filename.push(&file_id); + + let parsed_multipart = + multipart::parse_multipart(payload, &file_id, &filename, config.max_file_size).await; + let UploadConfig { + original_name, + valid_till, + kind, + delete_on_download, + } = match parsed_multipart { + Ok(data) => data, + Err(err) => { + if filename.exists().await { + fs::remove_file(filename).await.map_err(|_| { + error::ErrorInternalServerError( + "could not parse multipart; could not remove file", + ) + })?; + } + return Err(err); + } + }; + + let db_insert = sqlx::query( + "INSERT INTO Files (file_id, file_name, valid_till, kind, delete_on_download) \ + VALUES ($1, $2, $3, $4, $5)", + ) + .bind(&file_id) + .bind(&original_name) + .bind(valid_till.naive_local()) + .bind(kind.to_string()) + .bind(delete_on_download) + .execute(db.as_ref()) + .await; + if db_insert.is_err() { + fs::remove_file(filename).await.map_err(|_| { + error::ErrorInternalServerError( + "could not insert file into database; could not remove file", + ) + })?; + return Err(error::ErrorInternalServerError( + "could not insert file into database", + )); + } + + log::info!( + "{} create new file {} (valid_till: {}, kind: {}, delete_on_download: {})", + req.connection_info().realip_remote_addr().unwrap_or("-"), + file_id, + valid_till, + kind, + delete_on_download + ); + + expiry_watch_sender.send(()).await.unwrap(); + + let redirect = if kind == FileKind::Binary { + let encoded_name = urlencoding::encode(&original_name); + format!("/upload/{}/{}", file_id, encoded_name) + } else { + format!("/upload/{}", file_id) + }; + + let url = get_file_url(&req, &file_id, Some(&original_name)); + Ok(HttpResponse::SeeOther() + .header("location", redirect) + .body(format!("{}\n", url))) +} + +fn gen_file_id() -> String { + let mut rng = rand::thread_rng(); + let mut id = String::with_capacity(5); + for _ in 0..5 { + id.push(*ID_CHARS.choose(&mut rng).expect("ID_CHARS is not empty")); + } + id +} + +fn get_host_url(req: &web::HttpRequest) -> String { + let conn = req.connection_info(); + format!("{}://{}", conn.scheme(), conn.host()) +} + +fn get_file_url(req: &web::HttpRequest, id: &str, name: Option<&str>) -> String { + if let Some(name) = name { + let encoded_name = urlencoding::encode(name); + format!("{}/{}/{}", get_host_url(req), id, encoded_name) + } else { + format!("{}/{}", get_host_url(req), id) + } +} + +pub async fn uploaded(req: web::HttpRequest) -> Result { + let id = req.match_info().query("id"); + let name = req.match_info().get("name"); + let url = get_file_url(&req, id, name); + let upload_html = UPLOAD_HTML.replace("{url}", url.as_str()); + Ok(HttpResponse::Ok() + .content_type("text/html") + .body(upload_html)) +}