refactor config variables, add upload password

src/main.rs

@@ -3,20 +3,48 @@ mod store;
 mod upload;
 mod zip;
 
-use std::{fs::File, path::PathBuf};
+use std::{fmt::Debug, fs::File, path::PathBuf, str::FromStr};
 
 use actix_web::{
-    get, middleware::Logger, web, App, HttpRequest, HttpResponse, HttpServer, Responder,
+    get, middleware::Logger, post, web, App, HttpRequest, HttpResponse, HttpServer, Responder,
 };
 use actix_web_actors::ws;
-use log::error;
-use serde::Deserialize;
+use bytesize::ByteSize;
+use log::{error, warn};
+use serde::{Deserialize, Serialize};
 use store::FileStore;
 use tokio::sync::RwLock;
 
 const APP_NAME: &str = "transbeam";
 
-type AppData = web::Data<RwLock<FileStore>>;
+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,
+}
+
+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 {
@@ -27,16 +55,18 @@ struct DownloadRequest {
 async fn handle_download(
     req: HttpRequest,
     download: web::Query<DownloadRequest>,
-    data: AppData,
+    data: web::Data<AppState>,
 ) -> actix_web::Result<HttpResponse> {
+    let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
     let code = &download.code;
     if !store::is_valid_storage_code(code) {
+        log_auth_failure(&ip_addr);
         return Ok(HttpResponse::NotFound().finish());
     }
-    let data = data.read().await;
-    let info = data.lookup_file(code);
+    let file_store = data.file_store.read().await;
+    let info = file_store.lookup_file(code);
     if let Some(info) = info {
-        let storage_path = store::storage_dir().join(code);
+        let storage_path = data.config.storage_dir.join(code);
         let file = File::open(&storage_path)?;
         Ok(download::DownloadingFile {
             file,
@@ -45,50 +75,153 @@ async fn handle_download(
         }
         .into_response(&req))
     } else {
+        log_auth_failure(&ip_addr);
         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)
+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 data: AppData = web::Data::new(RwLock::new(FileStore::load().await?));
-    start_reaper(data.clone());
+    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 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);
+    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,
+        },
+    });
+    start_reaper(data.clone());
 
     HttpServer::new(move || {
         App::new()
             .app_data(data.clone())
-            .wrap(Logger::default())
-            .service(handle_upload)
+            .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(("127.0.0.1", port))?
+    .bind((
+        if reverse_proxy {
+            "127.0.0.1"
+        } else {
+            "0.0.0.0"
+        },
+        port,
+    ))?
     .run()
     .await?;
     Ok(())
 }
 
-fn start_reaper(data: AppData) {
+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.write().await.remove_expired_files().await {
+                if let Err(e) = data.file_store.write().await.remove_expired_files().await {
                     error!("Error reaping expired files: {}", e);
                 }
             }