datatrash/src/main.rs

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::Sender, task};
use file_kind::FileKind;
use sqlx::postgres::PgPool;
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<HttpResponse, Error> {
    Ok(HttpResponse::Ok()
        .content_type("text/html")
        .body(INDEX_HTML))
}

async fn upload(
    payload: Multipart,
    db: web::Data<PgPool>,
    sender: web::Data<Sender<()>>,
) -> Result<HttpResponse, Error> {
    let file_id = format!("{:x?}", rand::random::<u32>());
    let filename = PathBuf::from(format!("files/{}", 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)",
        file_id,
        original_name.unwrap_or_else(|| file_id.clone()),
        valid_till.naive_local(),
        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::Found()
        .header("location", format!("/upload/{}", file_id))
        .finish())
}

async fn uploaded(id: web::Path<String>) -> Result<HttpResponse, Error> {
    let upload_html = UPLOAD_HTML.replace("{id}", id.as_ref());
    Ok(HttpResponse::Ok()
        .content_type("text/html")
        .body(upload_html))
}

async fn download(
    req: HttpRequest,
    id: web::Path<String>,
    db: web::Data<PgPool>,
) -> Result<HttpResponse, Error> {
    let row = sqlx::query!(
        "SELECT file_id, file_name, kind from files WHERE file_id = $1",
        *id
    )
    .fetch_one(db.as_ref())
    .await
    .map_err(|_| error::ErrorNotFound("could not find file"))?;
    let path: PathBuf = PathBuf::from(format!("files/{}", row.file_id));

    if row.kind == FileKind::TEXT.to_string() {
        let content = fs::read_to_string(path).await?;
        let view_html = VIEW_HTML.replace("{text}", &content);
        let response = HttpResponse::Ok().content_type("text/html").body(view_html);
        Ok(response)
    } else {
        let file = NamedFile::open(path)?.set_content_disposition(ContentDisposition {
            disposition: DispositionType::Attachment,
            parameters: vec![DispositionParam::Filename(row.file_name)],
        });
        file.into_response(&req)
    }
}

async fn setup_db() -> PgPool {
    let pool = PgPool::builder()
        .max_size(5)
        .build(&env::var("DATABASE_URL").expect("DATABASE_URL environement variable not set"))
        .await
        .expect("could not create db pool");

    sqlx::query_file!("./init-db.sql")
        .execute(&pool)
        .await
        .expect("could not create table Files");

    pool
}

#[actix_rt::main]
async fn main() -> std::io::Result<()> {
    std::env::set_var("RUST_LOG", "warn,datatrash=info,actix_web=info");
    env_logger::init();

    let pool: PgPool = setup_db().await;

    log::info!("omnomnom");

    let (send, recv) = async_std::sync::channel::<()>(1);
    task::spawn(deleter::delete_old_files(recv, pool.clone()));

    let db = web::Data::new(pool);
    let send = web::Data::new(send);

    HttpServer::new(move || {
        App::new()
            .wrap(middleware::Logger::default())
            .app_data(db.clone())
            .app_data(send.clone())
            .app_data(Bytes::configure(|cfg| cfg.limit(8_388_608)))
            .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("0.0.0.0:8000")?
    .run()
    .await
}