diff --git a/Cargo.lock b/Cargo.lock index 601710f..992ff07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -416,7 +416,7 @@ dependencies = [ [[package]] name = "datatrash" -version = "2.0.0" +version = "2.0.1" dependencies = [ "actix-files", "actix-governor", diff --git a/Cargo.toml b/Cargo.toml index 36dd045..93e1bee 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "datatrash" -version = "2.0.0" +version = "2.0.1" authors = ["neri"] edition = "2021" diff --git a/src/multipart.rs b/src/multipart.rs index d0d42d0..0fa4a3e 100644 --- a/src/multipart.rs +++ b/src/multipart.rs @@ -3,7 +3,10 @@ use actix_multipart::{Field, Multipart}; use actix_web::{error, http::header::DispositionParam, Error}; use futures_util::{StreamExt, TryStreamExt}; use mime::{Mime, APPLICATION_OCTET_STREAM, TEXT_PLAIN}; -use std::path::Path; +use std::{ + cmp::{max, min}, + path::Path, +}; use time::{Duration, OffsetDateTime}; use tokio::{fs::File, io::AsyncWriteExt}; @@ -42,10 +45,10 @@ pub(crate) async fn parse_multipart( continue; } original_name = uploaded_name; - size = create_file(file_path, field, config.max_file_size).await?; - println!("mime: {}", mime); + let first_bytes; + (size, first_bytes) = create_file(file_path, field, config.max_file_size).await?; content_type = Some(if mime == APPLICATION_OCTET_STREAM { - get_content_type(file_path) + get_content_type(&first_bytes) } else { mime_relations::get_alias(mime) }); @@ -54,8 +57,9 @@ pub(crate) async fn parse_multipart( if original_name.is_some() { continue; } - size = create_file(file_path, field, config.max_file_size).await?; - content_type = Some(get_content_type(file_path)); + let first_bytes; + (size, first_bytes) = create_file(file_path, field, config.max_file_size).await?; + content_type = Some(get_content_type(&first_bytes)); } "delete_on_download" => { delete_on_download = parse_string(name, field).await? != "false"; @@ -149,23 +153,25 @@ async fn create_file( filename: &Path, field: Field, max_file_size: Option, -) -> Result { +) -> Result<(u64, Vec), Error> { let mut file = File::create(&filename).await.map_err(|file_err| { log::error!("could not create file {:?}", file_err); error::ErrorInternalServerError("could not create file") })?; - let written_bytes = write_to_file(&mut file, field, max_file_size).await?; - Ok(written_bytes) + write_to_file(&mut file, field, max_file_size).await } async fn write_to_file( file: &mut File, mut field: Field, max_size: Option, -) -> Result { +) -> Result<(u64, Vec), error::Error> { + let mut first_bytes = Vec::with_capacity(2048); let mut written_bytes: u64 = 0; while let Some(chunk) = field.next().await { let chunk = chunk.map_err(error::ErrorBadRequest)?; + let remaining_first_bytes = min(max(0, 2048 - written_bytes) as usize, chunk.len()); + first_bytes.extend_from_slice(&chunk[0..remaining_first_bytes]); written_bytes += chunk.len() as u64; if let Some(max_size) = max_size { if written_bytes > max_size { @@ -179,7 +185,7 @@ async fn write_to_file( error::ErrorInternalServerError("could not write file") })?; } - Ok(written_bytes) + Ok((written_bytes, first_bytes)) } fn get_file_metadata(field: &actix_multipart::Field) -> (Mime, Option) { @@ -195,9 +201,9 @@ fn get_file_metadata(field: &actix_multipart::Field) -> (Mime, Option) { (mime, filename) } -fn get_content_type(path: &Path) -> Mime { - let std_path = std::path::Path::new(path.as_os_str()); - tree_magic_mini::from_filepath(std_path) - .and_then(|mime| mime.parse().ok()) +fn get_content_type(bytes: &[u8]) -> Mime { + tree_magic_mini::from_u8(bytes) + .parse() + .ok() .unwrap_or(TEXT_PLAIN) }