transbeam/src/main.rs

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

use std::{fmt::Debug, fs::File, path::PathBuf, str::FromStr};

use actix_files::NamedFile;
use actix_web::{
    get, http::StatusCode, middleware::Logger, post, web, App, HttpRequest, HttpResponse,
    HttpServer, Responder,
};
use actix_web_actors::ws;
use bytesize::ByteSize;
use log::{error, warn};
use serde::{Deserialize, Serialize};
use store::FileStore;
use tokio::sync::RwLock;

const APP_NAME: &str = "transbeam";

struct AppState {
    file_store: RwLock<FileStore>,
    config: Config,
}

struct Config {
    max_upload_size: u64,
    max_lifetime: u16,
    upload_password: String,
    storage_dir: PathBuf,
    static_dir: PathBuf,
    reverse_proxy: bool,
    mnemonic_codes: bool,
}

pub fn get_ip_addr(req: &HttpRequest, reverse_proxy: bool) -> String {
    let conn_info = req.connection_info();
    if reverse_proxy {
        conn_info.realip_remote_addr()
    } else {
        conn_info.peer_addr()
    }
    .unwrap()
    .to_owned()
}

pub fn log_auth_failure(ip_addr: &str) {
    warn!("Incorrect authentication attempt from {}", ip_addr);
}

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

#[get("/download")]
async fn handle_download(
    req: HttpRequest,
    download: web::Query<DownloadRequest>,
    data: web::Data<AppState>,
) -> actix_web::Result<HttpResponse> {
    let code = &download.code;
    if !store::is_valid_storage_code(code) {
        return download_not_found(req, data);
    }
    let info = data.file_store.read().await.lookup_file(code);
    if let Some(info) = info {
        let storage_path = data.config.storage_dir.join(code);
        let file = File::open(&storage_path)?;
        Ok(download::DownloadingFile {
            file,
            storage_path,
            info,
        }
        .into_response(&req))
    } else {
        download_not_found(req, data)
    }
}

fn download_not_found(
    req: HttpRequest,
    data: web::Data<AppState>,
) -> actix_web::Result<HttpResponse> {
    let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
    log_auth_failure(&ip_addr);
    Ok(NamedFile::open(data.config.static_dir.join("404.html"))?
        .set_status_code(StatusCode::NOT_FOUND)
        .into_response(&req))
}

#[get("/upload")]
async fn handle_upload(
    req: HttpRequest,
    stream: web::Payload,
    data: web::Data<AppState>,
) -> impl Responder {
    if data.file_store.read().await.full() {
        return Ok(HttpResponse::BadRequest().finish());
    }
    let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
    ws::start(upload::Uploader::new(data, ip_addr), &req, stream)
}

#[derive(Deserialize)]
struct UploadPasswordCheck {
    password: String,
}

#[post("/upload/check_password")]
async fn check_upload_password(
    req: HttpRequest,
    body: web::Json<UploadPasswordCheck>,
    data: web::Data<AppState>,
) -> impl Responder {
    let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
    if body.password != data.config.upload_password {
        log_auth_failure(&ip_addr);
        HttpResponse::Forbidden().finish()
    } else {
        HttpResponse::NoContent().finish()
    }
}

#[derive(Serialize)]
struct UploadLimits {
    open: bool,
    max_size: u64,
    max_lifetime: u16,
}

#[get("/upload/limits.json")]
async fn upload_limits(data: web::Data<AppState>) -> impl Responder {
    let file_store = data.file_store.read().await;
    let open = !file_store.full();
    let available_size = file_store.available_size();
    let max_size = std::cmp::min(available_size, data.config.max_upload_size);
    web::Json(UploadLimits {
        open,
        max_size,
        max_lifetime: data.config.max_lifetime,
    })
}

fn env_or<T: FromStr>(var: &str, default: T) -> T
where
    <T as FromStr>::Err: Debug,
{
    std::env::var(var)
        .map(|val| {
            val.parse::<T>()
                .unwrap_or_else(|_| panic!("Invalid value {} for variable {}", val, var))
        })
        .unwrap_or(default)
}

fn env_or_else<T: FromStr, F: FnOnce() -> T>(var: &str, default: F) -> T
where
    <T as FromStr>::Err: Debug,
{
    std::env::var(var)
        .map(|val| {
            val.parse::<T>()
                .unwrap_or_else(|_| panic!("Invalid value {} for variable {}", val, var))
        })
        .ok()
        .unwrap_or_else(default)
}

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

    let static_dir: PathBuf = env_or_else("TRANSBEAM_STATIC_DIR", || PathBuf::from("static"));
    let storage_dir: PathBuf = env_or_else("TRANSBEAM_STORAGE_DIR", || PathBuf::from("storage"));
    let port: u16 = env_or("TRANSBEAM_PORT", 8080);
    let mnemonic_codes: bool = env_or("TRANSBEAM_MNEMONIC_CODES", true);
    let reverse_proxy: bool = env_or("TRANSBEAM_REVERSE_PROXY", false);
    let max_lifetime: u16 = env_or("TRANSBEAM_MAX_LIFETIME", 30);
    let max_upload_size: u64 =
        env_or::<ByteSize>("TRANSBEAM_MAX_UPLOAD_SIZE", ByteSize(16 * bytesize::GB)).as_u64();
    let max_storage_size: u64 =
        env_or::<ByteSize>("TRANSBEAM_MAX_STORAGE_SIZE", ByteSize(64 * bytesize::GB)).as_u64();
    let upload_password: String =
        std::env::var("TRANSBEAM_UPLOAD_PASSWORD").expect("TRANSBEAM_UPLOAD_PASSWORD must be set!");

    let data = web::Data::new(AppState {
        file_store: RwLock::new(FileStore::load(storage_dir.clone(), max_storage_size).await?),
        config: Config {
            max_upload_size,
            max_lifetime,
            upload_password,
            storage_dir,
            static_dir: static_dir.clone(),
            reverse_proxy,
            mnemonic_codes,
        },
    });
    start_reaper(data.clone());

    HttpServer::new(move || {
        App::new()
            .app_data(data.clone())
            .wrap(if data.config.reverse_proxy {
                Logger::new(r#"%{r}a "%r" %s %b "%{Referer}i" "%{User-Agent}i" %T"#)
            } else {
                Logger::default()
            })
            .service(handle_download)
            .service(handle_upload)
            .service(check_upload_password)
            .service(upload_limits)
            .service(actix_files::Files::new("/", static_dir.clone()).index_file("index.html"))
    })
    .bind((
        if reverse_proxy {
            "127.0.0.1"
        } else {
            "0.0.0.0"
        },
        port,
    ))?
    .run()
    .await?;
    Ok(())
}

fn start_reaper(data: web::Data<AppState>) {
    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.file_store.write().await.remove_expired_files().await {
                    error!("Error reaping expired files: {}", e);
                }
            }
        });
    });
}