transbeam/src/main.rs

mod download;
mod file;
mod upload;
mod zip;

use std::{
    collections::HashMap,
    path::PathBuf,
    fs::File,
};

use actix::Addr;
use actix_web::{
    get, middleware::Logger, web, App, HttpRequest, HttpServer, Responder, HttpResponse,
};
use actix_web_actors::ws;
use time::OffsetDateTime;
use tokio::sync::RwLock;

const APP_NAME: &'static str = "transbeam";

pub struct UploadedFile {
    name: String,
    size: u64,
    modtime: OffsetDateTime,
}

impl UploadedFile {
    fn new(name: &str, size: u64, modtime: OffsetDateTime) -> Self {
        Self {
            name: sanitise_file_name::sanitise(name),
            size,
            modtime,
        }
    }
}

#[derive(Clone)]
pub struct DownloadableFile {
    name: String,
    size: u64,
    modtime: OffsetDateTime,
    uploader: Option<Addr<upload::Uploader>>,
}

type AppData = web::Data<RwLock<HashMap<String, DownloadableFile>>>;

fn storage_dir() -> PathBuf {
    PathBuf::from(std::env::var("STORAGE_DIR").unwrap_or_else(|_| String::from("storage")))
}

#[get("/download/{file_code}")]
async fn handle_download(req: HttpRequest, path: web::Path<String>, data: AppData) -> actix_web::Result<HttpResponse> {
    let file_code = path.into_inner();
    if !file_code.as_bytes().iter().all(|c| c.is_ascii_alphanumeric()) {
        return Ok(HttpResponse::NotFound().finish());
    }
    let data = data.read().await;
    let info = data.get(&file_code);
    if let Some(info) = info {
        Ok(download::DownloadingFile {
            file: File::open(storage_dir().join(file_code))?,
            info: info.clone(),
        }.into_response(&req))
    } else {
        Ok(HttpResponse::NotFound().finish())
    }
}

#[get("/upload")]
async fn handle_upload(req: HttpRequest, stream: web::Payload, data: AppData) -> impl Responder {
    ws::start(upload::Uploader::new(data), &req, stream)
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    env_logger::init();

    let data: AppData = web::Data::new(RwLock::new(HashMap::new()));

    let static_dir =
        PathBuf::from(std::env::var("STATIC_DIR").unwrap_or_else(|_| String::from("static")));
    let port = std::env::var("PORT")
        .ok()
        .and_then(|p| p.parse::<u16>().ok())
        .unwrap_or(8080);

    HttpServer::new(move || {
        App::new()
            .app_data(data.clone())
            .wrap(Logger::default())
            .service(handle_upload)
            .service(handle_download)
            .service(actix_files::Files::new("/", static_dir.clone()).index_file("index.html"))
    })
    .bind(("127.0.0.1", port))?
    .run()
    .await?;
    Ok(())
}