...

src/storage/mod.rs

@@ -0,0 +1,126 @@
+use std::{
+    path::Path,
+    time::{SystemTime, UNIX_EPOCH},
+};
+
+use redb::{Database, TableDefinition};
+use serde::{Deserialize, Serialize};
+
+use crate::crypto::CryptoFactory;
+use crate::error::DBError;
+
+// Re-export modules
+mod storage_basic;
+mod storage_hset;
+mod storage_lists;
+mod storage_extra;
+
+// Re-export implementations
+// Note: These imports are used by the impl blocks in the submodules
+// The compiler shows them as unused because they're not directly used in this file
+// but they're needed for the Storage struct methods to be available
+pub use storage_extra::*;
+
+// Table definitions for different Redis data types
+const TYPES_TABLE: TableDefinition<&str, &str> = TableDefinition::new("types");
+const STRINGS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("strings");
+const HASHES_TABLE: TableDefinition<(&str, &str), &[u8]> = TableDefinition::new("hashes");
+const LISTS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("lists");
+const STREAMS_META_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("streams_meta");
+const STREAMS_DATA_TABLE: TableDefinition<(&str, &str), &[u8]> = TableDefinition::new("streams_data");
+const ENCRYPTED_TABLE: TableDefinition<&str, u8> = TableDefinition::new("encrypted");
+const EXPIRATION_TABLE: TableDefinition<&str, u64> = TableDefinition::new("expiration");
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct StreamEntry {
+    pub fields: Vec<(String, String)>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ListValue {
+    pub elements: Vec<String>,
+}
+
+#[inline]
+pub fn now_in_millis() -> u128 {
+    let start = SystemTime::now();
+    let duration_since_epoch = start.duration_since(UNIX_EPOCH).unwrap();
+    duration_since_epoch.as_millis()
+}
+
+pub struct Storage {
+    db: Database,
+    crypto: Option<CryptoFactory>,
+}
+
+impl Storage {
+    pub fn new(path: impl AsRef<Path>, should_encrypt: bool, master_key: Option<&str>) -> Result<Self, DBError> {
+        let db = Database::create(path)?;
+        
+        // Create tables if they don't exist
+        let write_txn = db.begin_write()?;
+        {
+            let _ = write_txn.open_table(TYPES_TABLE)?;
+            let _ = write_txn.open_table(STRINGS_TABLE)?;
+            let _ = write_txn.open_table(HASHES_TABLE)?;
+            let _ = write_txn.open_table(LISTS_TABLE)?;
+            let _ = write_txn.open_table(STREAMS_META_TABLE)?;
+            let _ = write_txn.open_table(STREAMS_DATA_TABLE)?;
+            let _ = write_txn.open_table(ENCRYPTED_TABLE)?;
+            let _ = write_txn.open_table(EXPIRATION_TABLE)?;
+        }
+        write_txn.commit()?;
+        
+        // Check if database was previously encrypted
+        let read_txn = db.begin_read()?;
+        let encrypted_table = read_txn.open_table(ENCRYPTED_TABLE)?;
+        let was_encrypted = encrypted_table.get("encrypted")?.map(|v| v.value() == 1).unwrap_or(false);
+        drop(read_txn);
+        
+        let crypto = if should_encrypt || was_encrypted {
+            if let Some(key) = master_key {
+                Some(CryptoFactory::new(key.as_bytes()))
+            } else {
+                return Err(DBError("Encryption requested but no master key provided".to_string()));
+            }
+        } else {
+            None
+        };
+        
+        // If we're enabling encryption for the first time, mark it
+        if should_encrypt && !was_encrypted {
+            let write_txn = db.begin_write()?;
+            {
+                let mut encrypted_table = write_txn.open_table(ENCRYPTED_TABLE)?;
+                encrypted_table.insert("encrypted", &1u8)?;
+            }
+            write_txn.commit()?;
+        }
+        
+        Ok(Storage {
+            db,
+            crypto,
+        })
+    }
+    
+    pub fn is_encrypted(&self) -> bool {
+        self.crypto.is_some()
+    }
+
+    // Helper methods for encryption
+    fn encrypt_if_needed(&self, data: &[u8]) -> Result<Vec<u8>, DBError> {
+        if let Some(crypto) = &self.crypto {
+            Ok(crypto.encrypt(data))
+        } else {
+            Ok(data.to_vec())
+        }
+    }
+    
+    fn decrypt_if_needed(&self, data: &[u8]) -> Result<Vec<u8>, DBError> {
+        if let Some(crypto) = &self.crypto {
+            Ok(crypto.decrypt(data)?)
+        } else {
+            Ok(data.to_vec())
+        }
+    }
+}