datatrash/src/main.rs

343 lines
11 KiB
Rust
Raw Normal View History

2020-07-09 17:27:24 +00:00
mod deleter;
mod file_kind;
2020-07-08 19:26:46 +00:00
mod multipart;
2020-07-09 17:27:24 +00:00
use actix_files::{Files, NamedFile};
2020-07-08 19:26:46 +00:00
use actix_multipart::Multipart;
2020-08-02 23:12:42 +00:00
use actix_web::{
error,
2021-03-09 22:36:24 +00:00
http::header::{Charset, ContentDisposition, DispositionParam, DispositionType, ExtendedValue},
middleware, web, App, Error, HttpRequest, HttpResponse, HttpServer,
2020-08-02 23:12:42 +00:00
};
2020-07-13 13:29:40 +00:00
use async_std::{
channel::{self, Sender},
2020-07-13 13:29:40 +00:00
fs,
2021-03-09 22:36:24 +00:00
path::{Path, PathBuf},
2020-07-13 13:29:40 +00:00
task,
};
2020-07-09 17:27:24 +00:00
use file_kind::FileKind;
use futures::TryStreamExt;
2021-03-09 22:36:24 +00:00
use mime::Mime;
use rand::prelude::SliceRandom;
use sqlx::{
postgres::{PgPool, PgPoolOptions, PgRow},
Row,
};
2020-07-08 19:26:46 +00:00
use std::env;
const INDEX_HTML: &str = include_str!("../template/index.html");
2020-07-09 17:27:24 +00:00
const UPLOAD_HTML: &str = include_str!("../template/upload.html");
const VIEW_HTML: &str = include_str!("../template/view.html");
2020-07-08 19:26:46 +00:00
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<HttpResponse, Error> {
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))
2020-07-08 19:26:46 +00:00
}
// multipart data
// required: either 'file' or 'text'
// optional: 'keep_for' default to 30 minutes
2020-07-09 17:27:24 +00:00
async fn upload(
req: web::HttpRequest,
2020-07-09 17:27:24 +00:00
payload: Multipart,
db: web::Data<PgPool>,
2020-12-03 22:30:37 +00:00
expiry_watch_sender: web::Data<Sender<()>>,
2020-07-11 21:27:15 +00:00
config: web::Data<Config>,
2020-07-09 17:27:24 +00:00
) -> Result<HttpResponse, Error> {
let file_id = gen_file_id();
2020-07-11 21:27:15 +00:00
let mut filename = config.files_dir.clone();
filename.push(&file_id);
2020-07-09 17:27:24 +00:00
2021-03-09 22:36:24 +00:00
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",
)
})?;
2020-07-08 19:26:46 +00:00
}
2021-03-09 22:36:24 +00:00
return Err(err);
}
};
2020-07-08 19:26:46 +00:00
2020-12-03 22:30:37 +00:00
let db_insert = sqlx::query(
"INSERT INTO Files (file_id, file_name, valid_till, kind) VALUES ($1, $2, $3, $4)",
)
.bind(&file_id)
.bind(original_name.as_ref().unwrap_or(&file_id))
.bind(valid_till.naive_local())
.bind(kind.to_string())
.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",
));
}
2020-07-08 19:26:46 +00:00
2020-07-09 17:27:24 +00:00
log::info!(
"create new file {} (valid_till: {}, kind: {})",
file_id,
valid_till,
kind
);
2020-07-08 19:26:46 +00:00
expiry_watch_sender.send(()).await.unwrap();
2020-07-08 19:26:46 +00:00
2020-08-03 00:42:27 +00:00
let redirect = if kind == FileKind::BINARY && original_name.is_some() {
let encoded_name = urlencoding::encode(original_name.as_ref().unwrap());
format!("/upload/{}/{}", file_id, encoded_name)
2020-08-03 00:42:27 +00:00
} else {
format!("/upload/{}", file_id)
};
let url = get_file_url(&req, &file_id, original_name.as_deref());
Ok(HttpResponse::SeeOther()
2020-08-03 00:42:27 +00:00
.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)
}
2020-07-08 19:26:46 +00:00
}
2020-08-03 00:42:27 +00:00
async fn uploaded(req: web::HttpRequest) -> Result<HttpResponse, Error> {
let id = req.match_info().query("id");
let name = req.match_info().get("name");
let url = get_file_url(&req, id, name);
2020-08-02 23:12:42 +00:00
let upload_html = UPLOAD_HTML.replace("{url}", url.as_str());
2020-07-08 19:26:46 +00:00
Ok(HttpResponse::Ok()
.content_type("text/html")
.body(upload_html))
}
2020-07-09 17:27:24 +00:00
async fn download(
req: HttpRequest,
db: web::Data<PgPool>,
2020-07-11 21:27:15 +00:00
config: web::Data<Config>,
2020-07-09 17:27:24 +00:00
) -> Result<HttpResponse, Error> {
2020-08-03 00:42:27 +00:00
let id = req.match_info().query("id");
let mut rows = sqlx::query("SELECT file_id, file_name from files WHERE file_id = $1")
2020-08-03 00:42:27 +00:00
.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");
2020-07-11 21:27:15 +00:00
let mut path = config.files_dir.clone();
path.push(&file_id);
2020-07-09 17:27:24 +00:00
let download = req.query_string().contains("dl");
let (content_type, mut content_disposition) = get_content_types(&path, &file_name);
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);
2020-07-09 17:27:24 +00:00
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")
})?
2021-03-09 22:36:24 +00:00
.set_content_type(content_type)
.set_content_disposition(content_disposition);
2020-07-09 17:27:24 +00:00
file.into_response(&req)
}
}
2021-03-09 22:36:24 +00:00
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 cd = ContentDisposition {
disposition,
parameters: get_disposition_params(filename),
};
(ct, cd)
}
2021-03-09 22:36:24 +00:00
fn get_disposition_params(filename: &str) -> Vec<DispositionParam> {
let mut parameters = vec![DispositionParam::Filename(filename.to_owned())];
2021-03-09 22:36:24 +00:00
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
2021-03-09 22:36:24 +00:00
}
2020-08-19 14:24:42 +00:00
async fn not_found() -> Result<HttpResponse, Error> {
Ok(HttpResponse::NotFound()
.content_type("text/plain")
.body("not found"))
}
2020-08-02 01:08:07 +00:00
fn get_db_url() -> String {
2020-08-02 23:28:42 +00:00
if let Ok(database_url) = env::var("DATABASE_URL") {
return database_url;
}
2020-08-02 01:08:07 +00:00
2020-08-02 23:28:42 +00:00
let auth = if let Ok(user) = env::var("DATABASE_USER") {
if let Ok(pass) = env::var("DATABASE_PASS") {
format!("{}:{}@", user, pass)
} else {
format!("{}@", user)
2020-08-02 01:08:07 +00:00
}
2020-08-02 23:28:42 +00:00
} else {
String::new()
};
format!(
2020-08-02 23:30:45 +00:00
"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())
2020-08-02 23:28:42 +00:00
)
2020-08-02 01:08:07 +00:00
}
2020-07-09 17:27:24 +00:00
async fn setup_db() -> PgPool {
2020-08-02 01:08:07 +00:00
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)
2020-07-09 17:27:24 +00:00
.await
.expect("could not create db pool");
sqlx::query(include_str!("../init-db.sql"))
2020-07-09 20:01:25 +00:00
.execute(&pool)
.await
.expect("could not create table Files");
2020-07-09 17:27:24 +00:00
pool
}
2020-07-11 21:27:15 +00:00
#[derive(Clone)]
struct Config {
files_dir: PathBuf,
2021-03-09 22:36:24 +00:00
max_file_size: Option<u64>,
2020-07-11 21:27:15 +00:00
}
#[actix_web::main]
2020-07-08 19:26:46 +00:00
async fn main() -> std::io::Result<()> {
2020-07-14 15:53:43 +00:00
if env::var("RUST_LOG").is_err() {
env::set_var("RUST_LOG", "info");
}
2020-07-08 19:26:46 +00:00
env_logger::init();
2020-07-09 17:27:24 +00:00
let pool: PgPool = setup_db().await;
2021-03-09 22:36:24 +00:00
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)
};
2020-07-11 21:27:15 +00:00
let config = Config {
files_dir: PathBuf::from(env::var("FILES_DIR").unwrap_or_else(|_| "./files".to_owned())),
2021-03-09 22:36:24 +00:00
max_file_size,
2020-07-11 21:27:15 +00:00
};
fs::create_dir_all(&config.files_dir)
.await
.expect("could not create directory for storing files");
let (sender, receiver) = channel::bounded(8);
2020-07-13 13:29:40 +00:00
log::info!("omnomnom");
2020-07-11 21:27:15 +00:00
task::spawn(deleter::delete_old_files(
2020-07-14 15:53:43 +00:00
receiver,
2020-07-11 21:27:15 +00:00
pool.clone(),
config.files_dir.clone(),
));
2020-07-09 17:27:24 +00:00
2020-07-08 19:26:46 +00:00
let db = web::Data::new(pool);
2020-12-03 22:30:37 +00:00
let expiry_watch_sender = web::Data::new(sender);
2020-07-11 21:27:15 +00:00
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())
2020-12-03 22:30:37 +00:00
.app_data(expiry_watch_sender.clone())
2020-07-11 21:27:15 +00:00
.data(config.clone())
.service(web::resource("/").route(web::get().to(index)))
.service(web::resource("/upload").route(web::post().to(upload)))
2020-08-02 23:12:42 +00:00
.service(
2020-08-03 00:42:27 +00:00
web::resource(["/upload/{id}", "/upload/{id}/{name}"])
.route(web::get().to(uploaded)),
)
.service(Files::new("/static", "static").disable_content_disposition())
2020-08-03 00:42:27 +00:00
.service(
web::resource(["/{id:[a-z0-9]{5}}", "/{id:[a-z0-9]{5}}/{name}"])
2020-08-02 23:12:42 +00:00
.route(web::get().to(download)),
)
2020-08-19 14:24:42 +00:00
.default_service(web::route().to(not_found))
2020-07-11 21:27:15 +00:00
}
2020-07-08 19:26:46 +00:00
})
2020-07-11 21:27:15 +00:00
.bind(bind_address)?
2020-07-08 19:26:46 +00:00
.run()
.await
}