transbeam/src/main.rs

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

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

use actix_web::{
    error::InternalError, get, middleware::Logger, post, web, App, HttpRequest, HttpResponse,
    HttpServer, Responder,
};
use actix_web_actors::ws;
use askama::Template;
use bytesize::ByteSize;
use log::{error, warn};
use serde::{Deserialize, Serialize};
use store::{FileStore, StoredFile};
use tokio::{fs::File, sync::RwLock};

const APP_NAME: &str = "transbeam";

const DATE_DISPLAY_FORMAT: &[time::format_description::FormatItem] =
    time::macros::format_description!("[year]-[month]-[day]");

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

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

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(Template)]
#[template(path = "index.html")]
struct IndexPage<'a> {
    cachebuster: &'a str,
}

#[get("/")]
async fn index(data: web::Data<AppState>) -> impl Responder {
    HttpResponse::Ok().body(
        IndexPage {
            cachebuster: &data.config.cachebuster,
        }
        .render()
        .unwrap(),
    )
}

#[derive(Deserialize)]
struct DownloadRequest {
    code: String,
    download: Option<download::DownloadSelection>,
}

#[derive(Template)]
#[template(path = "download.html")]
struct DownloadPage<'a> {
    info: DownloadInfo,
    cachebuster: &'a str,
}

#[derive(Serialize)]
struct DownloadInfo {
    file: StoredFile,
    code: String,
    available: u64,
    offsets: Option<Vec<u64>>,
}

#[get("/download")]
async fn handle_download(
    req: HttpRequest,
    query: web::Query<DownloadRequest>,
    data: web::Data<AppState>,
) -> actix_web::Result<HttpResponse> {
    let code = &query.code;
    if !store::is_valid_storage_code(code) {
        return not_found(req, data, true);
    }
    let info = data.file_store.read().await.lookup_file(code);
    let info = if let Some(i) = info {
        i
    } else {
        return not_found(req, data, true);
    };

    let storage_path = data.config.storage_dir.join(code);
    let file = File::open(&storage_path).await?;
    if let Some(selection) = query.download {
        if let download::DownloadSelection::One(n) = selection {
            if let Some(ref files) = info.contents {
                if n >= files.len() {
                    return not_found(req, data, false);
                }
            } else {
                return not_found(req, data, false);
            }
        }
        Ok(download::DownloadingFile {
            file: file.into_std().await,
            storage_path,
            info,
            selection,
        }
        .into_response(&req))
    } else {
        let offsets = info.contents.as_deref().map(zip::file_data_offsets);
        Ok(HttpResponse::Ok().body(
            DownloadPage {
                info: DownloadInfo {
                    file: info,
                    code: code.clone(),
                    available: file.metadata().await?.len(),
                    offsets,
                },
                cachebuster: &data.config.cachebuster,
            }
            .render()
            .unwrap(),
        ))
    }
}

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

#[get("/info")]
async fn download_info(
    req: HttpRequest,
    query: web::Query<InfoQuery>,
    data: web::Data<AppState>,
) -> actix_web::Result<impl Responder> {
    let code = &query.code;
    if !store::is_valid_storage_code(code) {
        return not_found(req, data, true);
    }
    let info = data.file_store.read().await.lookup_file(code);
    let info = if let Some(i) = info {
        i
    } else {
        return not_found(req, data, true);
    };

    let storage_path = data.config.storage_dir.join(code);
    let offsets = info.contents.as_deref().map(zip::file_data_offsets);
    Ok(web::Json(DownloadInfo {
        file: info,
        code: code.clone(),
        available: File::open(&storage_path).await?.metadata().await?.len(),
        offsets,
    }))
}

#[derive(Template)]
#[template(path = "404.html")]
struct NotFoundPage<'a> {
    cachebuster: &'a str,
}

fn not_found<T>(req: HttpRequest, data: web::Data<AppState>, report: bool) -> actix_web::Result<T> {
    if report {
        let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
        log_auth_failure(&ip_addr);
    }
    Err(InternalError::from_response(
        "Download not found",
        HttpResponse::NotFound().body(
            NotFoundPage {
                cachebuster: &data.config.cachebuster,
            }
            .render()
            .unwrap(),
        ),
    )
    .into())
}

#[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 cachebuster: String = env_or_else("TRANSBEAM_CACHEBUSTER", String::new);

    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,
            reverse_proxy,
            mnemonic_codes,
            cachebuster,
        },
    });
    start_reaper(data.clone());

    let server = 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(index)
            .service(handle_download)
            .service(download_info)
            .service(handle_upload)
            .service(check_upload_password)
            .service(upload_limits)
            .service(actix_files::Files::new("/", static_dir.clone()))
    });

    if reverse_proxy {
        server
            .bind((std::net::Ipv4Addr::LOCALHOST, port))?
            .bind((std::net::Ipv6Addr::LOCALHOST, port))?
    } else {
        // Looks like this also picks up IPv4?
        // Binding 0.0.0.0 and :: on the same port fails with an error.
        server.bind((std::net::Ipv6Addr::UNSPECIFIED, 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);
                }
            }
        });
    });
}