datatrash/src/main.rs

mod multipart;

use actix_files::Files;
use actix_multipart::Multipart;
use actix_web::{error, middleware, web, App, Error, HttpResponse, HttpServer};
use async_std::prelude::*;
use chrono::{prelude::*, Duration};
use futures::{StreamExt, TryStreamExt};
use sqlx::postgres::PgPool;
use std::env;

const INDEX_HTML: &str = include_str!("../static/index.html");
const UPLOAD_HTML: &str = include_str!("../static/upload.html");

async fn index() -> Result<HttpResponse, Error> {
    Ok(HttpResponse::Ok()
        .content_type("text/html")
        .body(INDEX_HTML))
}

async fn upload(mut payload: Multipart, db: web::Data<PgPool>) -> Result<HttpResponse, Error> {
    let id = format!("{:x?}", rand::random::<u32>());
    let filename = format!("files/{}", id);
    let mut timeout: Option<String> = None;
    let mut kind: Option<String> = None;

    while let Ok(Some(mut field)) = payload.try_next().await {
        let name = multipart::get_field_name(&field)?;
        match name.as_str() {
            "validity_secs" => {
                timeout = multipart::read_string(field)
                    .await
                    .map(Some)
                    .map_err(error::ErrorInternalServerError)?;
            }
            "kind" => {
                kind = multipart::read_string(field)
                    .await
                    .map(Some)
                    .map_err(error::ErrorInternalServerError)?;
            }
            "content" => {
                let mut file = async_std::fs::File::create(&filename)
                    .await
                    .map_err(error::ErrorInternalServerError)?;

                while let Some(chunk) = field.next().await {
                    let data = chunk.unwrap();
                    file = file.write_all(&data).await.map(|_| file)?;
                }
            }
            _ => {}
        };
    }

    println!("timeout = {:?}, kind = {:?}", timeout, kind);

    if timeout == None || kind == None {
        async_std::fs::remove_file(&filename)
            .await
            .expect("could not delete file");
        return Ok(HttpResponse::BadRequest().body("timeout or kind not specified"));
    }

    let validity_secs = timeout
        .unwrap()
        .parse::<i64>()
        .expect("could not parse validity as int");
    let valid_till = Local::now() + Duration::seconds(validity_secs);
    let kind = kind.unwrap();

    sqlx::query("INSERT INTO Files (valid_till, kind) VALUES ($1, $2)")
        .bind(valid_till)
        .bind(kind)
        .execute(db.as_ref())
        .await
        .expect("could not insert");

    Ok(HttpResponse::Found()
        .header("location", format!("/upload/{}", id))
        .finish())
}

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

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

    let pool: PgPool = PgPool::builder()
        .max_size(5) // maximum number of connections in the pool
        .build(&env::var("DATABASE_URL").expect("DATABASE_URL environement variable not set"))
        .await
        .expect("could not create db pool");
    sqlx::query!("CREATE TABLE IF NOT EXISTS Files ( id serial, valid_till timestamp, kind varchar(255), primary key (id) )")
        .execute(&pool)
        .await
        .expect("could not create table Files");

    log::info!("omnomnom");

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

    HttpServer::new(move || {
        App::new()
            .wrap(middleware::Logger::default())
            .app_data(db.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(Files::new("/static", "static").disable_content_disposition())
            .service(Files::new("/file", "files"))
    })
    .bind("0.0.0.0:8000")?
    .run()
    .await
}