do mime guessing, fix MAX_UPLOAD_BYTES

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"

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"]

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::<Mime>()
+        .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<HttpResponse, Error> {
     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<u64>,
 }
 
 #[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)))

src/multipart.rs

@@ -9,6 +9,7 @@ pub(crate) async fn parse_multipart(
     mut payload: Multipart,
     file_id: &str,
     filename: &Path,
+    max_size: Option<u64>,
 ) -> Result<(Option<String>, DateTime<Local>, FileKind), error::Error> {
     let mut original_name: Option<String> = None;
     let mut keep_for: Option<String> = 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<Vec<u8>, erro
 async fn write_to_file(
     file: &mut File,
     mut field: actix_multipart::Field,
+    max_size: Option<u64>,
 ) -> 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(())
 }