From dfd3387ac64514c661e441f3a18df0f650974643 Mon Sep 17 00:00:00 2001 From: neri Date: Sat, 11 Sep 2021 02:08:47 +0200 Subject: [PATCH] ensure file ids are unique --- src/deleter.rs | 2 +- src/download.rs | 2 +- src/multipart.rs | 33 +++++++++++++++++++-------------- src/upload.rs | 42 ++++++++++++++++++++++++++++++++++-------- 4 files changed, 55 insertions(+), 24 deletions(-) diff --git a/src/deleter.rs b/src/deleter.rs index 070bdb6..2a26bd2 100644 --- a/src/deleter.rs +++ b/src/deleter.rs @@ -36,7 +36,7 @@ pub(crate) async fn delete_by_id( file_id: &str, files_dir: &Path, ) -> Result<(), sqlx::Error> { - delete_content(file_id, &files_dir).await?; + delete_content(file_id, files_dir).await?; sqlx::query("DELETE FROM files WHERE file_id = $1") .bind(file_id) .execute(db) diff --git a/src/download.rs b/src/download.rs index 9783c26..da72bfe 100644 --- a/src/download.rs +++ b/src/download.rs @@ -102,7 +102,7 @@ fn build_file_response( } else { DispositionType::Inline }, - parameters: get_disposition_params(&file_name), + parameters: get_disposition_params(file_name), }; let file = NamedFile::open(path) .map_err(|file_err| { diff --git a/src/multipart.rs b/src/multipart.rs index 0b1e3ed..7656917 100644 --- a/src/multipart.rs +++ b/src/multipart.rs @@ -1,7 +1,7 @@ use crate::{config, file_kind::FileKind}; use actix_multipart::{Field, Multipart}; -use actix_web::{error, http::header::DispositionParam}; -use async_std::{fs, fs::File, path::Path, prelude::*}; +use actix_web::{error, http::header::DispositionParam, Error}; +use async_std::{fs::File, path::Path, prelude::*}; use chrono::{prelude::*, Duration}; use futures::{StreamExt, TryStreamExt}; @@ -18,7 +18,7 @@ pub(crate) struct UploadConfig { pub(crate) async fn parse_multipart( mut payload: Multipart, file_id: &str, - filename: &Path, + file_name: &Path, config: &config::Config, ) -> Result { let mut original_name: Option = None; @@ -42,11 +42,7 @@ pub(crate) async fn parse_multipart( } original_name = file_original_name; kind = Some(FileKind::Binary); - let mut file = fs::File::create(&filename).await.map_err(|file_err| { - log::error!("could not create file {:?}", file_err); - error::ErrorInternalServerError("could not create file") - })?; - size = write_to_file(&mut file, field, config.max_file_size).await?; + size = create_file(file_name, field, config.max_file_size).await?; } "text" => { if original_name.is_some() { @@ -54,11 +50,7 @@ pub(crate) async fn parse_multipart( } original_name = Some(format!("{}.txt", file_id)); kind = Some(FileKind::Text); - let mut file = fs::File::create(&filename).await.map_err(|file_err| { - log::error!("could not create file {:?}", file_err); - error::ErrorInternalServerError("could not create file") - })?; - size = write_to_file(&mut file, field, config.max_file_size).await?; + size = create_file(file_name, field, config.max_file_size).await?; } "delete_on_download" => { delete_on_download = dbg!(parse_string(name, field).await?) != "false"; @@ -145,9 +137,22 @@ async fn read_content(mut field: actix_multipart::Field) -> Result, erro Ok(data) } +async fn create_file( + filename: &Path, + field: Field, + max_file_size: Option, +) -> Result { + 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) +} + async fn write_to_file( file: &mut File, - mut field: actix_multipart::Field, + mut field: Field, max_size: Option, ) -> Result { let mut written_bytes: u64 = 0; diff --git a/src/upload.rs b/src/upload.rs index 3fe23f6..d96790c 100644 --- a/src/upload.rs +++ b/src/upload.rs @@ -1,3 +1,4 @@ +use std::io::ErrorKind; use std::{cmp, vec}; use crate::config::Config; @@ -6,7 +7,11 @@ 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 async_std::{ + channel::Sender, + fs::{self, OpenOptions}, + path::PathBuf, +}; use chrono::Duration; use rand::prelude::SliceRandom; use sqlx::postgres::PgPool; @@ -142,11 +147,12 @@ pub async fn upload( 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 (file_id, file_name) = create_unique_file(&config).await.map_err(|file_err| { + log::error!("could not create file {:?}", file_err); + error::ErrorInternalServerError("could not create file") + })?; - let parsed_multipart = multipart::parse_multipart(payload, &file_id, &filename, &config).await; + let parsed_multipart = multipart::parse_multipart(payload, &file_id, &file_name, &config).await; let UploadConfig { original_name, valid_till, @@ -155,8 +161,8 @@ pub async fn upload( } = match parsed_multipart { Ok(data) => data, Err(err) => { - if filename.exists().await { - fs::remove_file(filename).await.map_err(|file_err| { + if file_name.exists().await { + fs::remove_file(file_name).await.map_err(|file_err| { log::error!("could not remove file {:?}", file_err); error::ErrorInternalServerError( "could not parse multipart; could not remove file", @@ -180,7 +186,7 @@ pub async fn upload( .await; if let Err(db_err) = db_insert { log::error!("could not insert into datebase {:?}", db_err); - fs::remove_file(filename).await.map_err(|file_err| { + fs::remove_file(file_name).await.map_err(|file_err| { log::error!("could not remove file {:?}", file_err); error::ErrorInternalServerError( "could not insert file into database; could not remove file", @@ -215,6 +221,26 @@ pub async fn upload( .body(format!("{}\n", url))) } +async fn create_unique_file( + config: &web::Data, +) -> Result<(String, PathBuf), std::io::Error> { + loop { + let file_id = gen_file_id(); + let mut file_name = config.files_dir.clone(); + file_name.push(&file_id); + match OpenOptions::new() + .write(true) + .create_new(true) + .open(&file_name) + .await + { + Ok(_) => return Ok((file_id, file_name)), + Err(error) if error.kind() == ErrorKind::AlreadyExists => continue, + Err(error) => return Err(error), + } + } +} + fn gen_file_id() -> String { let mut rng = rand::thread_rng(); let mut id = String::with_capacity(5);