do mime guessing, fix MAX_UPLOAD_BYTES

This commit is contained in:
neri 2021-03-09 23:36:24 +01:00
parent 3a3174619d
commit 1d51c200d6
4 changed files with 112 additions and 35 deletions

30
Cargo.lock generated
View File

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

View File

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

View File

@ -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)))

View File

@ -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(())
}