From bf7e444b57b9f30b56dda16825ab954013b57941 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 18 Jul 2021 13:09:29 +0300 Subject: [PATCH 1/7] add support for Command filter --- src/commandfilter/mod.rs | 44 ++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 7 +++++-- src/macros.rs | 37 +++++++++++++++++++++++++++++++++ 3 files changed, 86 insertions(+), 2 deletions(-) create mode 100644 src/commandfilter/mod.rs diff --git a/src/commandfilter/mod.rs b/src/commandfilter/mod.rs new file mode 100644 index 00000000..9142909a --- /dev/null +++ b/src/commandfilter/mod.rs @@ -0,0 +1,44 @@ +use crate::raw; +use crate::RedisString; +use std::os::raw::c_int; +use std::ptr; + +pub struct CommandFilterContext { + pub ctx: *mut raw::RedisModuleCommandFilterCtx, +} + +impl CommandFilterContext { + pub fn new(ctx: *mut raw::RedisModuleCommandFilterCtx) -> Self { + CommandFilterContext { ctx } + } + + pub fn args_count(&self) -> usize { + unsafe { raw::RedisModule_CommandFilterArgsCount.unwrap()(self.ctx) as usize } + } + + pub fn args_get(&self, pos: usize) -> RedisString { + let arg = unsafe { + raw::RedisModule_CommandFilterArgGet.unwrap()(self.ctx, pos as c_int) + as *mut raw::RedisModuleString + }; + RedisString::new(ptr::null_mut(), arg) + } + + pub fn args_insert(&self, pos: usize, arg: RedisString) -> usize { + unsafe { + raw::RedisModule_CommandFilterArgInsert.unwrap()(self.ctx, pos as c_int, arg.inner) + as usize + } + } + + pub fn args_replace(&self, pos: usize, arg: RedisString) -> usize { + unsafe { + raw::RedisModule_CommandFilterArgReplace.unwrap()(self.ctx, pos as c_int, arg.inner) + as usize + } + } + + pub fn args_delete(&self, pos: usize) -> usize { + unsafe { raw::RedisModule_CommandFilterArgDelete.unwrap()(self.ctx, pos as c_int) as usize } + } +} diff --git a/src/lib.rs b/src/lib.rs index 04eab947..67b2e95b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,19 +17,22 @@ mod redismodule; pub mod redisraw; pub mod redisvalue; +#[cfg(feature = "experimental-api")] +mod commandfilter; mod context; pub mod key; pub mod logging; mod macros; +#[cfg(feature = "experimental-api")] +pub use crate::commandfilter::CommandFilterContext; #[cfg(feature = "experimental-api")] pub use crate::context::blocked::BlockedClient; #[cfg(feature = "experimental-api")] pub use crate::context::thread_safe::{DetachedFromClient, ThreadSafeContext}; +pub use crate::context::Context; #[cfg(feature = "experimental-api")] pub use crate::raw::NotifyEvent; - -pub use crate::context::Context; pub use crate::raw::Status; pub use crate::redismodule::*; diff --git a/src/macros.rs b/src/macros.rs index fb54eaf7..02a71436 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -82,6 +82,31 @@ macro_rules! redis_event_handler { }}; } +#[cfg(feature = "experimental-api")] +#[macro_export] +macro_rules! redis_commnad_filter { + ($ctx:expr, + $filter_handler:expr, + $filter_flags:expr) => {{ + ///////////////////// + extern "C" fn __do_filter(ctx: *mut $crate::raw::RedisModuleCommandFilterCtx) { + let context = $crate::CommandFilterContext::new(ctx); + $filter_handler(&context); + } + ///////////////////// + if unsafe { + $crate::raw::RedisModule_RegisterCommandFilter.unwrap()( + $ctx, + Some(__do_filter), + $filter_flags as c_int, + ) + } == std::ptr::null_mut() + { + return $crate::raw::Status::Err as c_int; + } + }}; +} + #[macro_export] macro_rules! redis_module { ( @@ -102,6 +127,12 @@ macro_rules! redis_module { $keystep:expr ]),* $(,)* ] $(,)* + $(filters: [ + $([ + $filter_handler:expr, + $filter_flags:expr + ]),* $(,)* + ] $(,)*)? $(event_handlers: [ $([ $(@$event_type:ident) +: @@ -162,6 +193,12 @@ macro_rules! redis_module { redis_command!(ctx, $name, $command, $flags, $firstkey, $lastkey, $keystep); )* + $( + $( + redis_commnad_filter!(ctx, $filter_handler, $filter_flags); + )* + )? + $( $( redis_event_handler!(ctx, $(raw::NotifyEvent::$event_type |)+ raw::NotifyEvent::empty(), $event_handler); From a87047fc30aac883dd809cff78ae03386b2158e2 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Sun, 18 Jul 2021 14:35:08 +0300 Subject: [PATCH 2/7] add retain string --- src/commandfilter/mod.rs | 4 ++++ src/redismodule.rs | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/commandfilter/mod.rs b/src/commandfilter/mod.rs index 9142909a..adf8e4ae 100644 --- a/src/commandfilter/mod.rs +++ b/src/commandfilter/mod.rs @@ -25,6 +25,8 @@ impl CommandFilterContext { } pub fn args_insert(&self, pos: usize, arg: RedisString) -> usize { + // retain arg to since RedisModule_CommandFilterArgInsert going to release it too + raw::string_retain_string(std::ptr::null_mut(), arg.inner); unsafe { raw::RedisModule_CommandFilterArgInsert.unwrap()(self.ctx, pos as c_int, arg.inner) as usize @@ -32,6 +34,8 @@ impl CommandFilterContext { } pub fn args_replace(&self, pos: usize, arg: RedisString) -> usize { + // retain arg to since RedisModule_CommandFilterArgInsert going to release it too + raw::string_retain_string(std::ptr::null_mut(), arg.inner); unsafe { raw::RedisModule_CommandFilterArgReplace.unwrap()(self.ctx, pos as c_int, arg.inner) as usize diff --git a/src/redismodule.rs b/src/redismodule.rs index 9ea7f64e..a0065c40 100644 --- a/src/redismodule.rs +++ b/src/redismodule.rs @@ -50,7 +50,7 @@ where } #[inline] - fn next_i64(&mut self) -> Result { + fn next_i64(&mut self) -> Result { self.next() .map_or(Err(RedisError::WrongArity), |v| v.parse_integer()) } From db9cc113a2e2805a4214a8514f17569f19d4b1be Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Mon, 19 Jul 2021 00:29:20 +0300 Subject: [PATCH 3/7] add example and test --- Cargo.toml | 4 ++++ examples/command_filter.rs | 37 +++++++++++++++++++++++++++++++++++++ tests/integration.rs | 23 +++++++++++++++++++++++ 3 files changed, 64 insertions(+) create mode 100644 examples/command_filter.rs diff --git a/Cargo.toml b/Cargo.toml index b2c6b2d2..08be24cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,10 @@ readme = "README.md" name = "hello" crate-type = ["cdylib"] +[[example]] +name = "command_filter" +crate-type = ["cdylib"] + [[example]] name = "keys_pos" crate-type = ["cdylib"] diff --git a/examples/command_filter.rs b/examples/command_filter.rs new file mode 100644 index 00000000..73ce42b1 --- /dev/null +++ b/examples/command_filter.rs @@ -0,0 +1,37 @@ +#[macro_use] +extern crate redis_module; + +use std::time::{SystemTime}; +use redis_module::{RedisString, CommandFilterContext}; + +fn filter(ctx: &CommandFilterContext) { + + // Prints every command to the console + let cmd = ctx.args_get(0); + eprint!("{} ", cmd); + let count = ctx.args_count(); + for index in 1..count { + eprint!("{} ", ctx.args_get(index)); + } + eprintln!(""); + + + // Add time field for every Hash + if let Ok("HSET") = cmd.try_as_str() { + ctx.args_insert(count, RedisString::create(std::ptr::null_mut(), "__TIME__")); + let now = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap(); + ctx.args_insert(count+1, RedisString::create(std::ptr::null_mut(), &format!("{}", now.as_millis()))); + } +} + +////////////////////////////////////////////////////// + +redis_module! { + name: "command_filter", + version: 1, + data_types: [], + commands: [], + filters:[ + [filter, 0], + ] +} diff --git a/tests/integration.rs b/tests/integration.rs index 15e24c13..e46a1870 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -48,3 +48,26 @@ fn test_keys_pos() -> Result<()> { Ok(()) } + +#[cfg(feature = "experimental-api")] +#[test] +fn test_command_filter() -> Result<()> { + let _guards = vec![start_redis_server_with_module("command_filter", 6481) + .with_context(|| "failed to start redis server")?]; + let mut con = + get_redis_connection(6481).with_context(|| "failed to connect to redis server")?; + + let res: i32 = redis::cmd("HSET") + .arg(&["mykey", "first", "Don"]) + .query(&mut con) + .with_context(|| "failed to run HSET")?; + assert_eq!(res, 2); + + let res: Vec = redis::cmd("HGETALL") + .arg(&["mykey"]) + .query(&mut con) + .with_context(|| "failed to run HSET")?; + assert_eq!(res, ["first", "Don"]); + + Ok(()) +} \ No newline at end of file From 36b8b92255836cd5fd24ff508e3a6fe52d2b55ca Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Mon, 19 Jul 2021 00:38:38 +0300 Subject: [PATCH 4/7] fix test --- tests/integration.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration.rs b/tests/integration.rs index e46a1870..0afbe7a1 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -67,7 +67,9 @@ fn test_command_filter() -> Result<()> { .arg(&["mykey"]) .query(&mut con) .with_context(|| "failed to run HSET")?; - assert_eq!(res, ["first", "Don"]); + + assert_eq!(res.len(), 4); + assert_eq!(res[0..3], ["first", "Don", "__TIME__"]); Ok(()) } \ No newline at end of file From 7800a0072ddb94a3ff82e2cb7d1601b659db65e9 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Mon, 19 Jul 2021 09:45:22 +0300 Subject: [PATCH 5/7] Update src/commandfilter/mod.rs Co-authored-by: Omer Shadmi <76992134+oshadmi@users.noreply.github.com> --- src/commandfilter/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/commandfilter/mod.rs b/src/commandfilter/mod.rs index adf8e4ae..237a6176 100644 --- a/src/commandfilter/mod.rs +++ b/src/commandfilter/mod.rs @@ -34,7 +34,7 @@ impl CommandFilterContext { } pub fn args_replace(&self, pos: usize, arg: RedisString) -> usize { - // retain arg to since RedisModule_CommandFilterArgInsert going to release it too + // retain arg to since RedisModule_CommandFilterArgReplace going to release it too raw::string_retain_string(std::ptr::null_mut(), arg.inner); unsafe { raw::RedisModule_CommandFilterArgReplace.unwrap()(self.ctx, pos as c_int, arg.inner) From 320b858eaf054ecb7d0c6ad74045cee3ce7dd92d Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Mon, 19 Jul 2021 15:34:00 +0300 Subject: [PATCH 6/7] fix review comments --- examples/command_filter.rs | 21 ++++++++++++--------- src/commandfilter/mod.rs | 33 ++++++++++++++++++++++----------- src/redisvalue.rs | 4 ++-- tests/integration.rs | 4 ++-- 4 files changed, 38 insertions(+), 24 deletions(-) diff --git a/examples/command_filter.rs b/examples/command_filter.rs index 73ce42b1..54b7abed 100644 --- a/examples/command_filter.rs +++ b/examples/command_filter.rs @@ -1,26 +1,29 @@ #[macro_use] extern crate redis_module; -use std::time::{SystemTime}; -use redis_module::{RedisString, CommandFilterContext}; +use redis_module::{CommandFilterContext, RedisString}; +use std::time::SystemTime; fn filter(ctx: &CommandFilterContext) { - // Prints every command to the console - let cmd = ctx.args_get(0); + let cmd = ctx.args_get(0).unwrap(); eprint!("{} ", cmd); let count = ctx.args_count(); for index in 1..count { - eprint!("{} ", ctx.args_get(index)); + eprint!("{} ", ctx.args_get(index).unwrap()); } eprintln!(""); - - // Add time field for every Hash + // Add time field for every Hash if let Ok("HSET") = cmd.try_as_str() { ctx.args_insert(count, RedisString::create(std::ptr::null_mut(), "__TIME__")); - let now = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap(); - ctx.args_insert(count+1, RedisString::create(std::ptr::null_mut(), &format!("{}", now.as_millis()))); + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap(); + ctx.args_insert( + count + 1, + RedisString::create(std::ptr::null_mut(), &format!("{}", now.as_millis())), + ); } } diff --git a/src/commandfilter/mod.rs b/src/commandfilter/mod.rs index 237a6176..c7af7556 100644 --- a/src/commandfilter/mod.rs +++ b/src/commandfilter/mod.rs @@ -1,5 +1,6 @@ use crate::raw; use crate::RedisString; +use crate::Status; use std::os::raw::c_int; use std::ptr; @@ -16,33 +17,43 @@ impl CommandFilterContext { unsafe { raw::RedisModule_CommandFilterArgsCount.unwrap()(self.ctx) as usize } } - pub fn args_get(&self, pos: usize) -> RedisString { + pub fn args_get(&self, pos: usize) -> Option { let arg = unsafe { raw::RedisModule_CommandFilterArgGet.unwrap()(self.ctx, pos as c_int) as *mut raw::RedisModuleString }; - RedisString::new(ptr::null_mut(), arg) + if arg.is_null() { + None + } else { + Some(RedisString::new(ptr::null_mut(), arg)) + } } - pub fn args_insert(&self, pos: usize, arg: RedisString) -> usize { - // retain arg to since RedisModule_CommandFilterArgInsert going to release it too + pub fn args_insert(&self, pos: usize, arg: RedisString) -> Status { + // retain arg since RedisModule_CommandFilterArgInsert going to release it too raw::string_retain_string(std::ptr::null_mut(), arg.inner); unsafe { raw::RedisModule_CommandFilterArgInsert.unwrap()(self.ctx, pos as c_int, arg.inner) - as usize + .into() } } - pub fn args_replace(&self, pos: usize, arg: RedisString) -> usize { - // retain arg to since RedisModule_CommandFilterArgReplace going to release it too + pub fn args_replace(&self, pos: usize, arg: RedisString) -> Status { + // retain arg since RedisModule_CommandFilterArgReplace going to release it too raw::string_retain_string(std::ptr::null_mut(), arg.inner); - unsafe { + let status = unsafe { raw::RedisModule_CommandFilterArgReplace.unwrap()(self.ctx, pos as c_int, arg.inner) - as usize + .into() + }; + + // If the string wasn't replaced we have to release it ourself + if status == Status::Err { + unsafe { raw::RedisModule_FreeString.unwrap()(std::ptr::null_mut(), arg.inner) }; } + status } - pub fn args_delete(&self, pos: usize) -> usize { - unsafe { raw::RedisModule_CommandFilterArgDelete.unwrap()(self.ctx, pos as c_int) as usize } + pub fn args_delete(&self, pos: usize) -> Status { + unsafe { raw::RedisModule_CommandFilterArgDelete.unwrap()(self.ctx, pos as c_int).into() } } } diff --git a/src/redisvalue.rs b/src/redisvalue.rs index f7cb4409..3b35d4e5 100644 --- a/src/redisvalue.rs +++ b/src/redisvalue.rs @@ -129,10 +129,10 @@ mod tests { #[test] fn from_vec() { - let v : Vec = vec![0,3,5,21,255]; + let v: Vec = vec![0, 3, 5, 21, 255]; assert_eq!( RedisValue::from(v), - RedisValue::StringBuffer(vec![0,3,5,21,255]) + RedisValue::StringBuffer(vec![0, 3, 5, 21, 255]) ); } diff --git a/tests/integration.rs b/tests/integration.rs index 0afbe7a1..fbcbfa9f 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -67,9 +67,9 @@ fn test_command_filter() -> Result<()> { .arg(&["mykey"]) .query(&mut con) .with_context(|| "failed to run HSET")?; - + assert_eq!(res.len(), 4); assert_eq!(res[0..3], ["first", "Don", "__TIME__"]); Ok(()) -} \ No newline at end of file +} From 6995ac28c66893b0b4d9bdee39a92ef75758e8c0 Mon Sep 17 00:00:00 2001 From: Guy Korland Date: Mon, 19 Jul 2021 16:39:58 +0300 Subject: [PATCH 7/7] handle failed insert --- src/commandfilter/mod.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/commandfilter/mod.rs b/src/commandfilter/mod.rs index c7af7556..bfe19349 100644 --- a/src/commandfilter/mod.rs +++ b/src/commandfilter/mod.rs @@ -32,10 +32,16 @@ impl CommandFilterContext { pub fn args_insert(&self, pos: usize, arg: RedisString) -> Status { // retain arg since RedisModule_CommandFilterArgInsert going to release it too raw::string_retain_string(std::ptr::null_mut(), arg.inner); - unsafe { + let status = unsafe { raw::RedisModule_CommandFilterArgInsert.unwrap()(self.ctx, pos as c_int, arg.inner) .into() + }; + + // If the string wasn't inserted we have to release it ourself + if status == Status::Err { + unsafe { raw::RedisModule_FreeString.unwrap()(std::ptr::null_mut(), arg.inner) }; } + status } pub fn args_replace(&self, pos: usize, arg: RedisString) -> Status {