245 lines
9.6 KiB
Rust
245 lines
9.6 KiB
Rust
use redb::{ReadableTable};
|
|
use crate::error::DBError;
|
|
use super::*;
|
|
|
|
impl Storage {
|
|
pub fn flushdb(&self) -> Result<(), DBError> {
|
|
let write_txn = self.db.begin_write()?;
|
|
{
|
|
let mut types_table = write_txn.open_table(TYPES_TABLE)?;
|
|
let mut strings_table = write_txn.open_table(STRINGS_TABLE)?;
|
|
let mut hashes_table = write_txn.open_table(HASHES_TABLE)?;
|
|
let mut lists_table = write_txn.open_table(LISTS_TABLE)?;
|
|
let mut streams_meta_table = write_txn.open_table(STREAMS_META_TABLE)?;
|
|
let mut streams_data_table = write_txn.open_table(STREAMS_DATA_TABLE)?;
|
|
let mut expiration_table = write_txn.open_table(EXPIRATION_TABLE)?;
|
|
|
|
// inefficient, but there is no other way
|
|
let keys: Vec<String> = types_table.iter()?.map(|item| item.unwrap().0.value().to_string()).collect();
|
|
for key in keys {
|
|
types_table.remove(key.as_str())?;
|
|
}
|
|
let keys: Vec<String> = strings_table.iter()?.map(|item| item.unwrap().0.value().to_string()).collect();
|
|
for key in keys {
|
|
strings_table.remove(key.as_str())?;
|
|
}
|
|
let keys: Vec<(String, String)> = hashes_table
|
|
.iter()?
|
|
.map(|item| {
|
|
let binding = item.unwrap();
|
|
let (k, f) = binding.0.value();
|
|
(k.to_string(), f.to_string())
|
|
})
|
|
.collect();
|
|
for (key, field) in keys {
|
|
hashes_table.remove((key.as_str(), field.as_str()))?;
|
|
}
|
|
let keys: Vec<String> = lists_table.iter()?.map(|item| item.unwrap().0.value().to_string()).collect();
|
|
for key in keys {
|
|
lists_table.remove(key.as_str())?;
|
|
}
|
|
let keys: Vec<String> = streams_meta_table.iter()?.map(|item| item.unwrap().0.value().to_string()).collect();
|
|
for key in keys {
|
|
streams_meta_table.remove(key.as_str())?;
|
|
}
|
|
let keys: Vec<(String,String)> = streams_data_table.iter()?.map(|item| {
|
|
let binding = item.unwrap();
|
|
let (key, field) = binding.0.value();
|
|
(key.to_string(), field.to_string())
|
|
}).collect();
|
|
for (key, field) in keys {
|
|
streams_data_table.remove((key.as_str(), field.as_str()))?;
|
|
}
|
|
let keys: Vec<String> = expiration_table.iter()?.map(|item| item.unwrap().0.value().to_string()).collect();
|
|
for key in keys {
|
|
expiration_table.remove(key.as_str())?;
|
|
}
|
|
}
|
|
write_txn.commit()?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn get_key_type(&self, key: &str) -> Result<Option<String>, DBError> {
|
|
let read_txn = self.db.begin_read()?;
|
|
let table = read_txn.open_table(TYPES_TABLE)?;
|
|
|
|
// Before returning type, check for expiration
|
|
if let Some(type_val) = table.get(key)? {
|
|
if type_val.value() == "string" {
|
|
let expiration_table = read_txn.open_table(EXPIRATION_TABLE)?;
|
|
if let Some(expires_at) = expiration_table.get(key)? {
|
|
if now_in_millis() > expires_at.value() as u128 {
|
|
// The key is expired, so it effectively has no type
|
|
return Ok(None);
|
|
}
|
|
}
|
|
}
|
|
Ok(Some(type_val.value().to_string()))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
// ✅ ENCRYPTION APPLIED: Value is encrypted/decrypted
|
|
pub fn get(&self, key: &str) -> Result<Option<String>, DBError> {
|
|
let read_txn = self.db.begin_read()?;
|
|
|
|
let types_table = read_txn.open_table(TYPES_TABLE)?;
|
|
match types_table.get(key)? {
|
|
Some(type_val) if type_val.value() == "string" => {
|
|
// Check expiration first (unencrypted)
|
|
let expiration_table = read_txn.open_table(EXPIRATION_TABLE)?;
|
|
if let Some(expires_at) = expiration_table.get(key)? {
|
|
if now_in_millis() > expires_at.value() as u128 {
|
|
drop(read_txn);
|
|
self.del(key.to_string())?;
|
|
return Ok(None);
|
|
}
|
|
}
|
|
|
|
// Get and decrypt value
|
|
let strings_table = read_txn.open_table(STRINGS_TABLE)?;
|
|
match strings_table.get(key)? {
|
|
Some(data) => {
|
|
let decrypted = self.decrypt_if_needed(data.value())?;
|
|
let value = String::from_utf8(decrypted)?;
|
|
Ok(Some(value))
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
_ => Ok(None),
|
|
}
|
|
}
|
|
|
|
// ✅ ENCRYPTION APPLIED: Value is encrypted before storage
|
|
pub fn set(&self, key: String, value: String) -> Result<(), DBError> {
|
|
let write_txn = self.db.begin_write()?;
|
|
|
|
{
|
|
let mut types_table = write_txn.open_table(TYPES_TABLE)?;
|
|
types_table.insert(key.as_str(), "string")?;
|
|
|
|
let mut strings_table = write_txn.open_table(STRINGS_TABLE)?;
|
|
// Only encrypt the value, not expiration
|
|
let encrypted = self.encrypt_if_needed(value.as_bytes())?;
|
|
strings_table.insert(key.as_str(), encrypted.as_slice())?;
|
|
|
|
// Remove any existing expiration since this is a regular SET
|
|
let mut expiration_table = write_txn.open_table(EXPIRATION_TABLE)?;
|
|
expiration_table.remove(key.as_str())?;
|
|
}
|
|
|
|
write_txn.commit()?;
|
|
Ok(())
|
|
}
|
|
|
|
// ✅ ENCRYPTION APPLIED: Value is encrypted before storage
|
|
pub fn setx(&self, key: String, value: String, expire_ms: u128) -> Result<(), DBError> {
|
|
let write_txn = self.db.begin_write()?;
|
|
|
|
{
|
|
let mut types_table = write_txn.open_table(TYPES_TABLE)?;
|
|
types_table.insert(key.as_str(), "string")?;
|
|
|
|
let mut strings_table = write_txn.open_table(STRINGS_TABLE)?;
|
|
// Only encrypt the value
|
|
let encrypted = self.encrypt_if_needed(value.as_bytes())?;
|
|
strings_table.insert(key.as_str(), encrypted.as_slice())?;
|
|
|
|
// Store expiration separately (unencrypted)
|
|
let mut expiration_table = write_txn.open_table(EXPIRATION_TABLE)?;
|
|
let expires_at = expire_ms + now_in_millis();
|
|
expiration_table.insert(key.as_str(), &(expires_at as u64))?;
|
|
}
|
|
|
|
write_txn.commit()?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn del(&self, key: String) -> Result<(), DBError> {
|
|
let write_txn = self.db.begin_write()?;
|
|
|
|
{
|
|
let mut types_table = write_txn.open_table(TYPES_TABLE)?;
|
|
let mut strings_table = write_txn.open_table(STRINGS_TABLE)?;
|
|
let mut hashes_table: redb::Table<(&str, &str), &[u8]> = write_txn.open_table(HASHES_TABLE)?;
|
|
let mut lists_table = write_txn.open_table(LISTS_TABLE)?;
|
|
|
|
// Remove from type table
|
|
types_table.remove(key.as_str())?;
|
|
|
|
// Remove from strings table
|
|
strings_table.remove(key.as_str())?;
|
|
|
|
// Remove all hash fields for this key
|
|
let mut to_remove = Vec::new();
|
|
let mut iter = hashes_table.iter()?;
|
|
while let Some(entry) = iter.next() {
|
|
let entry = entry?;
|
|
let (hash_key, field) = entry.0.value();
|
|
if hash_key == key.as_str() {
|
|
to_remove.push((hash_key.to_string(), field.to_string()));
|
|
}
|
|
}
|
|
drop(iter);
|
|
|
|
for (hash_key, field) in to_remove {
|
|
hashes_table.remove((hash_key.as_str(), field.as_str()))?;
|
|
}
|
|
|
|
// Remove from lists table
|
|
lists_table.remove(key.as_str())?;
|
|
|
|
// Also remove expiration
|
|
let mut expiration_table = write_txn.open_table(EXPIRATION_TABLE)?;
|
|
expiration_table.remove(key.as_str())?;
|
|
}
|
|
|
|
write_txn.commit()?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn keys(&self, pattern: &str) -> Result<Vec<String>, DBError> {
|
|
let read_txn = self.db.begin_read()?;
|
|
let table = read_txn.open_table(TYPES_TABLE)?;
|
|
|
|
let mut keys = Vec::new();
|
|
let mut iter = table.iter()?;
|
|
while let Some(entry) = iter.next() {
|
|
let key = entry?.0.value().to_string();
|
|
if pattern == "*" || super::storage_extra::glob_match(pattern, &key) {
|
|
keys.push(key);
|
|
}
|
|
}
|
|
|
|
Ok(keys)
|
|
}
|
|
}
|
|
|
|
impl Storage {
|
|
pub fn dbsize(&self) -> Result<i64, DBError> {
|
|
let read_txn = self.db.begin_read()?;
|
|
let types_table = read_txn.open_table(TYPES_TABLE)?;
|
|
let expiration_table = read_txn.open_table(EXPIRATION_TABLE)?;
|
|
|
|
let mut count: i64 = 0;
|
|
let mut iter = types_table.iter()?;
|
|
while let Some(entry) = iter.next() {
|
|
let entry = entry?;
|
|
let key = entry.0.value();
|
|
let ty = entry.1.value();
|
|
|
|
if ty == "string" {
|
|
if let Some(expires_at) = expiration_table.get(key)? {
|
|
if now_in_millis() > expires_at.value() as u128 {
|
|
// Skip logically expired string keys
|
|
continue;
|
|
}
|
|
}
|
|
}
|
|
count += 1;
|
|
}
|
|
Ok(count)
|
|
}
|
|
} |