move state to jsondb

src/main.rs

@@ -1,4 +1,5 @@
 mod download;
+mod state;
 mod store;
 mod timestamp;
 mod upload;
@@ -16,21 +17,23 @@ use askama_actix::{Template, TemplateToResponse};
 use bytesize::ByteSize;
 use log::{error, warn};
 use serde::{Deserialize, Serialize};
-use store::{FileStore, StoredFile};
-use tokio::{fs::File, sync::RwLock};
+use state::StateDb;
+use store::StoredFile;
+use tokio::fs::File;
 
 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>,
+struct AppData {
+    state: state::StateDb,
     config: Config,
 }
 
 struct Config {
     base_url: String,
+    max_storage_size: u64,
     max_upload_size: u64,
     max_lifetime: u16,
     upload_password: String,
@@ -63,7 +66,7 @@ struct IndexPage {
 }
 
 #[get("/")]
-async fn index(data: web::Data<AppState>) -> impl Responder {
+async fn index(data: web::Data<AppData>) -> impl Responder {
     IndexPage {
         cachebuster: data.config.cachebuster.clone(),
         base_url: data.config.base_url.clone(),
@@ -96,15 +99,16 @@ struct DownloadInfo {
 async fn handle_download(
     req: HttpRequest,
     query: web::Query<DownloadRequest>,
-    data: web::Data<AppState>,
+    data: web::Data<AppData>,
 ) -> 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 store = data.state.read().await;
+    let info = store.0.get(code);
     let info = if let Some(i) = info {
-        i
+        i.clone()
     } else {
         return not_found(req, data, true);
     };
@@ -153,15 +157,16 @@ struct InfoQuery {
 async fn download_info(
     req: HttpRequest,
     query: web::Query<InfoQuery>,
-    data: web::Data<AppState>,
+    data: web::Data<AppData>,
 ) -> 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 store = data.state.read().await;
+    let info = store.0.get(code);
     let info = if let Some(i) = info {
-        i
+        i.clone()
     } else {
         return not_found(req, data, true);
     };
@@ -183,7 +188,7 @@ struct NotFoundPage<'a> {
     base_url: &'a str,
 }
 
-fn not_found<T>(req: HttpRequest, data: web::Data<AppState>, report: bool) -> actix_web::Result<T> {
+fn not_found<T>(req: HttpRequest, data: web::Data<AppData>, report: bool) -> actix_web::Result<T> {
     if report {
         let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
         log_auth_failure(&ip_addr);
@@ -202,9 +207,9 @@ fn not_found<T>(req: HttpRequest, data: web::Data<AppState>, report: bool) -> ac
 async fn handle_upload(
     req: HttpRequest,
     stream: web::Payload,
-    data: web::Data<AppState>,
+    data: web::Data<AppData>,
 ) -> impl Responder {
-    if data.file_store.read().await.full() {
+    if data.state.read().await.full(data.config.max_storage_size) {
         return Ok(HttpResponse::BadRequest().finish());
     }
     let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
@@ -220,7 +225,7 @@ struct UploadPasswordCheck {
 async fn check_upload_password(
     req: HttpRequest,
     body: web::Json<UploadPasswordCheck>,
-    data: web::Data<AppState>,
+    data: web::Data<AppData>,
 ) -> impl Responder {
     let ip_addr = get_ip_addr(&req, data.config.reverse_proxy);
     if body.password != data.config.upload_password {
@@ -239,10 +244,10 @@ struct UploadLimits {
 }
 
 #[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();
+async fn upload_limits(data: web::Data<AppData>) -> impl Responder {
+    let file_store = data.state.read().await;
+    let open = !file_store.full(data.config.max_storage_size);
+    let available_size = file_store.available_size(data.config.max_storage_size);
     let max_size = std::cmp::min(available_size, data.config.max_upload_size);
     web::Json(UploadLimits {
         open,
@@ -304,11 +309,24 @@ async fn main() -> std::io::Result<()> {
     let upload_password: String = env_or_panic("TRANSBEAM_UPLOAD_PASSWORD");
     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?),
+    let state_file: PathBuf = match std::env::var("TRANSBEAM_STATE_FILE") {
+        Ok(v) => v.parse().unwrap_or_else(|_| panic!("Invalid value {} for variable TRANSBEAM_STATE_FILE", v)),
+        Err(_) => {
+            let legacy_state_file = storage_dir.join("files.json");
+            if legacy_state_file.is_file() {
+                legacy_state_file
+            } else {
+                PathBuf::from("transbeam.json")
+            }
+        }
+    };
+
+    let data = web::Data::new(AppData {
+        state: StateDb::load(state_file).await.expect("Failed to load state file"),
         config: Config {
             base_url,
             max_upload_size,
+            max_storage_size,
             max_lifetime,
             upload_password,
             storage_dir,
@@ -317,6 +335,7 @@ async fn main() -> std::io::Result<()> {
             cachebuster,
         },
     });
+    data.cleanup().await?;
     start_reaper(data.clone());
 
     let server = HttpServer::new(move || {
@@ -351,12 +370,12 @@ async fn main() -> std::io::Result<()> {
     Ok(())
 }
 
-fn start_reaper(data: web::Data<AppState>) {
+fn start_reaper(data: web::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.file_store.write().await.remove_expired_files().await {
+                if let Err(e) = data.remove_expired_files().await {
                     error!("Error reaping expired files: {}", e);
                 }
             }