transbeam/src/main.rs

mod download;
mod store;
mod upload;
mod zip;

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

use actix_web::{
    get, middleware::Logger, web, App, HttpRequest, HttpResponse, HttpServer, Responder,
};
use actix_web_actors::ws;
use log::error;
use serde::Deserialize;
use store::FileStore;
use tokio::sync::RwLock;

const APP_NAME: &str = "transbeam";

type AppData = web::Data<RwLock<FileStore>>;

#[derive(Deserialize)]
struct DownloadRequest {
    code: String,
}

#[get("/download")]
async fn handle_download(
    req: HttpRequest,
    download: web::Query<DownloadRequest>,
    data: AppData,
) -> actix_web::Result<HttpResponse> {
    let code = &download.code;
    if !store::is_valid_storage_code(code) {
        return Ok(HttpResponse::NotFound().finish());
    }
    let data = data.read().await;
    let info = data.lookup_file(code);
    if let Some(info) = info {
        let storage_path = store::storage_dir().join(code);
        let file = File::open(&storage_path)?;
        Ok(download::DownloadingFile {
            file,
            storage_path,
            info,
        }
        .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(FileStore::load().await?));
    start_reaper(data.clone());

    let static_dir =
        PathBuf::from(std::env::var("TRANSBEAM_STATIC_DIR").unwrap_or_else(|_| String::from("static")));
    let port = std::env::var("TRANSBEAM_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(())
}

fn start_reaper(data: AppData) {
    std::thread::spawn(move || {
        actix_web::rt::System::new().block_on(async {
            loop {
                actix_web::rt::time::sleep(core::time::Duration::from_secs(86400)).await;
                if let Err(e) = data.write().await.remove_expired_files().await {
                    error!("Error reaping expired files: {}", e);
                }
            }
        });
    });
}