move state to jsondb

src/store.rs

@@ -1,7 +1,7 @@
 use std::{
     collections::HashMap,
     io::ErrorKind,
-    path::{Path, PathBuf},
+    path::{Path, PathBuf}, ops::DerefMut,
 };
 
 use log::{debug, error, info};
@@ -13,15 +13,11 @@ use serde::{Deserialize, Serialize};
 use serde_with::skip_serializing_none;
 use serde_with::{serde_as, FromInto, PickFirst};
 use time::OffsetDateTime;
-use tokio::{
-    fs::File,
-    io::{AsyncReadExt, AsyncWriteExt},
-};
+use tokio::fs::File;
 
 use crate::upload::UploadedFile;
 use crate::zip::FileSet;
 
-const STATE_FILE_NAME: &str = "files.json";
 const MAX_STORAGE_FILES: usize = 1024;
 
 pub fn gen_storage_code(use_mnemonic: bool) -> String {
@@ -40,7 +36,7 @@ pub fn is_valid_storage_code(s: &str) -> bool {
 
 #[serde_as]
 #[skip_serializing_none]
-#[derive(Clone, Deserialize, Serialize)]
+#[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct StoredFile {
     pub name: String,
     pub size: u64,
@@ -53,6 +49,9 @@ pub struct StoredFile {
     pub contents: Option<FileSet>,
 }
 
+#[derive(Debug, Default, Deserialize, Serialize)]
+pub struct StoredFiles(pub HashMap<String, StoredFile>);
+
 async fn is_valid_entry(key: &str, info: &StoredFile, storage_dir: &Path) -> bool {
     if info.expiry < OffsetDateTime::now_utc() {
         info!("File {} has expired", key);
@@ -104,124 +103,87 @@ pub enum FileAddError {
     Full,
 }
 
-pub struct FileStore {
-    files: HashMap<String, StoredFile>,
-    storage_dir: PathBuf,
-    max_storage_size: u64,
-}
-
-impl FileStore {
-    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();
-                f.read_to_string(&mut buf).await?;
-                let map: HashMap<String, StoredFile> = serde_json::from_str(&buf)?;
-                info!("Loaded {} file entries from persistent storage", map.len());
-                let mut filtered: HashMap<String, StoredFile> = HashMap::new();
-                for (key, info) in map.into_iter() {
-                    // Handle this case separately, because we don't
-                    // want to try to delete it if it's not the sort
-                    // of path we're expecting
-                    if !is_valid_storage_code(&key) {
-                        error!("Invalid key in persistent storage: {}", key);
-                        continue;
-                    }
-                    if is_valid_entry(&key, &info, &storage_dir).await {
-                        filtered.insert(key, info);
-                    } else {
-                        info!("Deleting file {}", key);
-                        delete_file_if_exists(&storage_dir.join(&key)).await?;
-                    }
-                }
-                let mut loaded = Self {
-                    files: filtered,
-                    storage_dir,
-                    max_storage_size,
-                };
-                loaded.save().await?;
-                Ok(loaded)
+impl crate::AppData {
+    // This must only be run at startup, or it will delete in-progress
+    // uploads.
+    pub async fn cleanup(&self) -> std::io::Result<()> {
+        let mut store = self.state.write().await;
+        let old = std::mem::take(store.deref_mut());
+        for (key, info) in old.0.into_iter() {
+            // Handle this case separately, because we don't
+            // want to try to delete it if it's not the sort
+            // of path we're expecting
+            if !is_valid_storage_code(&key) {
+                error!("Invalid key in persistent storage: {}", key);
+                continue;
             }
-            Err(e) => {
-                if let ErrorKind::NotFound = e.kind() {
-                    Ok(Self {
-                        files: HashMap::new(),
-                        storage_dir,
-                        max_storage_size,
-                    })
-                } else {
-                    Err(e)
-                }
+            if is_valid_entry(&key, &info, &self.config.storage_dir).await {
+                store.0.insert(key, info);
+            } else {
+                info!("Deleting file {}", key);
+                delete_file_if_exists(&self.config.storage_dir.join(&key)).await?;
             }
         }
-    }
-
-    fn total_size(&self) -> u64 {
-        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.files.len());
-        File::create(self.storage_dir.join(STATE_FILE_NAME))
-            .await?
-            .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
+        Ok(())
     }
 
     /// 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.
-    pub(crate) async fn add_file(
-        &mut self,
+    pub async fn add_file(
+        &self,
         key: String,
         file: StoredFile,
     ) -> Result<(), FileAddError> {
-        if self.full() {
+        let mut store = self.state.write().await;
+        if store.full(self.config.max_storage_size) {
             return Err(FileAddError::Full);
         }
-        let available_size = self.available_size();
+        let available_size = store.available_size(self.config.max_storage_size);
         if file.size > available_size {
             return Err(FileAddError::TooBig(available_size));
         }
-        self.files.insert(key, file);
-        self.save().await?;
+        store.0.insert(key, file);
         Ok(())
     }
 
-    pub(crate) fn lookup_file(&self, key: &str) -> Option<StoredFile> {
-        self.files.get(key).cloned()
-    }
-
-    pub(crate) async fn remove_file(&mut self, key: &str) -> std::io::Result<()> {
+    pub async fn remove_file(&self, key: &str) -> std::io::Result<()> {
         debug!("removing entry {} from state", key);
-        self.files.remove(key);
-        self.save().await?;
+        let mut store = self.state.write().await;
+        store.0.remove(key);
         if is_valid_storage_code(key) {
-            delete_file_if_exists(&self.storage_dir.join(key)).await?;
+            delete_file_if_exists(&self.config.storage_dir.join(key)).await?;
         }
         Ok(())
     }
 
-    pub(crate) async fn remove_expired_files(&mut self) -> std::io::Result<()> {
+    pub async fn remove_expired_files(&self) -> std::io::Result<()> {
         info!("Checking for expired files");
         let now = OffsetDateTime::now_utc();
-        for (key, file) in std::mem::take(&mut self.files).into_iter() {
+        let mut store = self.state.write().await;
+        let old = std::mem::take(store.deref_mut());
+        for (key, file) in old.0.into_iter() {
             if file.expiry > now {
-                self.files.insert(key, file);
+                store.0.insert(key, file);
             } else {
                 info!("Deleting expired file {}", key);
-                delete_file_if_exists(&self.storage_dir.join(&key)).await?;
+                delete_file_if_exists(&self.config.storage_dir.join(&key)).await?;
             }
         }
-        self.save().await
+        Ok(())
+    }
+}
+
+impl StoredFiles {
+    fn total_size(&self) -> u64 {
+        self.0.iter().fold(0, |acc, (_, f)| acc + f.size)
+    }
+
+    pub fn available_size(&self, max_storage_size: u64) -> u64 {
+        max_storage_size.saturating_sub(self.total_size())
+    }
+
+    pub fn full(&self, max_storage_size: u64) -> bool {
+        self.available_size(max_storage_size) == 0 || self.0.len() >= MAX_STORAGE_FILES
     }
 }