herocode/herodb
12
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

...

Browse Source
This commit is contained in:
despiegk
2025-08-16 09:29:18 +02:00
parent 5502ff4bc5
commit 0000d82799
13 changed files with 719 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

272
Cargo.lock generated
View File

@@ -17,6 +17,16 @@ version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
[[package]]
name = "aead"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0"
dependencies = [
"crypto-common",
"generic-array",
]
[[package]] [[package]]
name = "anstream" name = "anstream"
version = "0.6.15" version = "0.6.15"
@@ -72,6 +82,17 @@ version = "1.0.86"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da"
[[package]]
name = "async-trait"
version = "0.1.88"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e539d3fca749fcee5236ab05e93a52867dd549cc157c8cb7f99595f3cedffdb5"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]] [[package]]
name = "autocfg" name = "autocfg"
version = "1.3.0" version = "1.3.0"
@@ -108,6 +129,15 @@ version = "2.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de"
[[package]]
name = "block-buffer"
version = "0.10.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
dependencies = [
"generic-array",
]
[[package]] [[package]]
name = "byteorder" name = "byteorder"
version = "1.5.0" version = "1.5.0"
@@ -132,6 +162,41 @@ version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
[[package]]
name = "chacha20"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818"
dependencies = [
"cfg-if",
"cipher",
"cpufeatures",
]
[[package]]
name = "chacha20poly1305"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35"
dependencies = [
"aead",
"chacha20",
"cipher",
"poly1305",
"zeroize",
]
[[package]]
name = "cipher"
version = "0.4.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
dependencies = [
"crypto-common",
"inout",
"zeroize",
]
[[package]] [[package]]
name = "clap" name = "clap"
version = "4.5.20" version = "4.5.20"
@@ -185,7 +250,41 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd"
dependencies = [ dependencies = [
"bytes", "bytes",
"futures-core",
"memchr", "memchr",
"pin-project-lite",
"tokio",
"tokio-util",
]
[[package]]
name = "cpufeatures"
version = "0.2.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
dependencies = [
"libc",
]
[[package]]
name = "crypto-common"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3"
dependencies = [
"generic-array",
"rand_core",
"typenum",
]
[[package]]
name = "digest"
version = "0.10.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
dependencies = [
"block-buffer",
"crypto-common",
] ]
[[package]] [[package]]
@@ -297,6 +396,27 @@ dependencies = [
"slab", "slab",
] ]
[[package]]
name = "generic-array"
version = "0.14.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
dependencies = [
"typenum",
"version_check",
]
[[package]]
name = "getrandom"
version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592"
dependencies = [
"cfg-if",
"libc",
"wasi",
]
[[package]] [[package]]
name = "gimli" name = "gimli"
version = "0.29.0" version = "0.29.0"
@@ -422,6 +542,15 @@ dependencies = [
"icu_properties", "icu_properties",
] ]
[[package]]
name = "inout"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
dependencies = [
"generic-array",
]
[[package]] [[package]]
name = "is_terminal_polyfill" name = "is_terminal_polyfill"
version = "1.70.1" version = "1.70.1"
@@ -501,6 +630,12 @@ dependencies = [
"memchr", "memchr",
] ]
[[package]]
name = "opaque-debug"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381"
[[package]] [[package]]
name = "parking_lot" name = "parking_lot"
version = "0.12.3" version = "0.12.3"
@@ -542,6 +677,17 @@ version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
[[package]]
name = "poly1305"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf"
dependencies = [
"cpufeatures",
"opaque-debug",
"universal-hash",
]
[[package]] [[package]]
name = "potential_utf" name = "potential_utf"
version = "0.1.2" version = "0.1.2"
@@ -551,6 +697,15 @@ dependencies = [
"zerovec", "zerovec",
] ]
[[package]]
name = "ppv-lite86"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
dependencies = [
"zerocopy",
]
[[package]] [[package]]
name = "proc-macro2" name = "proc-macro2"
version = "1.0.86" version = "1.0.86"
@@ -569,6 +724,36 @@ dependencies = [
"proc-macro2", "proc-macro2",
] ]
[[package]]
name = "rand"
version = "0.8.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
dependencies = [
"libc",
"rand_chacha",
"rand_core",
]
[[package]]
name = "rand_chacha"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
dependencies = [
"ppv-lite86",
"rand_core",
]
[[package]]
name = "rand_core"
version = "0.6.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
dependencies = [
"getrandom",
]
[[package]] [[package]]
name = "redb" name = "redb"
version = "2.6.2" version = "2.6.2"
@@ -584,12 +769,18 @@ version = "0.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c580d9cbbe1d1b479e8d67cf9daf6a62c957e6846048408b80b43ac3f6af84cd" checksum = "c580d9cbbe1d1b479e8d67cf9daf6a62c957e6846048408b80b43ac3f6af84cd"
dependencies = [ dependencies = [
"async-trait",
"bytes",
"combine", "combine",
"futures-util",
"itoa", "itoa",
"percent-encoding", "percent-encoding",
"pin-project-lite",
"ryu", "ryu",
"sha1_smol", "sha1_smol",
"socket2 0.4.10", "socket2 0.4.10",
"tokio",
"tokio-util",
"url", "url",
] ]
@@ -601,11 +792,14 @@ dependencies = [
"bincode", "bincode",
"byteorder", "byteorder",
"bytes", "bytes",
"chacha20poly1305",
"clap", "clap",
"futures", "futures",
"rand",
"redb", "redb",
"redis", "redis",
"serde", "serde",
"sha2",
"thiserror", "thiserror",
"tokio", "tokio",
] ]
@@ -663,6 +857,17 @@ version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]]
name = "sha2"
version = "0.10.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
dependencies = [
"cfg-if",
"cpufeatures",
"digest",
]
[[package]] [[package]]
name = "signal-hook-registry" name = "signal-hook-registry"
version = "1.4.2" version = "1.4.2"
@@ -719,6 +924,12 @@ version = "0.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
[[package]]
name = "subtle"
version = "2.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
[[package]] [[package]]
name = "syn" name = "syn"
version = "2.0.69" version = "2.0.69"
@@ -801,12 +1012,41 @@ dependencies = [
"syn", "syn",
] ]
[[package]]
name = "tokio-util"
version = "0.7.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "14307c986784f72ef81c89db7d9e28d6ac26d16213b109ea501696195e6e3ce5"
dependencies = [
"bytes",
"futures-core",
"futures-sink",
"pin-project-lite",
"tokio",
]
[[package]]
name = "typenum"
version = "1.18.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f"
[[package]] [[package]]
name = "unicode-ident" name = "unicode-ident"
version = "1.0.12" version = "1.0.12"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
[[package]]
name = "universal-hash"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea"
dependencies = [
"crypto-common",
"subtle",
]
[[package]] [[package]]
name = "url" name = "url"
version = "2.5.4" version = "2.5.4"
@@ -830,6 +1070,12 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "version_check"
version = "0.9.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
[[package]] [[package]]
name = "wasi" name = "wasi"
version = "0.11.0+wasi-snapshot-preview1" version = "0.11.0+wasi-snapshot-preview1"
@@ -1027,6 +1273,26 @@ dependencies = [
"synstructure", "synstructure",
] ]
[[package]]
name = "zerocopy"
version = "0.8.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f"
dependencies = [
"zerocopy-derive",
]
[[package]]
name = "zerocopy-derive"
version = "0.8.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]] [[package]]
name = "zerofrom" name = "zerofrom"
version = "0.1.6" version = "0.1.6"
@@ -1048,6 +1314,12 @@ dependencies = [
"synstructure", "synstructure",
] ]
[[package]]
name = "zeroize"
version = "1.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde"
[[package]] [[package]]
name = "zerotrie" name = "zerotrie"
version = "0.2.2" version = "0.2.2"

