1 files changed, 34 insertions, 20 deletions

diff --git a/shared/src/store.rs b/shared/src/store.rs
index 44e5e3d..7130021 100644
--- a/shared/src/store.rs
+++ b/shared/src/store.rs
@@ -2,35 +2,49 @@ use crate::packets::Resource;
 use anyhow::Result;
 use redb::{Database, TableDefinition};
 use sha2::{Digest, Sha256};
-use std::path::Path;
+use std::{collections::HashMap, path::Path, sync::Mutex};
 
 const T_ENTRIES: TableDefinition<[u8; 32], &[u8]> = TableDefinition::new("e");
 
-pub struct ResourceStore {
-    db: Database,
+pub enum ResourceStore {
+    Redb(Database),
+    Memory(Mutex<HashMap<Resource, Vec<u8>>>),
 }
 impl ResourceStore {
-    pub fn new(path: &Path) -> Result<Self> {
-        Ok(Self {
-            db: Database::create(path)?,
-        })
+    pub fn new_persistent(path: &Path) -> Result<Self> {
+        Ok(Self::Redb(Database::create(path)?))
+    }
+    pub fn new_memory() -> Self {
+        Self::Memory(HashMap::new().into())
     }
     pub fn get(&self, key: Resource) -> Result<Option<Vec<u8>>> {
-        let txn = self.db.begin_read()?;
-        let ent = txn.open_table(T_ENTRIES)?;
-        match ent.get(key.0)? {
-            Some(x) => Ok(Some(x.value().to_vec())),
-            None => Ok(None),
+        match self {
+            ResourceStore::Redb(database) => {
+                let txn = database.begin_read()?;
+                let ent = txn.open_table(T_ENTRIES)?;
+                match ent.get(key.0)? {
+                    Some(x) => Ok(Some(x.value().to_vec())),
+                    None => Ok(None),
+                }
+            }
+            ResourceStore::Memory(map) => Ok(map.lock().unwrap().get(&key).map(|x| x.to_vec())),
         }
     }
-    pub fn set(&self, value: &[u8]) -> Result<()> {
-        let key = sha256(value);
-        let txn = self.db.begin_write()?;
-        let mut ent = txn.open_table(T_ENTRIES)?;
-        ent.insert(key, value)?;
-        drop(ent);
-        txn.commit()?;
-        Ok(())
+    pub fn set(&self, value: &[u8]) -> Result<Resource> {
+        let key = Resource(sha256(value));
+        match self {
+            ResourceStore::Redb(database) => {
+                let txn = database.begin_write()?;
+                let mut ent = txn.open_table(T_ENTRIES)?;
+                ent.insert(key.0, value)?;
+                drop(ent);
+                txn.commit()?;
+            }
+            ResourceStore::Memory(map) => {
+                map.lock().unwrap().insert(key, value.to_vec());
+            }
+        }
+        Ok(key)
     }
 }