datatrash/src/main.rs

208 lines
6.5 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-07-09 17:27:24 +00:00
use actix_web::{
error,
http::header::{ContentDisposition, DispositionParam, DispositionType},
middleware,
web::{self, Bytes},
App, Error, FromRequest, HttpRequest, HttpResponse, HttpServer,
};
2020-07-13 13:29:40 +00:00
use async_std::{
fs,
path::PathBuf,
sync::{channel, Sender},
task,
};
2020-07-09 17:27:24 +00:00
use file_kind::FileKind;
use sqlx::{
postgres::{PgPool, PgRow},
Cursor, Row,
};
2020-07-08 19:26:46 +00:00
use std::env;
2020-07-09 17:27:24 +00:00
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");
2020-07-08 19:26:46 +00:00
async fn index() -> Result<HttpResponse, Error> {
Ok(HttpResponse::Ok()
.content_type("text/html")
.body(INDEX_HTML))
}
2020-07-09 17:27:24 +00:00
async fn upload(
payload: Multipart,
db: web::Data<PgPool>,
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 = format!("{:x?}", rand::random::<u32>());
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
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"))?;
2020-07-08 19:26:46 +00:00
}
2020-07-09 17:27:24 +00:00
return Err(err);
2020-07-08 19:26:46 +00:00
}
};
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"))?;
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
2020-07-09 17:27:24 +00:00
sender.send(()).await;
2020-07-08 19:26:46 +00:00
Ok(HttpResponse::Found()
2020-07-09 17:27:24 +00:00
.header("location", format!("/upload/{}", file_id))
2020-07-08 19:26:46 +00:00
.finish())
}
2020-07-11 21:27:15 +00:00
async fn uploaded(id: web::Path<String>, config: web::Data<Config>) -> Result<HttpResponse, Error> {
let upload_html = UPLOAD_HTML
.replace("{id}", id.as_ref())
.replace("{server}", &config.server_url);
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,
id: web::Path<String>,
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> {
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");
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
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);
2020-07-09 17:27:24 +00:00
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)],
});
2020-07-09 17:27:24 +00:00
file.into_response(&req)
}
}
async fn setup_db() -> PgPool {
let pool = PgPool::builder()
.max_size(5)
.connect_timeout(std::time::Duration::from_secs(5))
2020-07-11 21:27:15 +00:00
.build(&env::var("DATABASE_URL").unwrap_or_else(|_| "postgresql://localhost".to_owned()))
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 {
server_url: String,
files_dir: PathBuf,
}
2020-07-08 19:26:46 +00:00
#[actix_rt::main]
async fn main() -> std::io::Result<()> {
std::env::set_var("RUST_LOG", "warn,datatrash=info,actix_web=info");
env_logger::init();
2020-07-09 17:27:24 +00:00
let pool: PgPool = setup_db().await;
2020-07-08 19:26:46 +00:00
2020-07-11 21:27:15 +00:00
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");
2020-07-13 13:29:40 +00:00
log::info!("omnomnom");
let (send, recv) = channel(1);
2020-07-11 21:27:15 +00:00
task::spawn(deleter::delete_old_files(
recv,
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-07-09 17:27:24 +00:00
let send = web::Data::new(send);
2020-07-11 21:27:15 +00:00
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())
}
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
}