From f3855be2c929d55b3a64383d173f3830ba5682ee Mon Sep 17 00:00:00 2001 From: orhun Date: Wed, 4 Aug 2021 17:35:54 +0300 Subject: [PATCH] feat(paste): support shortening URLs --- Cargo.lock | 1 + Cargo.toml | 1 + src/file.rs | 109 --------------------------- src/lib.rs | 6 +- src/main.rs | 3 +- src/paste.rs | 198 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/server.rs | 40 +++++++--- 7 files changed, 235 insertions(+), 123 deletions(-) delete mode 100644 src/file.rs create mode 100644 src/paste.rs diff --git a/Cargo.lock b/Cargo.lock index 7236f6d..3ac1e90 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1523,6 +1523,7 @@ dependencies = [ "petname", "rand 0.8.4", "serde", + "url", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index c244ff5..fdc7223 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ futures-util = "0.3.15" petname = "1.1.0" rand = "0.8.4" dotenv = "0.15.0" +url = "2.2.2" [dependencies.config] version = "0.11.0" diff --git a/src/file.rs b/src/file.rs deleted file mode 100644 index d97669f..0000000 --- a/src/file.rs +++ /dev/null @@ -1,109 +0,0 @@ -use crate::config::Config; -use std::fs::File; -use std::io::{Result as IoResult, Write}; -use std::path::PathBuf; - -/// Writes the bytes to a file in upload directory. -/// -/// - If `file_name` does not have an extension, it is replaced with [`default_extension`]. -/// - If `file_name` is "-", it is replaced with "stdin". -/// - If [`random_url.enabled`] is `true`, `file_name` is replaced with a pet name or random string. -/// -/// [`default_extension`]: crate::config::PasteConfig::default_extension -/// [`random_url.enabled`]: crate::random::RandomURLConfig::enabled -pub fn save(file_name: &str, bytes: &[u8], config: &Config) -> IoResult { - let file_name = match PathBuf::from(file_name) - .file_name() - .map(|v| v.to_str()) - .flatten() - { - Some("-") => String::from("stdin"), - Some(v) => v.to_string(), - None => String::from("file"), - }; - let mut path = config.server.upload_path.join(file_name); - match path.clone().extension() { - Some(extension) => { - if let Some(url) = config.paste.random_url.generate() { - path.set_file_name(url); - path.set_extension(extension); - } - } - None => { - if let Some(url) = config.paste.random_url.generate() { - path.set_file_name(url); - } - path.set_extension( - infer::get(bytes) - .map(|t| t.extension()) - .unwrap_or(&config.paste.default_extension), - ); - } - } - let mut buffer = File::create(&path)?; - buffer.write_all(bytes)?; - Ok(path - .file_name() - .map(|v| v.to_string_lossy()) - .unwrap_or_default() - .to_string()) -} - -#[cfg(test)] -mod test { - use super::*; - use crate::random::{RandomURLConfig, RandomURLType}; - use std::env; - use std::fs; - use std::path::PathBuf; - - #[test] - fn test_save_file() -> IoResult<()> { - let mut config = Config::default(); - config.server.upload_path = env::current_dir()?; - config.paste.random_url = RandomURLConfig { - enabled: true, - words: Some(3), - separator: Some(String::from("_")), - type_: RandomURLType::PetName, - ..RandomURLConfig::default() - }; - let file_name = save("test.txt", &[65, 66, 67], &config)?; - assert_eq!("ABC", fs::read_to_string(&file_name)?); - assert_eq!( - Some("txt"), - PathBuf::from(&file_name) - .extension() - .map(|v| v.to_str()) - .flatten() - ); - fs::remove_file(file_name)?; - - config.paste.default_extension = String::from("bin"); - config.paste.random_url.enabled = false; - config.paste.random_url = RandomURLConfig { - enabled: true, - length: Some(10), - type_: RandomURLType::Alphanumeric, - ..RandomURLConfig::default() - }; - let file_name = save("random", &[120, 121, 122], &config)?; - assert_eq!("xyz", fs::read_to_string(&file_name)?); - assert_eq!( - Some("bin"), - PathBuf::from(&file_name) - .extension() - .map(|v| v.to_str()) - .flatten() - ); - fs::remove_file(file_name)?; - - config.paste.random_url.enabled = false; - let file_name = save("test.file", &[116, 101, 115, 116], &config)?; - assert_eq!("test.file", &file_name); - assert_eq!("test", fs::read_to_string(&file_name)?); - fs::remove_file(file_name)?; - - Ok(()) - } -} diff --git a/src/lib.rs b/src/lib.rs index 330fe38..8bd023b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,8 +13,8 @@ pub mod server; /// HTTP headers. pub mod header; -/// File handler. -pub mod file; - /// Auth handler. pub mod auth; + +/// Storage handler. +pub mod paste; diff --git a/src/main.rs b/src/main.rs index 5eeb74e..d62ea10 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,7 +13,8 @@ async fn main() -> IoResult<()> { let config = Config::parse(env::var("CONFIG").as_deref().unwrap_or("config")) .expect("failed to parse config"); let server_config = config.server.clone(); - fs::create_dir_all(server_config.upload_path)?; + fs::create_dir_all(&server_config.upload_path)?; + fs::create_dir_all(&server_config.upload_path.join("url"))?; let mut http_server = HttpServer::new(move || { App::new() .data(config.clone()) diff --git a/src/paste.rs b/src/paste.rs new file mode 100644 index 0000000..c745d4b --- /dev/null +++ b/src/paste.rs @@ -0,0 +1,198 @@ +use crate::config::Config; +use crate::header::ContentDisposition; +use std::convert::TryFrom; +use std::fs::{self, File}; +use std::io::{Error as IoError, ErrorKind as IoErrorKind, Result as IoResult, Write}; +use std::path::PathBuf; +use std::str; +use url::Url; + +/// Type of the data to store. +#[derive(Clone, Copy, Debug)] +pub enum PasteType { + /// Any type of file. + File, + /// A file that only contains an URL. + Url, +} + +impl<'a> TryFrom<&'a ContentDisposition> for PasteType { + type Error = (); + fn try_from(content_disposition: &'a ContentDisposition) -> Result { + if content_disposition.has_form_field("file") { + Ok(Self::File) + } else if content_disposition.has_form_field("url") { + Ok(Self::Url) + } else { + Err(()) + } + } +} + +/// Representation of a single paste. +#[derive(Debug)] +pub struct Paste { + /// Data to store. + pub data: Vec, + /// Type of the data. + pub type_: PasteType, +} + +impl Paste { + /// Writes the bytes to a file in upload directory. + /// + /// - If `file_name` does not have an extension, it is replaced with [`default_extension`]. + /// - If `file_name` is "-", it is replaced with "stdin". + /// - If [`random_url.enabled`] is `true`, `file_name` is replaced with a pet name or random string. + /// + /// [`default_extension`]: crate::config::PasteConfig::default_extension + /// [`random_url.enabled`]: crate::random::RandomURLConfig::enabled + pub fn store_file(&self, file_name: &str, config: &Config) -> IoResult { + let file_name = match PathBuf::from(file_name) + .file_name() + .map(|v| v.to_str()) + .flatten() + { + Some("-") => String::from("stdin"), + Some(v) => v.to_string(), + None => String::from("file"), + }; + let mut path = config.server.upload_path.join(file_name); + match path.clone().extension() { + Some(extension) => { + if let Some(file_name) = config.paste.random_url.generate() { + path.set_file_name(file_name); + path.set_extension(extension); + } + } + None => { + if let Some(file_name) = config.paste.random_url.generate() { + path.set_file_name(file_name); + } + path.set_extension( + infer::get(&self.data) + .map(|t| t.extension()) + .unwrap_or(&config.paste.default_extension), + ); + } + } + let mut buffer = File::create(&path)?; + buffer.write_all(&self.data)?; + Ok(path + .file_name() + .map(|v| v.to_string_lossy()) + .unwrap_or_default() + .to_string()) + } + + /// Writes an URL to a file in upload directory. + /// + /// - Checks if the data is a valid URL. + /// - If [`random_url.enabled`] is `true`, file name is set to a pet name or random string. + /// + /// [`random_url.enabled`]: crate::random::RandomURLConfig::enabled + pub fn store_url(&self, config: &Config) -> IoResult { + let data = str::from_utf8(&self.data) + .map_err(|e| IoError::new(IoErrorKind::Other, e.to_string()))?; + let url = Url::parse(data).map_err(|e| IoError::new(IoErrorKind::Other, e.to_string()))?; + let file_name = config + .paste + .random_url + .generate() + .unwrap_or_else(|| String::from("url")); + let path = config.server.upload_path.join("url").join(&file_name); + fs::write(&path, url.to_string())?; + Ok(file_name) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::random::{RandomURLConfig, RandomURLType}; + use std::env; + + #[test] + fn test_paste_data() -> IoResult<()> { + let mut config = Config::default(); + config.server.upload_path = env::current_dir()?; + config.paste.random_url = RandomURLConfig { + enabled: true, + words: Some(3), + separator: Some(String::from("_")), + type_: RandomURLType::PetName, + ..RandomURLConfig::default() + }; + let paste = Paste { + data: vec![65, 66, 67], + type_: PasteType::File, + }; + let file_name = paste.store_file("test.txt", &config)?; + assert_eq!("ABC", fs::read_to_string(&file_name)?); + assert_eq!( + Some("txt"), + PathBuf::from(&file_name) + .extension() + .map(|v| v.to_str()) + .flatten() + ); + fs::remove_file(file_name)?; + + config.paste.default_extension = String::from("bin"); + config.paste.random_url.enabled = false; + config.paste.random_url = RandomURLConfig { + enabled: true, + length: Some(10), + type_: RandomURLType::Alphanumeric, + ..RandomURLConfig::default() + }; + let paste = Paste { + data: vec![120, 121, 122], + type_: PasteType::File, + }; + let file_name = paste.store_file("random", &config)?; + assert_eq!("xyz", fs::read_to_string(&file_name)?); + assert_eq!( + Some("bin"), + PathBuf::from(&file_name) + .extension() + .map(|v| v.to_str()) + .flatten() + ); + fs::remove_file(file_name)?; + + config.paste.random_url.enabled = false; + let paste = Paste { + data: vec![116, 101, 115, 116], + type_: PasteType::File, + }; + let file_name = paste.store_file("test.file", &config)?; + assert_eq!("test.file", &file_name); + assert_eq!("test", fs::read_to_string(&file_name)?); + fs::remove_file(file_name)?; + + fs::create_dir_all(config.server.upload_path.join("url"))?; + + config.paste.random_url.enabled = true; + let url = String::from("https://orhun.dev/"); + let paste = Paste { + data: url.as_bytes().to_vec(), + type_: PasteType::Url, + }; + let file_name = paste.store_url(&config)?; + let file_path = config.server.upload_path.join("url").join(&file_name); + assert_eq!(url, fs::read_to_string(&file_path)?); + fs::remove_file(file_path)?; + + let url = String::from("testurl.com"); + let paste = Paste { + data: url.as_bytes().to_vec(), + type_: PasteType::Url, + }; + assert!(paste.store_url(&config).is_err()); + + fs::remove_dir(config.server.upload_path.join("url"))?; + + Ok(()) + } +} diff --git a/src/server.rs b/src/server.rs index b980cb2..7813cce 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,7 +1,7 @@ use crate::auth; use crate::config::Config; -use crate::file; use crate::header::ContentDisposition; +use crate::paste::{Paste, PasteType}; use actix_files::NamedFile; use actix_multipart::Multipart; use actix_web::{error, get, post, web, Error, HttpRequest, HttpResponse, Responder}; @@ -9,6 +9,7 @@ use byte_unit::Byte; use futures_util::stream::StreamExt; use std::convert::TryFrom; use std::env; +use std::fs; /// Shows the landing page. #[get("/")] @@ -22,15 +23,27 @@ async fn index() -> impl Responder { #[get("/{file}")] async fn serve( request: HttpRequest, - path: web::Path, + file: web::Path, config: web::Data, ) -> Result { - let path = config.server.upload_path.join(&*path); - let file = NamedFile::open(&path)? - .disable_content_disposition() - .prefer_utf8(true); - let response = file.into_response(&request)?; - Ok(response) + let mut path = config.server.upload_path.join(&*file); + let mut paste_type = PasteType::File; + for (type_, alt_path) in &[(PasteType::Url, "url")] { + if !path.exists() || path.file_name().map(|v| v.to_str()).flatten() == Some(alt_path) { + path = config.server.upload_path.join(alt_path).join(&*file); + paste_type = *type_; + break; + } + } + match paste_type { + PasteType::File => Ok(NamedFile::open(&path)? + .disable_content_disposition() + .prefer_utf8(true) + .into_response(&request)?), + PasteType::Url => Ok(HttpResponse::Found() + .header("Location", fs::read_to_string(&path)?) + .finish()), + } } /// Handles file upload by processing `multipart/form-data`. @@ -47,7 +60,7 @@ async fn upload( while let Some(item) = payload.next().await { let mut field = item?; let content = ContentDisposition::try_from(field.content_disposition())?; - if content.has_form_field("file") { + if let Ok(paste_type) = PasteType::try_from(&content) { let mut bytes = Vec::::new(); while let Some(chunk) = field.next().await { bytes.append(&mut chunk?.to_vec()); @@ -61,7 +74,14 @@ async fn upload( return Err(error::ErrorBadRequest("invalid file size")); } let bytes_unit = Byte::from_bytes(bytes.len() as u128).get_appropriate_unit(false); - let file_name = &file::save(content.get_file_name()?, &bytes, &config)?; + let paste = Paste { + data: bytes.to_vec(), + type_: paste_type, + }; + let file_name = match paste_type { + PasteType::File => paste.store_file(content.get_file_name()?, &config)?, + PasteType::Url => paste.store_url(&config)?, + }; log::info!("{} ({}) is uploaded from {}", file_name, bytes_unit, host); urls.push(format!( "{}://{}/{}\n",