5
Cargo.toml
View File

@@ -15,6 +15,9 @@ futures = "0.3"
redb = "2.1.3" redb = "2.1.3"
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
bincode = "1.3.3" bincode = "1.3.3"
chacha20poly1305 = "0.10.1"
rand = "0.8"
sha2 = "0.10"
[dev-dependencies] [dev-dependencies]
redis = "0.24" redis = { version = "0.24", features = ["aio", "tokio-comp"] }

99
src/cmd.rs
View File

@@ -4,7 +4,7 @@ use crate::{error::DBError, protocol::Protocol, server::Server};
pub enum Cmd { pub enum Cmd {
Ping, Ping,
Echo(String), Echo(String),
Select(u16), Select(u64), // Changed from u16 to u64
Get(String), Get(String),
Set(String, String), Set(String, String),
SetPx(String, String, u128), SetPx(String, String, u128),
@@ -47,6 +47,7 @@ pub enum Cmd {
LTrim(String, i64, i64), LTrim(String, i64, i64),
LIndex(String, i64), LIndex(String, i64),
LRange(String, i64, i64), LRange(String, i64, i64),
FlushDb,
Unknow(String), Unknow(String),
} }
@@ -65,7 +66,7 @@ impl Cmd {
if cmd.len() != 2 { if cmd.len() != 2 {
return Err(DBError("wrong number of arguments for SELECT".to_string())); return Err(DBError("wrong number of arguments for SELECT".to_string()));
} }
let idx = cmd[1].parse::<u16>().map_err(|_| DBError("ERR DB index is not an integer".to_string()))?; let idx = cmd[1].parse::<u64>().map_err(|_| DBError("ERR DB index is not an integer".to_string()))?;
Cmd::Select(idx) Cmd::Select(idx)
} }
"echo" => Cmd::Echo(cmd[1].clone()), "echo" => Cmd::Echo(cmd[1].clone()),
@@ -394,6 +395,12 @@ impl Cmd {
let stop = cmd[3].parse::<i64>().map_err(|_| DBError("ERR value is not an integer or out of range".to_string()))?; let stop = cmd[3].parse::<i64>().map_err(|_| DBError("ERR value is not an integer or out of range".to_string()))?;
Cmd::LRange(cmd[1].clone(), start, stop) Cmd::LRange(cmd[1].clone(), start, stop)
} }
"flushdb" => {
if cmd.len() != 1 {
return Err(DBError("wrong number of arguments for FLUSHDB command".to_string()));
}
Cmd::FlushDb
}
_ => Cmd::Unknow(cmd[0].clone()), _ => Cmd::Unknow(cmd[0].clone()),
}, },
protocol, protocol,
@@ -482,6 +489,7 @@ impl Cmd {
Cmd::LTrim(key, start, stop) => ltrim_cmd(server, key, *start, *stop).await, Cmd::LTrim(key, start, stop) => ltrim_cmd(server, key, *start, *stop).await,
Cmd::LIndex(key, index) => lindex_cmd(server, key, *index).await, Cmd::LIndex(key, index) => lindex_cmd(server, key, *index).await,
Cmd::LRange(key, start, stop) => lrange_cmd(server, key, *start, *stop).await, Cmd::LRange(key, start, stop) => lrange_cmd(server, key, *start, *stop).await,
Cmd::FlushDb => flushdb_cmd(server).await,
Cmd::Unknow(s) => { Cmd::Unknow(s) => {
println!("\x1b[31;1munknown command: {}\x1b[0m", s); println!("\x1b[31;1munknown command: {}\x1b[0m", s);
Ok(Protocol::err(&format!("ERR unknown command '{}'", s))) Ok(Protocol::err(&format!("ERR unknown command '{}'", s)))
@@ -489,17 +497,25 @@ impl Cmd {
} }
} }
} }
async fn select_cmd(server: &mut Server, db: u16) -> Result<Protocol, DBError> {
let idx = db as usize; async fn flushdb_cmd(server: &mut Server) -> Result<Protocol, DBError> {
if idx >= server.storages.len() { match server.current_storage()?.flushdb() {
return Ok(Protocol::err("ERR DB index is out of range")); Ok(_) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
}
}
async fn select_cmd(server: &mut Server, db: u64) -> Result<Protocol, DBError> {
// Test if we can access the database (this will create it if needed)
server.selected_db = db;
match server.current_storage() {
Ok(_) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)),
} }
server.selected_db = idx;
Ok(Protocol::SimpleString("OK".to_string()))
} }
async fn lindex_cmd(server: &Server, key: &str, index: i64) -> Result<Protocol, DBError> { async fn lindex_cmd(server: &Server, key: &str, index: i64) -> Result<Protocol, DBError> {
match server.current_storage().lindex(key, index) { match server.current_storage()?.lindex(key, index) {
Ok(Some(element)) => Ok(Protocol::BulkString(element)), Ok(Some(element)) => Ok(Protocol::BulkString(element)),
Ok(None) => Ok(Protocol::Null), Ok(None) => Ok(Protocol::Null),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
@@ -507,35 +523,35 @@ async fn lindex_cmd(server: &Server, key: &str, index: i64) -> Result<Protocol,
} }
async fn lrange_cmd(server: &Server, key: &str, start: i64, stop: i64) -> Result<Protocol, DBError> { async fn lrange_cmd(server: &Server, key: &str, start: i64, stop: i64) -> Result<Protocol, DBError> {
match server.current_storage().lrange(key, start, stop) { match server.current_storage()?.lrange(key, start, stop) {
Ok(elements) => Ok(Protocol::Array(elements.into_iter().map(Protocol::BulkString).collect())), Ok(elements) => Ok(Protocol::Array(elements.into_iter().map(Protocol::BulkString).collect())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn ltrim_cmd(server: &Server, key: &str, start: i64, stop: i64) -> Result<Protocol, DBError> { async fn ltrim_cmd(server: &Server, key: &str, start: i64, stop: i64) -> Result<Protocol, DBError> {
match server.current_storage().ltrim(key, start, stop) { match server.current_storage()?.ltrim(key, start, stop) {
Ok(_) => Ok(Protocol::SimpleString("OK".to_string())), Ok(_) => Ok(Protocol::SimpleString("OK".to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn lrem_cmd(server: &Server, key: &str, count: i64, element: &str) -> Result<Protocol, DBError> { async fn lrem_cmd(server: &Server, key: &str, count: i64, element: &str) -> Result<Protocol, DBError> {
match server.current_storage().lrem(key, count, element) { match server.current_storage()?.lrem(key, count, element) {
Ok(removed_count) => Ok(Protocol::SimpleString(removed_count.to_string())), Ok(removed_count) => Ok(Protocol::SimpleString(removed_count.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn llen_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> { async fn llen_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
match server.current_storage().llen(key) { match server.current_storage()?.llen(key) {
Ok(len) => Ok(Protocol::SimpleString(len.to_string())), Ok(len) => Ok(Protocol::SimpleString(len.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn lpop_cmd(server: &Server, key: &str, count: &Option<u64>) -> Result<Protocol, DBError> { async fn lpop_cmd(server: &Server, key: &str, count: &Option<u64>) -> Result<Protocol, DBError> {
match server.current_storage().lpop(key, *count) { match server.current_storage()?.lpop(key, *count) {
Ok(Some(elements)) => { Ok(Some(elements)) => {
if count.is_some() { if count.is_some() {
Ok(Protocol::Array(elements.into_iter().map(Protocol::BulkString).collect())) Ok(Protocol::Array(elements.into_iter().map(Protocol::BulkString).collect()))
@@ -555,7 +571,7 @@ async fn lpop_cmd(server: &Server, key: &str, count: &Option<u64>) -> Result<Pro
} }
async fn rpop_cmd(server: &Server, key: &str, count: &Option<u64>) -> Result<Protocol, DBError> { async fn rpop_cmd(server: &Server, key: &str, count: &Option<u64>) -> Result<Protocol, DBError> {
match server.current_storage().rpop(key, *count) { match server.current_storage()?.rpop(key, *count) {
Ok(Some(elements)) => { Ok(Some(elements)) => {
if count.is_some() { if count.is_some() {
Ok(Protocol::Array(elements.into_iter().map(Protocol::BulkString).collect())) Ok(Protocol::Array(elements.into_iter().map(Protocol::BulkString).collect()))
@@ -575,14 +591,14 @@ async fn rpop_cmd(server: &Server, key: &str, count: &Option<u64>) -> Result<Pro
} }
async fn lpush_cmd(server: &Server, key: &str, elements: &[String]) -> Result<Protocol, DBError> { async fn lpush_cmd(server: &Server, key: &str, elements: &[String]) -> Result<Protocol, DBError> {
match server.current_storage().lpush(key, elements.to_vec()) { match server.current_storage()?.lpush(key, elements.to_vec()) {
Ok(len) => Ok(Protocol::SimpleString(len.to_string())), Ok(len) => Ok(Protocol::SimpleString(len.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn rpush_cmd(server: &Server, key: &str, elements: &[String]) -> Result<Protocol, DBError> { async fn rpush_cmd(server: &Server, key: &str, elements: &[String]) -> Result<Protocol, DBError> {
match server.current_storage().rpush(key, elements.to_vec()) { match server.current_storage()?.rpush(key, elements.to_vec()) {
Ok(len) => Ok(Protocol::SimpleString(len.to_string())), Ok(len) => Ok(Protocol::SimpleString(len.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
@@ -606,7 +622,8 @@ async fn exec_cmd(
} }
async fn incr_cmd(server: &Server, key: &String) -> Result<Protocol, DBError> { async fn incr_cmd(server: &Server, key: &String) -> Result<Protocol, DBError> {
let current_value = server.current_storage().get(key)?; let storage = server.current_storage()?;
let current_value = storage.get(key)?;
let new_value = match current_value { let new_value = match current_value {
Some(v) => { Some(v) => {
@@ -618,7 +635,7 @@ async fn incr_cmd(server: &Server, key: &String) -> Result<Protocol, DBError> {
None => 1, None => 1,
}; };
server.current_storage().set(key.clone(), new_value.to_string())?; storage.set(key.clone(), new_value.to_string())?;
Ok(Protocol::SimpleString(new_value.to_string())) Ok(Protocol::SimpleString(new_value.to_string()))
} }
@@ -634,14 +651,14 @@ fn config_get_cmd(name: &String, server: &Server) -> Result<Protocol, DBError> {
])), ])),
"databases" => Ok(Protocol::Array(vec![ "databases" => Ok(Protocol::Array(vec![
Protocol::BulkString(name.clone()), Protocol::BulkString(name.clone()),
Protocol::BulkString(server.option.databases.to_string()), Protocol::BulkString(server.option.max_databases.unwrap_or(0).to_string()),
])), ])),
_ => Ok(Protocol::Array(vec![])), _ => Ok(Protocol::Array(vec![])),
} }
} }
async fn keys_cmd(server: &Server) -> Result<Protocol, DBError> { async fn keys_cmd(server: &Server) -> Result<Protocol, DBError> {
let keys = server.current_storage().keys("*")?; let keys = server.current_storage()?.keys("*")?;
Ok(Protocol::Array( Ok(Protocol::Array(
keys.into_iter().map(Protocol::BulkString).collect(), keys.into_iter().map(Protocol::BulkString).collect(),
)) ))
@@ -660,14 +677,14 @@ fn info_cmd(section: &Option<String>) -> Result<Protocol, DBError> {
} }
async fn type_cmd(server: &Server, k: &String) -> Result<Protocol, DBError> { async fn type_cmd(server: &Server, k: &String) -> Result<Protocol, DBError> {
match server.current_storage().get_key_type(k)? { match server.current_storage()?.get_key_type(k)? {
Some(type_str) => Ok(Protocol::SimpleString(type_str)), Some(type_str) => Ok(Protocol::SimpleString(type_str)),
None => Ok(Protocol::SimpleString("none".to_string())), None => Ok(Protocol::SimpleString("none".to_string())),
} }
} }
async fn del_cmd(server: &Server, k: &str) -> Result<Protocol, DBError> { async fn del_cmd(server: &Server, k: &str) -> Result<Protocol, DBError> {
server.current_storage().del(k.to_string())?; server.current_storage()?.del(k.to_string())?;
Ok(Protocol::SimpleString("1".to_string())) Ok(Protocol::SimpleString("1".to_string()))
} }
@@ -677,7 +694,7 @@ async fn set_ex_cmd(
v: &str, v: &str,
x: &u128, x: &u128,
) -> Result<Protocol, DBError> { ) -> Result<Protocol, DBError> {
server.current_storage().setx(k.to_string(), v.to_string(), *x * 1000)?; server.current_storage()?.setx(k.to_string(), v.to_string(), *x * 1000)?;
Ok(Protocol::SimpleString("OK".to_string())) Ok(Protocol::SimpleString("OK".to_string()))
} }
@@ -687,28 +704,28 @@ async fn set_px_cmd(
v: &str, v: &str,
x: &u128, x: &u128,
) -> Result<Protocol, DBError> { ) -> Result<Protocol, DBError> {
server.current_storage().setx(k.to_string(), v.to_string(), *x)?; server.current_storage()?.setx(k.to_string(), v.to_string(), *x)?;
Ok(Protocol::SimpleString("OK".to_string())) Ok(Protocol::SimpleString("OK".to_string()))
} }
async fn set_cmd(server: &Server, k: &str, v: &str) -> Result<Protocol, DBError> { async fn set_cmd(server: &Server, k: &str, v: &str) -> Result<Protocol, DBError> {
server.current_storage().set(k.to_string(), v.to_string())?; server.current_storage()?.set(k.to_string(), v.to_string())?;
Ok(Protocol::SimpleString("OK".to_string())) Ok(Protocol::SimpleString("OK".to_string()))
} }
async fn get_cmd(server: &Server, k: &str) -> Result<Protocol, DBError> { async fn get_cmd(server: &Server, k: &str) -> Result<Protocol, DBError> {
let v = server.current_storage().get(k)?; let v = server.current_storage()?.get(k)?;
Ok(v.map_or(Protocol::Null, Protocol::BulkString)) Ok(v.map_or(Protocol::Null, Protocol::BulkString))
} }
// Hash command implementations // Hash command implementations
async fn hset_cmd(server: &Server, key: &str, pairs: &[(String, String)]) -> Result<Protocol, DBError> { async fn hset_cmd(server: &Server, key: &str, pairs: &[(String, String)]) -> Result<Protocol, DBError> {
let new_fields = server.current_storage().hset(key, pairs)?; let new_fields = server.current_storage()?.hset(key, pairs)?;
Ok(Protocol::SimpleString(new_fields.to_string())) Ok(Protocol::SimpleString(new_fields.to_string()))
} }
async fn hget_cmd(server: &Server, key: &str, field: &str) -> Result<Protocol, DBError> { async fn hget_cmd(server: &Server, key: &str, field: &str) -> Result<Protocol, DBError> {
match server.current_storage().hget(key, field) { match server.current_storage()?.hget(key, field) {
Ok(Some(value)) => Ok(Protocol::BulkString(value)), Ok(Some(value)) => Ok(Protocol::BulkString(value)),
Ok(None) => Ok(Protocol::Null), Ok(None) => Ok(Protocol::Null),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
@@ -716,7 +733,7 @@ async fn hget_cmd(server: &Server, key: &str, field: &str) -> Result<Protocol, D
} }
async fn hgetall_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> { async fn hgetall_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
match server.current_storage().hgetall(key) { match server.current_storage()?.hgetall(key) {
Ok(pairs) => { Ok(pairs) => {
let mut result = Vec::new(); let mut result = Vec::new();
for (field, value) in pairs { for (field, value) in pairs {
@@ -730,21 +747,21 @@ async fn hgetall_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
} }
async fn hdel_cmd(server: &Server, key: &str, fields: &[String]) -> Result<Protocol, DBError> { async fn hdel_cmd(server: &Server, key: &str, fields: &[String]) -> Result<Protocol, DBError> {
match server.current_storage().hdel(key, fields) { match server.current_storage()?.hdel(key, fields) {
Ok(deleted) => Ok(Protocol::SimpleString(deleted.to_string())), Ok(deleted) => Ok(Protocol::SimpleString(deleted.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn hexists_cmd(server: &Server, key: &str, field: &str) -> Result<Protocol, DBError> { async fn hexists_cmd(server: &Server, key: &str, field: &str) -> Result<Protocol, DBError> {
match server.current_storage().hexists(key, field) { match server.current_storage()?.hexists(key, field) {
Ok(exists) => Ok(Protocol::SimpleString(if exists { "1" } else { "0" }.to_string())), Ok(exists) => Ok(Protocol::SimpleString(if exists { "1" } else { "0" }.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn hkeys_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> { async fn hkeys_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
match server.current_storage().hkeys(key) { match server.current_storage()?.hkeys(key) {
Ok(keys) => Ok(Protocol::Array( Ok(keys) => Ok(Protocol::Array(
keys.into_iter().map(Protocol::BulkString).collect(), keys.into_iter().map(Protocol::BulkString).collect(),
)), )),
@@ -753,7 +770,7 @@ async fn hkeys_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
} }
async fn hvals_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> { async fn hvals_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
match server.current_storage().hvals(key) { match server.current_storage()?.hvals(key) {
Ok(values) => Ok(Protocol::Array( Ok(values) => Ok(Protocol::Array(
values.into_iter().map(Protocol::BulkString).collect(), values.into_iter().map(Protocol::BulkString).collect(),
)), )),
@@ -762,14 +779,14 @@ async fn hvals_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
} }
async fn hlen_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> { async fn hlen_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
match server.current_storage().hlen(key) { match server.current_storage()?.hlen(key) {
Ok(len) => Ok(Protocol::SimpleString(len.to_string())), Ok(len) => Ok(Protocol::SimpleString(len.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn hmget_cmd(server: &Server, key: &str, fields: &[String]) -> Result<Protocol, DBError> { async fn hmget_cmd(server: &Server, key: &str, fields: &[String]) -> Result<Protocol, DBError> {
match server.current_storage().hmget(key, fields) { match server.current_storage()?.hmget(key, fields) {
Ok(values) => { Ok(values) => {
let result: Vec<Protocol> = values let result: Vec<Protocol> = values
.into_iter() .into_iter()
@@ -782,14 +799,14 @@ async fn hmget_cmd(server: &Server, key: &str, fields: &[String]) -> Result<Prot
} }
async fn hsetnx_cmd(server: &Server, key: &str, field: &str, value: &str) -> Result<Protocol, DBError> { async fn hsetnx_cmd(server: &Server, key: &str, field: &str, value: &str) -> Result<Protocol, DBError> {
match server.current_storage().hsetnx(key, field, value) { match server.current_storage()?.hsetnx(key, field, value) {
Ok(was_set) => Ok(Protocol::SimpleString(if was_set { "1" } else { "0" }.to_string())), Ok(was_set) => Ok(Protocol::SimpleString(if was_set { "1" } else { "0" }.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn scan_cmd(server: &Server, cursor: &u64, pattern: Option<&str>, count: &Option<u64>) -> Result<Protocol, DBError> { async fn scan_cmd(server: &Server, cursor: &u64, pattern: Option<&str>, count: &Option<u64>) -> Result<Protocol, DBError> {
match server.current_storage().scan(*cursor, pattern, *count) { match server.current_storage()?.scan(*cursor, pattern, *count) {
Ok((next_cursor, keys)) => { Ok((next_cursor, keys)) => {
let mut result = Vec::new(); let mut result = Vec::new();
result.push(Protocol::BulkString(next_cursor.to_string())); result.push(Protocol::BulkString(next_cursor.to_string()));
@@ -803,7 +820,7 @@ async fn scan_cmd(server: &Server, cursor: &u64, pattern: Option<&str>, count: &
} }
async fn hscan_cmd(server: &Server, key: &str, cursor: &u64, pattern: Option<&str>, count: &Option<u64>) -> Result<Protocol, DBError> { async fn hscan_cmd(server: &Server, key: &str, cursor: &u64, pattern: Option<&str>, count: &Option<u64>) -> Result<Protocol, DBError> {
match server.current_storage().hscan(key, *cursor, pattern, *count) { match server.current_storage()?.hscan(key, *cursor, pattern, *count) {
Ok((next_cursor, fields)) => { Ok((next_cursor, fields)) => {
let mut result = Vec::new(); let mut result = Vec::new();
result.push(Protocol::BulkString(next_cursor.to_string())); result.push(Protocol::BulkString(next_cursor.to_string()));
@@ -817,14 +834,14 @@ async fn hscan_cmd(server: &Server, key: &str, cursor: &u64, pattern: Option<&st
} }
async fn ttl_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> { async fn ttl_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
match server.current_storage().ttl(key) { match server.current_storage()?.ttl(key) {
Ok(ttl) => Ok(Protocol::SimpleString(ttl.to_string())), Ok(ttl) => Ok(Protocol::SimpleString(ttl.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }
} }
async fn exists_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> { async fn exists_cmd(server: &Server, key: &str) -> Result<Protocol, DBError> {
match server.current_storage().exists(key) { match server.current_storage()?.exists(key) {
Ok(exists) => Ok(Protocol::SimpleString(if exists { "1" } else { "0" }.to_string())), Ok(exists) => Ok(Protocol::SimpleString(if exists { "1" } else { "0" }.to_string())),
Err(e) => Ok(Protocol::err(&e.0)), Err(e) => Ok(Protocol::err(&e.0)),
} }

73
src/crypto.rs Normal file
View File

@@ -0,0 +1,73 @@
use chacha20poly1305::{
aead::{Aead, KeyInit, OsRng},
XChaCha20Poly1305, XNonce,
};
use rand::RngCore;
use sha2::{Digest, Sha256};
const VERSION: u8 = 1;
const NONCE_LEN: usize = 24;
const TAG_LEN: usize = 16;
#[derive(Debug)]
pub enum CryptoError {
Format, // wrong length / header
Version(u8), // unknown version
Decrypt, // wrong key or corrupted data
}
impl From<CryptoError> for crate::error::DBError {
fn from(e: CryptoError) -> Self {
crate::error::DBError(format!("Crypto error: {:?}", e))
}
}
/// Super-simple factory: new(secret) + encrypt(bytes) + decrypt(bytes)
pub struct CryptoFactory {
key: chacha20poly1305::Key,
}
impl CryptoFactory {
/// Accepts any secret bytes; turns them into a 32-byte key (SHA-256).
pub fn new<S: AsRef<[u8]>>(secret: S) -> Self {
let mut h = Sha256::new();
h.update(b"xchacha20poly1305-factory:v1"); // domain separation
h.update(secret.as_ref());
let digest = h.finalize(); // 32 bytes
let key = chacha20poly1305::Key::from_slice(&digest).to_owned();
Self { key }
}
/// Output layout: [version:1][nonce:24][ciphertext||tag]
pub fn encrypt(&self, plaintext: &[u8]) -> Vec<u8> {
let cipher = XChaCha20Poly1305::new(&self.key);
let mut nonce_bytes = [0u8; NONCE_LEN];
OsRng.fill_bytes(&mut nonce_bytes);
let nonce = XNonce::from_slice(&nonce_bytes);
let mut out = Vec::with_capacity(1 + NONCE_LEN + plaintext.len() + TAG_LEN);
out.push(VERSION);
out.extend_from_slice(&nonce_bytes);
let ct = cipher.encrypt(nonce, plaintext).expect("encrypt");
out.extend_from_slice(&ct);
out
}
pub fn decrypt(&self, blob: &[u8]) -> Result<Vec<u8>, CryptoError> {
if blob.len() < 1 + NONCE_LEN + TAG_LEN {
return Err(CryptoError::Format);
}
let ver = blob[0];
if ver != VERSION {
return Err(CryptoError::Version(ver));
}
let nonce = XNonce::from_slice(&blob[1..1 + NONCE_LEN]);
let ct = &blob[1 + NONCE_LEN..];
let cipher = XChaCha20Poly1305::new(&self.key);
cipher.decrypt(nonce, ct).map_err(|_| CryptoError::Decrypt)
}
}

1
src/lib.rs
View File

@@ -1,4 +1,5 @@
pub mod cmd; pub mod cmd;
pub mod crypto;
pub mod error; pub mod error;
pub mod options; pub mod options;
pub mod protocol; pub mod protocol;

14
src/main.rs
View File

@@ -14,7 +14,6 @@ struct Args {
#[arg(long)] #[arg(long)]
dir: String, dir: String,
/// The port of the Redis server, default is 6379 if not specified /// The port of the Redis server, default is 6379 if not specified
#[arg(long)] #[arg(long)]
port: Option<u16>, port: Option<u16>,
@@ -23,9 +22,13 @@ struct Args {
#[arg(long)] #[arg(long)]
debug: bool, debug: bool,
/// Number of logical databases (SELECT 0..N-1) /// Maximum number of logical databases (None = unlimited)
#[arg(long, default_value_t = 16)] #[arg(long)]
databases: u16, max_databases: Option<u64>,
/// Master encryption key for encrypted databases
#[arg(long)]
encryption_key: Option<String>,
} }
#[tokio::main] #[tokio::main]
@@ -45,7 +48,8 @@ async fn main() {
dir: args.dir, dir: args.dir,
port, port,
debug: args.debug, debug: args.debug,
databases: args.databases, max_databases: args.max_databases,
encryption_key: args.encryption_key,
}; };
// new server // new server

3
src/options.rs
View File

@@ -3,5 +3,6 @@ pub struct DBOption {
pub dir: String, pub dir: String,
pub port: u16, pub port: u16,
pub debug: bool, pub debug: bool,
pub databases: u16, // number of logical DBs (default 16) pub max_databases: Option<u64>, // None = unlimited, Some(n) = limit to n
pub encryption_key: Option<String>, // Master encryption key
} }

55
src/server.rs
View File

@@ -1,5 +1,5 @@
use core::str; use core::str;
use std::path::PathBuf; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use tokio::io::AsyncReadExt; use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt; use tokio::io::AsyncWriteExt;
@@ -12,34 +12,56 @@ use crate::storage::Storage;
#[derive(Clone)] #[derive(Clone)]
pub struct Server { pub struct Server {
pub storages: Vec<Arc<Storage>>, pub db_cache: std::sync::Arc<std::sync::RwLock<HashMap<u64, Arc<Storage>>>>,
pub option: options::DBOption, pub option: options::DBOption,
pub client_name: Option<String>, pub client_name: Option<String>,
pub selected_db: usize, // per-connection pub selected_db: u64, // Changed from usize to u64
} }
impl Server { impl Server {
pub async fn new(option: options::DBOption) -> Self { pub async fn new(option: options::DBOption) -> Self {
// Eagerly create N db files: <dir>/<index>.db
let mut storages = Vec::with_capacity(option.databases as usize);
for i in 0..option.databases {
let db_file_path = PathBuf::from(option.dir.clone()).join(format!("{}.db", i));
println!("will open db file path (db {}): {}", i, db_file_path.display());
let storage = Storage::new(db_file_path).expect("Failed to initialize storage");
storages.push(Arc::new(storage));
}
Server { Server {
storages, db_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
option, option,
client_name: None, client_name: None,
selected_db: 0, selected_db: 0,
} }
} }
#[inline] pub fn current_storage(&self) -> Result<Arc<Storage>, DBError> {
pub fn current_storage(&self) -> &Storage { let mut cache = self.db_cache.write().unwrap();
self.storages[self.selected_db].as_ref()
if let Some(storage) = cache.get(&self.selected_db) {
return Ok(storage.clone());
}
// Check database limit if set
if let Some(max_db) = self.option.max_databases {
if self.selected_db >= max_db {
return Err(DBError(format!("DB index {} is out of range (max: {})", self.selected_db, max_db - 1)));
}
}
// Create new database file
let db_file_path = std::path::PathBuf::from(self.option.dir.clone())
.join(format!("{}.db", self.selected_db));
println!("Creating new db file: {}", db_file_path.display());
let storage = Arc::new(Storage::new(
db_file_path,
self.should_encrypt_db(self.selected_db),
self.option.encryption_key.as_deref()
)?);
cache.insert(self.selected_db, storage.clone());
Ok(storage)
}
fn should_encrypt_db(&self, db_index: u64) -> bool {
// You can implement logic here to determine which databases should be encrypted
// For now, let's say databases with even numbers are encrypted if key is provided
self.option.encryption_key.is_some() && db_index % 2 == 0
} }
pub async fn handle( pub async fn handle(
@@ -104,6 +126,5 @@ impl Server {
} }
} }
} }
Ok(())
} }
} }

281
src/storage.rs
View File

@@ -6,8 +6,101 @@ use std::{
use redb::{Database, ReadableTable, TableDefinition}; use redb::{Database, ReadableTable, TableDefinition};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::crypto::CryptoFactory;
use crate::error::DBError; use crate::error::DBError;
// Add this glob matching function
fn glob_match(pattern: &str, text: &str) -> bool {
fn match_recursive(pattern: &[char], text: &[char], p_idx: usize, t_idx: usize) -> bool {
if p_idx >= pattern.len() {
return t_idx >= text.len();
}
match pattern[p_idx] {
'*' => {
// Try matching zero characters
if match_recursive(pattern, text, p_idx + 1, t_idx) {
return true;
}
// Try matching one or more characters
for i in t_idx..text.len() {
if match_recursive(pattern, text, p_idx + 1, i + 1) {
return true;
}
}
false
}
'?' => {
if t_idx >= text.len() {
false
} else {
match_recursive(pattern, text, p_idx + 1, t_idx + 1)
}
}
'[' => {
// Find the closing bracket
let mut bracket_end = p_idx + 1;
while bracket_end < pattern.len() && pattern[bracket_end] != ']' {
bracket_end += 1;
}
if bracket_end >= pattern.len() || t_idx >= text.len() {
return false;
}
let bracket_content = &pattern[p_idx + 1..bracket_end];
let char_to_match = text[t_idx];
let mut matched = false;
let mut i = 0;
while i < bracket_content.len() {
if i + 2 < bracket_content.len() && bracket_content[i + 1] == '-' {
// Range like [a-z]
if char_to_match >= bracket_content[i] && char_to_match <= bracket_content[i + 2] {
matched = true;
break;
}
i += 3;
} else {
// Single character
if char_to_match == bracket_content[i] {
matched = true;
break;
}
i += 1;
}
}
if matched {
match_recursive(pattern, text, bracket_end + 1, t_idx + 1)
} else {
false
}
}
'\\' => {
// Escape next character
if p_idx + 1 >= pattern.len() || t_idx >= text.len() {
false
} else if pattern[p_idx + 1] == text[t_idx] {
match_recursive(pattern, text, p_idx + 2, t_idx + 1)
} else {
false
}
}
c => {
if t_idx >= text.len() || c != text[t_idx] {
false
} else {
match_recursive(pattern, text, p_idx + 1, t_idx + 1)
}
}
}
}
let pattern_chars: Vec<char> = pattern.chars().collect();
let text_chars: Vec<char> = text.chars().collect();
match_recursive(&pattern_chars, &text_chars, 0, 0)
}
// Table definitions for different Redis data types // Table definitions for different Redis data types
const TYPES_TABLE: TableDefinition<&str, &str> = TableDefinition::new("types"); const TYPES_TABLE: TableDefinition<&str, &str> = TableDefinition::new("types");
const STRINGS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("strings"); const STRINGS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("strings");
@@ -15,6 +108,7 @@ const HASHES_TABLE: TableDefinition<(&str, &str), &str> = TableDefinition::new("
const LISTS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("lists"); const LISTS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("lists");
const STREAMS_META_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("streams_meta"); const STREAMS_META_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("streams_meta");
const STREAMS_DATA_TABLE: TableDefinition<(&str, &str), &[u8]> = TableDefinition::new("streams_data"); const STREAMS_DATA_TABLE: TableDefinition<(&str, &str), &[u8]> = TableDefinition::new("streams_data");
const ENCRYPTED_TABLE: TableDefinition<&str, u8> = TableDefinition::new("encrypted");
#[derive(Serialize, Deserialize, Debug, Clone)] #[derive(Serialize, Deserialize, Debug, Clone)]
pub struct StringValue { pub struct StringValue {
@@ -42,10 +136,11 @@ pub fn now_in_millis() -> u128 {
pub struct Storage { pub struct Storage {
db: Database, db: Database,
crypto: Option<CryptoFactory>,
} }
impl Storage { impl Storage {
pub fn new(path: impl AsRef<Path>) -> Result<Self, DBError> { pub fn new(path: impl AsRef<Path>, should_encrypt: bool, master_key: Option<&str>) -> Result<Self, DBError> {
let db = Database::create(path)?; let db = Database::create(path)?;
// Create tables if they don't exist // Create tables if they don't exist
@@ -57,10 +152,109 @@ impl Storage {
let _ = write_txn.open_table(LISTS_TABLE)?; let _ = write_txn.open_table(LISTS_TABLE)?;
let _ = write_txn.open_table(STREAMS_META_TABLE)?; let _ = write_txn.open_table(STREAMS_META_TABLE)?;
let _ = write_txn.open_table(STREAMS_DATA_TABLE)?; let _ = write_txn.open_table(STREAMS_DATA_TABLE)?;
let _ = write_txn.open_table(ENCRYPTED_TABLE)?;
} }
write_txn.commit()?; write_txn.commit()?;
Ok(Storage { db }) // 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())
}
}
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)?;
// 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 (key, field) = binding.0.value();
(key.to_string(), field.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()))?;
}
}
write_txn.commit()?;
Ok(())
} }
pub fn get_key_type(&self, key: &str) -> Result<Option<String>, DBError> { pub fn get_key_type(&self, key: &str) -> Result<Option<String>, DBError> {
@@ -73,22 +267,22 @@ impl Storage {
} }
} }
// Update the get method to use decryption
pub fn get(&self, key: &str) -> Result<Option<String>, DBError> { pub fn get(&self, key: &str) -> Result<Option<String>, DBError> {
let read_txn = self.db.begin_read()?; let read_txn = self.db.begin_read()?;
// Check if key exists and is of string type
let types_table = read_txn.open_table(TYPES_TABLE)?; let types_table = read_txn.open_table(TYPES_TABLE)?;
match types_table.get(key)? { match types_table.get(key)? {
Some(type_val) if type_val.value() == "string" => { Some(type_val) if type_val.value() == "string" => {
let strings_table = read_txn.open_table(STRINGS_TABLE)?; let strings_table = read_txn.open_table(STRINGS_TABLE)?;
match strings_table.get(key)? { match strings_table.get(key)? {
Some(data) => { Some(data) => {
let string_value: StringValue = bincode::deserialize(data.value())?; let decrypted = self.decrypt_if_needed(data.value())?;
let string_value: StringValue = bincode::deserialize(&decrypted)?;
// Check if expired // Check if expired
if let Some(expires_at) = string_value.expires_at_ms { if let Some(expires_at) = string_value.expires_at_ms {
if now_in_millis() > expires_at { if now_in_millis() > expires_at {
// Key expired, remove it
drop(read_txn); drop(read_txn);
self.del(key.to_string())?; self.del(key.to_string())?;
return Ok(None); return Ok(None);
@@ -104,6 +298,10 @@ impl Storage {
} }
} }
// Apply similar encryption/decryption to other methods (setx, hset, lpush, etc.)
// ... (you'll need to update all methods that store/retrieve serialized data)
// Update the set method to use encryption
pub fn set(&self, key: String, value: String) -> Result<(), DBError> { pub fn set(&self, key: String, value: String) -> Result<(), DBError> {
let write_txn = self.db.begin_write()?; let write_txn = self.db.begin_write()?;
@@ -117,7 +315,8 @@ impl Storage {
expires_at_ms: None, expires_at_ms: None,
}; };
let serialized = bincode::serialize(&string_value)?; let serialized = bincode::serialize(&string_value)?;
strings_table.insert(key.as_str(), serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
strings_table.insert(key.as_str(), encrypted.as_slice())?;
} }
write_txn.commit()?; write_txn.commit()?;
@@ -137,7 +336,8 @@ impl Storage {
expires_at_ms: Some(expire_ms + now_in_millis()), expires_at_ms: Some(expire_ms + now_in_millis()),
}; };
let serialized = bincode::serialize(&string_value)?; let serialized = bincode::serialize(&string_value)?;
strings_table.insert(key.as_str(), serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
strings_table.insert(key.as_str(), encrypted.as_slice())?;
} }
write_txn.commit()?; write_txn.commit()?;
@@ -191,7 +391,7 @@ impl Storage {
let mut iter = table.iter()?; let mut iter = table.iter()?;
while let Some(entry) = iter.next() { while let Some(entry) = iter.next() {
let key = entry?.0.value().to_string(); let key = entry?.0.value().to_string();
if pattern == "*" || key.contains(pattern) { if pattern == "*" || glob_match(pattern, &key) {
keys.push(key); keys.push(key);
} }
} }
@@ -251,7 +451,7 @@ impl Storage {
None => Ok(None), None => Ok(None),
} }
} }
Some(_) => Err(DBError("WRONGTYPE Operation against a key holding the wrong kind of value".to_string())), Some(_) => return Err(DBError("WRONGTYPE Operation against a key holding the wrong kind of value".to_string())),
None => Ok(None), None => Ok(None),
} }
} }
@@ -649,7 +849,7 @@ impl Storage {
// List operations // List operations
pub fn lpush(&self, key: &str, elements: Vec<String>) -> Result<u64, DBError> { pub fn lpush(&self, key: &str, elements: Vec<String>) -> Result<u64, DBError> {
let write_txn = self.db.begin_write()?; let write_txn = self.db.begin_write()?;
let mut new_len = 0u64; let new_len;
{ {
let mut types_table = write_txn.open_table(TYPES_TABLE)?; let mut types_table = write_txn.open_table(TYPES_TABLE)?;
@@ -671,7 +871,10 @@ impl Storage {
} }
let mut list_value: ListValue = match lists_table.get(key)? { let mut list_value: ListValue = match lists_table.get(key)? {
Some(data) => bincode::deserialize(data.value())?, Some(data) => {
let decrypted = self.decrypt_if_needed(data.value())?;
bincode::deserialize(&decrypted)?
},
None => ListValue { elements: Vec::new() }, None => ListValue { elements: Vec::new() },
}; };
@@ -681,7 +884,8 @@ impl Storage {
new_len = list_value.elements.len() as u64; new_len = list_value.elements.len() as u64;
let serialized = bincode::serialize(&list_value)?; let serialized = bincode::serialize(&list_value)?;
lists_table.insert(key, serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
lists_table.insert(key, encrypted.as_slice())?;
} }
write_txn.commit()?; write_txn.commit()?;
@@ -690,7 +894,7 @@ impl Storage {
pub fn rpush(&self, key: &str, elements: Vec<String>) -> Result<u64, DBError> { pub fn rpush(&self, key: &str, elements: Vec<String>) -> Result<u64, DBError> {
let write_txn = self.db.begin_write()?; let write_txn = self.db.begin_write()?;
let mut new_len = 0u64; let new_len;
{ {
let mut types_table = write_txn.open_table(TYPES_TABLE)?; let mut types_table = write_txn.open_table(TYPES_TABLE)?;
@@ -712,7 +916,10 @@ impl Storage {
} }
let mut list_value: ListValue = match lists_table.get(key)? { let mut list_value: ListValue = match lists_table.get(key)? {
Some(data) => bincode::deserialize(data.value())?, Some(data) => {
let decrypted = self.decrypt_if_needed(data.value())?;
bincode::deserialize(&decrypted)?
},
None => ListValue { elements: Vec::new() }, None => ListValue { elements: Vec::new() },
}; };
@@ -722,7 +929,8 @@ impl Storage {
new_len = list_value.elements.len() as u64; new_len = list_value.elements.len() as u64;
let serialized = bincode::serialize(&list_value)?; let serialized = bincode::serialize(&list_value)?;
lists_table.insert(key, serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
lists_table.insert(key, encrypted.as_slice())?;
} }
write_txn.commit()?; write_txn.commit()?;
@@ -748,7 +956,10 @@ impl Storage {
} }
Some(_) => { Some(_) => {
let mut list_value: ListValue = match lists_table.get(key)? { let mut list_value: ListValue = match lists_table.get(key)? {
Some(data) => bincode::deserialize(data.value())?, Some(data) => {
let decrypted = self.decrypt_if_needed(data.value())?;
bincode::deserialize(&decrypted)?
},
None => return Ok(None), // Key exists but list is empty (shouldn't happen if type is "list") None => return Ok(None), // Key exists but list is empty (shouldn't happen if type is "list")
}; };
@@ -766,7 +977,8 @@ impl Storage {
types_table.remove(key)?; types_table.remove(key)?;
} else { } else {
let serialized = bincode::serialize(&list_value)?; let serialized = bincode::serialize(&list_value)?;
lists_table.insert(key, serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
lists_table.insert(key, encrypted.as_slice())?;
} }
} }
None => return Ok(None), None => return Ok(None),
@@ -800,7 +1012,10 @@ impl Storage {
} }
Some(_) => { Some(_) => {
let mut list_value: ListValue = match lists_table.get(key)? { let mut list_value: ListValue = match lists_table.get(key)? {
Some(data) => bincode::deserialize(data.value())?, Some(data) => {
let decrypted = self.decrypt_if_needed(data.value())?;
bincode::deserialize(&decrypted)?
}
None => return Ok(None), None => return Ok(None),
}; };
@@ -818,7 +1033,8 @@ impl Storage {
types_table.remove(key)?; types_table.remove(key)?;
} else { } else {
let serialized = bincode::serialize(&list_value)?; let serialized = bincode::serialize(&list_value)?;
lists_table.insert(key, serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
lists_table.insert(key, encrypted.as_slice())?;
} }
} }
None => return Ok(None), None => return Ok(None),
@@ -842,7 +1058,8 @@ impl Storage {
let lists_table = read_txn.open_table(LISTS_TABLE)?; let lists_table = read_txn.open_table(LISTS_TABLE)?;
match lists_table.get(key)? { match lists_table.get(key)? {
Some(data) => { Some(data) => {
let list_value: ListValue = bincode::deserialize(data.value())?; let decrypted = self.decrypt_if_needed(data.value())?;
let list_value: ListValue = bincode::deserialize(&decrypted)?;
Ok(list_value.elements.len() as u64) Ok(list_value.elements.len() as u64)
} }
None => Ok(0), // Key exists but list is empty None => Ok(0), // Key exists but list is empty
@@ -855,7 +1072,7 @@ impl Storage {
pub fn lrem(&self, key: &str, count: i64, element: &str) -> Result<u64, DBError> { pub fn lrem(&self, key: &str, count: i64, element: &str) -> Result<u64, DBError> {
let write_txn = self.db.begin_write()?; let write_txn = self.db.begin_write()?;
let mut removed_count = 0u64; let removed_count;
{ {
let mut types_table = write_txn.open_table(TYPES_TABLE)?; let mut types_table = write_txn.open_table(TYPES_TABLE)?;
@@ -872,7 +1089,10 @@ impl Storage {
} }
Some(_) => { Some(_) => {
let mut list_value: ListValue = match lists_table.get(key)? { let mut list_value: ListValue = match lists_table.get(key)? {
Some(data) => bincode::deserialize(data.value())?, Some(data) => {
let decrypted = self.decrypt_if_needed(data.value())?;
bincode::deserialize(&decrypted)?
}
None => return Ok(0), None => return Ok(0),
}; };
@@ -910,7 +1130,8 @@ impl Storage {
types_table.remove(key)?; types_table.remove(key)?;
} else { } else {
let serialized = bincode::serialize(&list_value)?; let serialized = bincode::serialize(&list_value)?;
lists_table.insert(key, serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
lists_table.insert(key, encrypted.as_slice())?;
} }
} }
None => return Ok(0), None => return Ok(0),
@@ -939,7 +1160,10 @@ impl Storage {
} }
Some(_) => { Some(_) => {
let mut list_value: ListValue = match lists_table.get(key)? { let mut list_value: ListValue = match lists_table.get(key)? {
Some(data) => bincode::deserialize(data.value())?, Some(data) => {
let decrypted = self.decrypt_if_needed(data.value())?;
bincode::deserialize(&decrypted)?
}
None => return Ok(()), None => return Ok(()),
}; };
@@ -974,7 +1198,8 @@ impl Storage {
types_table.remove(key)?; types_table.remove(key)?;
} else { } else {
let serialized = bincode::serialize(&list_value)?; let serialized = bincode::serialize(&list_value)?;
lists_table.insert(key, serialized.as_slice())?; let encrypted = self.encrypt_if_needed(&serialized)?;
lists_table.insert(key, encrypted.as_slice())?;
} }
} }
None => {} None => {}
@@ -994,7 +1219,8 @@ impl Storage {
let lists_table = read_txn.open_table(LISTS_TABLE)?; let lists_table = read_txn.open_table(LISTS_TABLE)?;
match lists_table.get(key)? { match lists_table.get(key)? {
Some(data) => { Some(data) => {
let list_value: ListValue = bincode::deserialize(data.value())?; let decrypted = self.decrypt_if_needed(data.value())?;
let list_value: ListValue = bincode::deserialize(&decrypted)?;
let len = list_value.elements.len() as i64; let len = list_value.elements.len() as i64;
let mut index = index; let mut index = index;
if index < 0 { if index < 0 {
@@ -1023,7 +1249,8 @@ impl Storage {
let lists_table = read_txn.open_table(LISTS_TABLE)?; let lists_table = read_txn.open_table(LISTS_TABLE)?;
match lists_table.get(key)? { match lists_table.get(key)? {
Some(data) => { Some(data) => {
let list_value: ListValue = bincode::deserialize(data.value())?; let decrypted = self.decrypt_if_needed(data.value())?;
let list_value: ListValue = bincode::deserialize(&decrypted)?;
let len = list_value.elements.len() as i64; let len = list_value.elements.len() as i64;
let mut start = start; let mut start = start;
let mut stop = stop; let mut stop = stop;

3
tests/debug_hset.rs
View File

@@ -25,7 +25,8 @@ async fn debug_hset_simple() {
dir: test_dir.to_string(), dir: test_dir.to_string(),
port, port,
debug: false, debug: false,
databases: 16, max_databases: Some(16),
encryption_key: None,
}; };
let mut server = Server::new(option).await; let mut server = Server::new(option).await;

3
tests/debug_hset_simple.rs
View File

@@ -16,7 +16,8 @@ async fn debug_hset_return_value() {
dir: test_dir.to_string(), dir: test_dir.to_string(),
port: 16390, port: 16390,
debug: false, debug: false,
databases: 16, max_databases: Some(16),
encryption_key: None,
}; };
let mut server = Server::new(option).await; let mut server = Server::new(option).await;

3
tests/redis_tests.rs
View File

@@ -20,7 +20,8 @@ async fn start_test_server(test_name: &str) -> (Server, u16) {
dir: test_dir, dir: test_dir,
port, port,
debug: true, debug: true,
databases: 16, max_databases: Some(16),
encryption_key: None,
}; };
let server = Server::new(option).await; let server = Server::new(option).await;

3
tests/simple_integration_test.rs
View File

@@ -22,7 +22,8 @@ async fn start_test_server(test_name: &str) -> (Server, u16) {
dir: test_dir, dir: test_dir,
port, port,
debug: true, debug: true,
databases: 16, max_databases: Some(16),
encryption_key: None,
}; };
let server = Server::new(option).await; let server = Server::new(option).await;