mod deleter; mod file_kind; mod multipart; 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, }; use async_std::{ fs, path::PathBuf, sync::{channel, Sender}, task, }; use file_kind::FileKind; use sqlx::{ postgres::{PgPool, PgRow}, Cursor, 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"); async fn index() -> Result { Ok(HttpResponse::Ok() .content_type("text/html") .body(INDEX_HTML)) } async fn upload( payload: Multipart, db: web::Data, sender: web::Data>, config: web::Data, ) -> Result { let file_id = format!("{:x?}", rand::random::()); 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 remove file"))?; } return Err(err); } }; sqlx::query("INSERT INTO Files (file_id, file_name, valid_till, kind) VALUES ($1, $2, $3, $4)") .bind(&file_id) .bind(original_name.unwrap_or_else(|| file_id.clone())) .bind(valid_till.naive_local()) .bind(kind.to_string()) .execute(db.as_ref()) .await .map_err(|_| error::ErrorInternalServerError("could not insert file into database"))?; log::info!( "create new file {} (valid_till: {}, kind: {})", file_id, valid_till, kind ); sender.send(()).await; Ok(HttpResponse::SeeOther() .header("location", format!("/upload/{}", file_id)) .finish()) } async fn uploaded(id: web::Path, config: web::Data) -> Result { let upload_html = UPLOAD_HTML .replace("{id}", id.as_ref()) .replace("{server}", &config.server_url); Ok(HttpResponse::Ok() .content_type("text/html") .body(upload_html)) } async fn download( req: HttpRequest, id: web::Path, db: web::Data, config: web::Data, ) -> Result { let mut cursor = sqlx::query("SELECT file_id, file_name, kind from files WHERE file_id = $1") .bind(id.as_ref()) .fetch(db.as_ref()); let row: PgRow = cursor .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 kind: String = row.get("kind"); let mut path = config.files_dir.clone(); path.push(&file_id); if kind == FileKind::TEXT.to_string() { 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 { 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)], }); file.into_response(&req) } } async fn setup_db() -> PgPool { let pool = PgPool::builder() .max_size(5) .connect_timeout(std::time::Duration::from_secs(5)) .build(&env::var("DATABASE_URL").unwrap_or_else(|_| "postgresql://localhost".to_owned())) .await .expect("could not create db pool"); sqlx::query(include_str!("../init-db.sql")) .execute(&pool) .await .expect("could not create table Files"); pool } #[derive(Clone)] struct Config { server_url: String, files_dir: PathBuf, } #[actix_rt::main] async fn main() -> std::io::Result<()> { env::set_var( "RUST_LOG", env::var("RUST_LOG").unwrap_or_else(|_| "info".to_string()), ); env_logger::init(); let pool: PgPool = setup_db().await; let config = Config { server_url: env::var("SERVER_URL").unwrap_or_else(|_| "http://localhost:8000".to_owned()), files_dir: PathBuf::from(env::var("FILES_DIR").unwrap_or_else(|_| "./files".to_owned())), }; fs::create_dir_all(&config.files_dir) .await .expect("could not create directory for storing files"); log::info!("omnomnom"); let (send, recv) = channel(1); task::spawn(deleter::delete_old_files( recv, pool.clone(), config.files_dir.clone(), )); let db = web::Data::new(pool); let send = web::Data::new(send); let max_bytes: usize = env::var("UPLOAD_MAX_BYTES") .ok() .and_then(|variable| variable.parse().ok()) .unwrap_or(8_388_608); 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::default()) .app_data(db.clone()) .app_data(send.clone()) .app_data(Bytes::configure(|cfg| cfg.limit(max_bytes))) .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}").route(web::get().to(uploaded))) .service(web::resource("/file/{id}").route(web::get().to(download))) .service(Files::new("/static", "static").disable_content_disposition()) } }) .bind(bind_address)? .run() .await }