From 1d51c200d686872da1aea67fd746992a3079052e Mon Sep 17 00:00:00 2001 From: neri Date: Tue, 9 Mar 2021 23:36:24 +0100 Subject: [PATCH] do mime guessing, fix MAX_UPLOAD_BYTES --- Cargo.lock | 30 ++++++++++++++++ Cargo.toml | 2 ++ src/main.rs | 89 +++++++++++++++++++++++++++++++++--------------- src/multipart.rs | 26 +++++++++----- 4 files changed, 112 insertions(+), 35 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3411230..3898da1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -830,9 +830,11 @@ dependencies = [ "futures", "htmlescape", "log", + "mime", "openssl-sys", "rand 0.8.3", "sqlx", + "tree_magic_mini", "urlencoding", ] @@ -923,6 +925,12 @@ dependencies = [ "instant", ] +[[package]] +name = "fixedbitset" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37ab347416e802de484e4d03c7316c48f1ecb56574dfd4a46a80f173ce1de04d" + [[package]] name = "flate2" version = "1.0.20" @@ -1656,6 +1664,16 @@ version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4fd5641d01c8f18a23da7b6fe29298ff4b55afcccdf78973b24cf3175fee32e" +[[package]] +name = "petgraph" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "467d164a6de56270bd7c4d070df81d07beace25012d5103ced4e9ff08d6afdb7" +dependencies = [ + "fixedbitset", + "indexmap", +] + [[package]] name = "pin-project" version = "0.4.27" @@ -2474,6 +2492,18 @@ dependencies = [ "tracing", ] +[[package]] +name = "tree_magic_mini" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "287a760af78a7d607cb231500c06009ea1e0599cd57b10c32760061979bab1f1" +dependencies = [ + "fnv", + "lazy_static", + "nom 6.1.2", + "petgraph", +] + [[package]] name = "trust-dns-proto" version = "0.19.6" diff --git a/Cargo.toml b/Cargo.toml index 7cdf033..bd7c3f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,8 @@ chrono = "0.4.19" openssl-sys = "0.9.60" htmlescape = "0.3.1" urlencoding = "1.1.1" +tree_magic_mini = "1.0.1" +mime = "0.3.16" [features] vendored = ["openssl-sys/vendored"] diff --git a/src/main.rs b/src/main.rs index be46102..483ccdc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,19 +6,18 @@ use actix_files::{Files, NamedFile}; use actix_multipart::Multipart; use actix_web::{ error, - http::header::{ContentDisposition, DispositionParam, DispositionType}, - middleware, - web::{self, Bytes}, - App, Error, FromRequest, HttpRequest, HttpResponse, HttpServer, + http::header::{Charset, ContentDisposition, DispositionParam, DispositionType, ExtendedValue}, + middleware, web, App, Error, HttpRequest, HttpResponse, HttpServer, }; use async_std::{ channel::{self, Sender}, fs, - path::PathBuf, + path::{Path, PathBuf}, task, }; use file_kind::FileKind; use futures::TryStreamExt; +use mime::Mime; use rand::prelude::SliceRandom; use sqlx::{ postgres::{PgPool, PgPoolOptions, PgRow}, @@ -57,20 +56,21 @@ async fn upload( let mut filename = config.files_dir.clone(); filename.push(&file_id); - let (original_name, valid_till, kind) = - match multipart::parse_multipart(payload, &file_id, &filename).await { - 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 parsed_multipart = + multipart::parse_multipart(payload, &file_id, &filename, config.max_file_size).await; + let (original_name, valid_till, kind) = 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) VALUES ($1, $2, $3, $4)", @@ -177,18 +177,47 @@ async fn download( let response = HttpResponse::Ok().content_type("text/html").body(view_html); Ok(response) } else { + let (content_type, content_disposition) = get_content_types(&path, &file_name); let file = NamedFile::open(path) .map_err(|_| { error::ErrorInternalServerError("this file should be here but could not be found") })? - .set_content_disposition(ContentDisposition { - disposition: DispositionType::Attachment, - parameters: vec![DispositionParam::Filename(file_name)], - }); + .set_content_type(content_type) + .set_content_disposition(content_disposition); file.into_response(&req) } } +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 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(), + })) + } + + let cd = ContentDisposition { + disposition, + parameters, + }; + + (ct, cd) +} + async fn not_found() -> Result { Ok(HttpResponse::NotFound() .content_type("text/plain") @@ -240,6 +269,7 @@ async fn setup_db() -> PgPool { #[derive(Clone)] struct Config { files_dir: PathBuf, + max_file_size: Option, } #[actix_web::main] @@ -250,8 +280,18 @@ async fn main() -> std::io::Result<()> { env_logger::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 @@ -268,10 +308,6 @@ async fn main() -> std::io::Result<()> { let db = web::Data::new(pool); let expiry_watch_sender = web::Data::new(sender); - let upload_max_bytes: usize = env::var("UPLOAD_MAX_BYTES") - .ok() - .and_then(|variable| variable.parse().ok()) - .unwrap_or(8 * 1024 * 1024); let bind_address = env::var("BIND_ADDRESS").unwrap_or_else(|_| "0.0.0.0:8000".to_owned()); HttpServer::new({ @@ -280,7 +316,6 @@ async fn main() -> std::io::Result<()> { .wrap(middleware::Logger::default()) .app_data(db.clone()) .app_data(expiry_watch_sender.clone()) - .app_data(Bytes::configure(|cfg| cfg.limit(upload_max_bytes))) .data(config.clone()) .service(web::resource("/").route(web::get().to(index))) .service(web::resource("/upload").route(web::post().to(upload))) diff --git a/src/multipart.rs b/src/multipart.rs index 5103f9b..538a235 100644 --- a/src/multipart.rs +++ b/src/multipart.rs @@ -9,6 +9,7 @@ pub(crate) async fn parse_multipart( mut payload: Multipart, file_id: &str, filename: &Path, + max_size: Option, ) -> Result<(Option, DateTime, FileKind), error::Error> { let mut original_name: Option = None; let mut keep_for: Option = None; @@ -31,9 +32,7 @@ pub(crate) async fn parse_multipart( let mut file = fs::File::create(&filename) .await .map_err(|_| error::ErrorInternalServerError("could not create file"))?; - write_to_file(&mut file, field) - .await - .map_err(|_| error::ErrorInternalServerError("could not write file"))?; + write_to_file(&mut file, field, max_size).await?; } "text" => { if original_name.is_some() { @@ -44,9 +43,7 @@ pub(crate) async fn parse_multipart( let mut file = fs::File::create(&filename) .await .map_err(|_| error::ErrorInternalServerError("could not create file"))?; - write_to_file(&mut file, field) - .await - .map_err(|_| error::ErrorInternalServerError("could not write file"))?; + write_to_file(&mut file, field, max_size).await?; } _ => {} }; @@ -101,10 +98,23 @@ async fn read_content(mut field: actix_multipart::Field) -> Result, erro async fn write_to_file( file: &mut File, mut field: actix_multipart::Field, + max_size: Option, ) -> Result<(), error::Error> { + let mut written_bytes: u64 = 0; while let Some(chunk) = field.next().await { - file.write_all(chunk.map_err(error::ErrorBadRequest)?.as_ref()) - .await?; + let chunk = chunk.map_err(error::ErrorBadRequest)?; + if let Some(max_size) = max_size { + written_bytes += chunk.len() as u64; + if written_bytes > max_size { + return Err(error::ErrorBadRequest(format!( + "exceeded maximum file size of {} bytes", + max_size + ))); + } + } + file.write_all(chunk.as_ref()) + .await + .map_err(|_| error::ErrorInternalServerError("could not write file"))?; } Ok(()) }