mod deleter; mod file_kind; mod multipart; use actix_files::{Files, NamedFile}; use actix_multipart::Multipart; use actix_web::{ error, http::header::{Charset, ContentDisposition, DispositionParam, DispositionType, ExtendedValue}, middleware, web, App, Error, HttpRequest, HttpResponse, HttpServer, }; use async_std::{ channel::{self, Sender}, fs, path::{Path, PathBuf}, task, }; use env_logger::Env; use file_kind::FileKind; use futures::TryStreamExt; use mime::Mime; use multipart::UploadConfig; use rand::prelude::SliceRandom; use sqlx::{ postgres::{PgPool, PgPoolOptions, PgRow}, Row, }; use std::env; const INDEX_HTML: &str = include_str!("../template/index.html"); const UPLOAD_HTML: &str = include_str!("../template/upload.html"); const VIEW_HTML: &str = include_str!("../template/view.html"); const ID_CHARS: &[char] = &[ 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'j', 'k', 'm', 'n', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '1', '2', '3', '4', '5', '6', '7', '8', '9', ]; async fn index(req: web::HttpRequest) -> Result { let upload_url = format!("{}/upload", get_host_url(&req)); let index_html = INDEX_HTML.replace("{upload_url}", upload_url.as_str()); Ok(HttpResponse::Ok() .content_type("text/html") .body(index_html)) } // multipart data // required: either 'file' or 'text' // optional: 'keep_for' default to 30 minutes async fn upload( req: web::HttpRequest, payload: Multipart, db: web::Data, 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 parsed_multipart = multipart::parse_multipart(payload, &file_id, &filename, config.max_file_size).await; let UploadConfig { original_name, valid_till, kind, delete_on_download, } = 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, delete_on_download) \ VALUES ($1, $2, $3, $4, $5)", ) .bind(&file_id) .bind(&original_name) .bind(valid_till.naive_local()) .bind(kind.to_string()) .bind(delete_on_download) .execute(db.as_ref()) .await; if db_insert.is_err() { fs::remove_file(filename).await.map_err(|_| { error::ErrorInternalServerError( "could not insert file into database; could not remove file", ) })?; return Err(error::ErrorInternalServerError( "could not insert file into database", )); } log::info!( "{} create new file {} (valid_till: {}, kind: {}, delete_on_download: {})", req.connection_info().realip_remote_addr().unwrap_or("-"), file_id, valid_till, kind, delete_on_download ); expiry_watch_sender.send(()).await.unwrap(); let redirect = if kind == FileKind::BINARY { let encoded_name = urlencoding::encode(&original_name); format!("/upload/{}/{}", file_id, encoded_name) } else { format!("/upload/{}", file_id) }; let url = get_file_url(&req, &file_id, Some(&original_name)); Ok(HttpResponse::SeeOther() .header("location", redirect) .body(format!("{}\n", url))) } fn gen_file_id() -> String { let mut rng = rand::thread_rng(); let mut id = String::with_capacity(5); for _ in 0..5 { id.push(*ID_CHARS.choose(&mut rng).expect("ID_CHARS is not empty")); } id } fn get_host_url(req: &web::HttpRequest) -> String { let conn = req.connection_info(); format!("{}://{}", conn.scheme(), conn.host()) } fn get_file_url(req: &web::HttpRequest, id: &str, name: Option<&str>) -> String { if let Some(name) = name { let encoded_name = urlencoding::encode(name); format!("{}/{}/{}", get_host_url(req), id, encoded_name) } else { format!("{}/{}", get_host_url(req), id) } } async fn uploaded(req: web::HttpRequest) -> Result { let id = req.match_info().query("id"); let name = req.match_info().get("name"); let url = get_file_url(&req, id, name); let upload_html = UPLOAD_HTML.replace("{url}", url.as_str()); Ok(HttpResponse::Ok() .content_type("text/html") .body(upload_html)) } async fn download( req: HttpRequest, db: web::Data, config: web::Data, ) -> Result { let id = req.match_info().query("id"); let mut rows = sqlx::query("SELECT file_id, file_name, delete_on_download from files WHERE file_id = $1") .bind(id) .fetch(db.as_ref()); let row: PgRow = rows .try_next() .await .map_err(|_| error::ErrorInternalServerError("could not run select statement"))? .ok_or_else(|| error::ErrorNotFound("file does not exist or has expired"))?; let file_id: String = row.get("file_id"); let file_name: String = row.get("file_name"); let delete_on_download: bool = row.get("delete_on_download"); let mut path = config.files_dir.clone(); path.push(&file_id); let download = req.query_string().contains("dl"); let (content_type, mut content_disposition) = get_content_types(&path, &file_name); let response = if content_type.type_() == mime::TEXT && !download { let content = fs::read_to_string(path).await.map_err(|_| { error::ErrorInternalServerError("this file should be here but could not be found") })?; let encoded = htmlescape::encode_minimal(&content); let view_html = VIEW_HTML.replace("{text}", &encoded); let response = HttpResponse::Ok().content_type("text/html").body(view_html); Ok(response) } else { if download { content_disposition.disposition = DispositionType::Attachment; } let file = NamedFile::open(path) .map_err(|_| { error::ErrorInternalServerError("this file should be here but could not be found") })? .set_content_type(content_type) .set_content_disposition(content_disposition); file.into_response(&req) }; if delete_on_download { deleter::delete_by_id(&db, &file_id, &config.files_dir) .await .map_err(|_| error::ErrorInternalServerError("could not delete file"))?; } response } 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::() .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 cd = ContentDisposition { disposition, parameters: get_disposition_params(filename), }; (ct, cd) } fn get_disposition_params(filename: &str) -> Vec { 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(), })) } parameters } async fn not_found() -> Result { Ok(HttpResponse::NotFound() .content_type("text/plain") .body("not found")) } fn get_db_url() -> String { if let Ok(database_url) = env::var("DATABASE_URL") { return database_url; } let auth = if let Ok(user) = env::var("DATABASE_USER") { if let Ok(pass) = env::var("DATABASE_PASS") { format!("{}:{}@", user, pass) } else { format!("{}@", user) } } else { String::new() }; format!( "postgresql://{auth}{host}/{name}", auth = auth, host = env::var("DATABASE_HOST").unwrap_or_else(|_| "localhost".to_string()), name = env::var("DATABASE_NAME").unwrap_or_else(|_| "datatrash".to_string()) ) } async fn setup_db() -> PgPool { let conn_url = &get_db_url(); log::info!("Using Connection string {}", conn_url); let pool = PgPoolOptions::new() .max_connections(5) .connect_timeout(std::time::Duration::from_secs(5)) .connect(conn_url) .await .expect("could not create db pool"); for query in include_str!("../init-db.sql").split_inclusive(";") { sqlx::query(query) .execute(&pool) .await .expect("could not initialize database schema"); } pool } #[derive(Clone)] struct Config { files_dir: PathBuf, max_file_size: Option, } #[actix_web::main] async fn main() -> std::io::Result<()> { env_logger::Builder::from_env(Env::default().default_filter_or("info,sqlx=warn")).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 .expect("could not create directory for storing files"); let (sender, receiver) = channel::bounded(8); log::info!("omnomnom"); task::spawn(deleter::delete_old_files( receiver, pool.clone(), config.files_dir.clone(), )); let db = web::Data::new(pool); let expiry_watch_sender = web::Data::new(sender); let bind_address = env::var("BIND_ADDRESS").unwrap_or_else(|_| "0.0.0.0:8000".to_owned()); HttpServer::new({ move || { App::new() .wrap(middleware::Logger::new(r#"%{r}a "%r" =%s %bbytes %Tsec"#)) .app_data(db.clone()) .app_data(expiry_watch_sender.clone()) .data(config.clone()) .service(web::resource("/").route(web::get().to(index))) .service(web::resource("/upload").route(web::post().to(upload))) .service( web::resource(["/upload/{id}", "/upload/{id}/{name}"]) .route(web::get().to(uploaded)), ) .service(Files::new("/static", "static").disable_content_disposition()) .service( web::resource(["/{id:[a-z0-9]{5}}", "/{id:[a-z0-9]{5}}/{name}"]) .route(web::get().to(download)), ) .default_service(web::route().to(not_found)) } }) .bind(bind_address)? .run() .await }