refactor config variables, add upload password

src/store.rs

@@ -1,6 +1,10 @@
-use std::{collections::HashMap, io::ErrorKind, path::PathBuf, str::FromStr};
+use std::{
+    collections::HashMap,
+    io::ErrorKind,
+    path::{Path, PathBuf},
+};
 
-use log::{debug, error, info, warn};
+use log::{debug, error, info};
 use rand::{
     distributions::{Alphanumeric, DistString},
     thread_rng, Rng,
@@ -13,17 +17,13 @@ use tokio::{
 };
 
 const STATE_FILE_NAME: &str = "files.json";
-const DEFAULT_STORAGE_DIR: &str = "storage";
-const DEFAULT_MAX_LIFETIME: u32 = 30;
-const GIGA: u64 = 1024 * 1024 * 1024;
-const DEFAULT_MAX_UPLOAD_SIZE: u64 = 16 * GIGA;
-const DEFAULT_MAX_STORAGE_SIZE: u64 = 64 * GIGA;
+const MAX_STORAGE_FILES: usize = 1024;
 
-pub fn gen_storage_code() -> String {
-    if std::env::var("TRANSBEAM_MNEMONIC_CODES").as_deref() == Ok("false") {
-        Alphanumeric.sample_string(&mut thread_rng(), 8)
-    } else {
+pub fn gen_storage_code(use_mnemonic: bool) -> String {
+    if use_mnemonic {
         mnemonic::to_string(thread_rng().gen::<[u8; 4]>())
+    } else {
+        Alphanumeric.sample_string(&mut thread_rng(), 8)
     }
 }
 
@@ -33,32 +33,6 @@ pub fn is_valid_storage_code(s: &str) -> bool {
         .all(|c| c.is_ascii_alphanumeric() || c == &b'-')
 }
 
-pub(crate) fn storage_dir() -> PathBuf {
-    PathBuf::from(
-        std::env::var("TRANSBEAM_STORAGE_DIR")
-            .unwrap_or_else(|_| String::from(DEFAULT_STORAGE_DIR)),
-    )
-}
-
-fn parse_env_var<T: FromStr>(var: &str, default: T) -> T {
-    std::env::var(var)
-        .ok()
-        .and_then(|val| val.parse::<T>().ok())
-        .unwrap_or(default)
-}
-
-pub(crate) fn max_lifetime() -> u32 {
-    parse_env_var("TRANSBEAM_MAX_LIFETIME", DEFAULT_MAX_LIFETIME)
-}
-
-pub(crate) fn max_single_size() -> u64 {
-    parse_env_var("TRANSBEAM_MAX_UPLOAD_SIZE", DEFAULT_MAX_UPLOAD_SIZE)
-}
-
-pub(crate) fn max_total_size() -> u64 {
-    parse_env_var("TRANSBEAM_MAX_STORAGE_SIZE", DEFAULT_MAX_STORAGE_SIZE)
-}
-
 #[derive(Clone, Deserialize, Serialize)]
 pub struct StoredFile {
     pub name: String,
@@ -110,13 +84,13 @@ pub(crate) mod timestamp {
     }
 }
 
-async fn is_valid_entry(key: &str, info: &StoredFile) -> bool {
+async fn is_valid_entry(key: &str, info: &StoredFile, storage_dir: &Path) -> bool {
     if info.expiry < OffsetDateTime::now_utc() {
         info!("File {} has expired", key);
         return false;
     }
 
-    let file = if let Ok(f) = File::open(storage_dir().join(&key)).await {
+    let file = if let Ok(f) = File::open(storage_dir.join(&key)).await {
         f
     } else {
         error!(
@@ -141,10 +115,35 @@ async fn is_valid_entry(key: &str, info: &StoredFile) -> bool {
     true
 }
 
-pub(crate) struct FileStore(HashMap<String, StoredFile>);
+async fn delete_file_if_exists(file: &PathBuf) -> std::io::Result<()> {
+    if let Err(e) = tokio::fs::remove_file(file).await {
+        if e.kind() != ErrorKind::NotFound {
+            error!("Failed to delete file {}: {}", file.to_string_lossy(), e);
+            return Err(e);
+        }
+    }
+    Ok(())
+}
+
+#[derive(thiserror::Error, Debug)]
+pub enum FileAddError {
+    #[error("Failed to write metadata to filesystem")]
+    FileSystem(#[from] std::io::Error),
+    #[error("File was too large, available space is {0} bytes")]
+    TooBig(u64),
+    #[error("File store is full")]
+    Full,
+}
+
+pub struct FileStore {
+    files: HashMap<String, StoredFile>,
+    storage_dir: PathBuf,
+    max_storage_size: u64,
+}
+
 impl FileStore {
-    pub(crate) async fn load() -> std::io::Result<Self> {
-        let open_result = File::open(storage_dir().join(STATE_FILE_NAME)).await;
+    pub(crate) async fn load(storage_dir: PathBuf, max_storage_size: u64) -> std::io::Result<Self> {
+        let open_result = File::open(storage_dir.join(STATE_FILE_NAME)).await;
         match open_result {
             Ok(mut f) => {
                 let mut buf = String::new();
@@ -160,22 +159,28 @@ impl FileStore {
                         error!("Invalid key in persistent storage: {}", key);
                         continue;
                     }
-                    if is_valid_entry(&key, &info).await {
+                    if is_valid_entry(&key, &info, &storage_dir).await {
                         filtered.insert(key, info);
                     } else {
                         info!("Deleting file {}", key);
-                        if let Err(e) = tokio::fs::remove_file(storage_dir().join(&key)).await {
-                            warn!("Failed to delete file {}: {}", key, e);
-                        }
+                        delete_file_if_exists(&storage_dir.join(&key)).await?;
                     }
                 }
-                let mut loaded = Self(filtered);
+                let mut loaded = Self {
+                    files: filtered,
+                    storage_dir,
+                    max_storage_size,
+                };
                 loaded.save().await?;
                 Ok(loaded)
             }
             Err(e) => {
                 if let ErrorKind::NotFound = e.kind() {
-                    Ok(Self(HashMap::new()))
+                    Ok(Self {
+                        files: HashMap::new(),
+                        storage_dir,
+                        max_storage_size,
+                    })
                 } else {
                     Err(e)
                 }
@@ -184,17 +189,25 @@ impl FileStore {
     }
 
     fn total_size(&self) -> u64 {
-        self.0.iter().fold(0, |acc, (_, f)| acc + f.size)
+        self.files.iter().fold(0, |acc, (_, f)| acc + f.size)
+    }
+
+    pub fn available_size(&self) -> u64 {
+        self.max_storage_size.saturating_sub(self.total_size())
     }
 
     async fn save(&mut self) -> std::io::Result<()> {
-        info!("saving updated state: {} entries", self.0.len());
-        File::create(storage_dir().join(STATE_FILE_NAME))
+        info!("saving updated state: {} entries", self.files.len());
+        File::create(self.storage_dir.join(STATE_FILE_NAME))
             .await?
-            .write_all(&serde_json::to_vec_pretty(&self.0)?)
+            .write_all(&serde_json::to_vec_pretty(&self.files)?)
             .await
     }
 
+    pub fn full(&self) -> bool {
+        self.available_size() == 0 || self.files.len() >= MAX_STORAGE_FILES
+    }
+
     /// Attempts to add a file to the store. Returns an I/O error if
     /// something's broken, or a u64 of the maximum allowed file size
     /// if the file was too big, or a unit if everything worked.
@@ -202,37 +215,42 @@ impl FileStore {
         &mut self,
         key: String,
         file: StoredFile,
-    ) -> std::io::Result<Result<(), u64>> {
-        let remaining_size = max_total_size().saturating_sub(self.total_size());
-        let allowed_size = std::cmp::min(remaining_size, max_single_size());
-        if file.size > allowed_size {
-            return Ok(Err(allowed_size));
+    ) -> Result<(), FileAddError> {
+        if self.full() {
+            return Err(FileAddError::Full);
         }
-        self.0.insert(key, file);
-        self.save().await.map(Ok)
+        let available_size = self.available_size();
+        if file.size > available_size {
+            return Err(FileAddError::TooBig(available_size));
+        }
+        self.files.insert(key, file);
+        self.save().await?;
+        Ok(())
     }
 
     pub(crate) fn lookup_file(&self, key: &str) -> Option<StoredFile> {
-        self.0.get(key).cloned()
+        self.files.get(key).cloned()
     }
 
     pub(crate) async fn remove_file(&mut self, key: &str) -> std::io::Result<()> {
         debug!("removing entry {} from state", key);
-        self.0.remove(key);
-        self.save().await
+        self.files.remove(key);
+        self.save().await?;
+        if is_valid_storage_code(key) {
+            delete_file_if_exists(&self.storage_dir.join(key)).await?;
+        }
+        Ok(())
     }
 
     pub(crate) async fn remove_expired_files(&mut self) -> std::io::Result<()> {
         info!("Checking for expired files");
         let now = OffsetDateTime::now_utc();
-        for (key, file) in std::mem::take(&mut self.0).into_iter() {
+        for (key, file) in std::mem::take(&mut self.files).into_iter() {
             if file.expiry > now {
-                self.0.insert(key, file);
+                self.files.insert(key, file);
             } else {
                 info!("Deleting expired file {}", key);
-                if let Err(e) = tokio::fs::remove_file(storage_dir().join(&key)).await {
-                    warn!("Failed to delete expired file {}: {}", key, e);
-                }
+                delete_file_if_exists(&self.storage_dir.join(&key)).await?;
             }
         }
         self.save().await