From 13d5d1d27d76c12ca4a724be444df54e56c5d2cd Mon Sep 17 00:00:00 2001 From: sachnvmurthy Date: Thu, 27 Nov 2025 11:25:46 +0530 Subject: [PATCH 1/2] added a post auth hook --- Cargo.toml | 4 +++ examples/post_auth.rs | 34 ++++++++++++++++++++++ src/macros.rs | 66 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+) create mode 100644 examples/post_auth.rs diff --git a/Cargo.toml b/Cargo.toml index 8ca4cf2..356d0f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -150,6 +150,10 @@ crate-type = ["cdylib"] name = "preload" crate-type = ["cdylib"] +[[example]] +name = "post_auth" +crate-type = ["cdylib"] + [[example]] name = "subcmd" crate-type = ["cdylib"] diff --git a/examples/post_auth.rs b/examples/post_auth.rs new file mode 100644 index 0000000..fa64783 --- /dev/null +++ b/examples/post_auth.rs @@ -0,0 +1,34 @@ +use std::sync::{LazyLock, Mutex}; +use valkey_module::alloc::ValkeyAlloc; + +use valkey_module::{valkey_module, Context, ValkeyResult, ValkeyString, ValkeyValue}; + +static LAST_AUTH_USER: LazyLock> = LazyLock::new(|| Mutex::new((String::new(), String::new()))); + +// This callback will be scheduled to run after the auth handler completes. +// It receives retained `ValkeyString` previous and new usernames. +fn post_auth_callback(ctx: &Context, prev_user: ValkeyString, new_user: ValkeyString) { + ctx.log_notice(&format!("post_auth: prev_user='{}', new_user='{}'", prev_user, new_user)); + let mut lock = LAST_AUTH_USER.lock().unwrap(); + *lock = (prev_user.to_string_lossy(), new_user.to_string_lossy()); +} + +fn whoami(_ctx: &Context, _args: Vec) -> ValkeyResult { + let (prev, new) = LAST_AUTH_USER.lock().unwrap().clone(); + if new.is_empty() { + Ok(ValkeyValue::Null) + } else { + Ok(ValkeyValue::Array(vec![ValkeyValue::BulkString(prev), ValkeyValue::BulkString(new)])) + } +} + +valkey_module! { + name: "post_auth_example", + version: 1, + allocator: (ValkeyAlloc, ValkeyAlloc), + data_types: [], + post_auth: [ post_auth_callback ], + commands: [ + ["whoami_postauth", whoami, "readonly", 0, 0, 0], + ] +} diff --git a/src/macros.rs b/src/macros.rs index 4836682..5eeee67 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -133,6 +133,9 @@ macro_rules! valkey_module { $(auth: [ $($auth_callback:expr),* $(,)* ],)? + $(post_auth: [ + $($post_auth_callback:expr),* $(,)* + ],)? $(acl_categories: [ $($acl_category:expr),* $(,)* ])? @@ -461,6 +464,10 @@ macro_rules! valkey_module { raw::register_info_function(ctx, Some(__info_func)); + $( + $crate::valkey_module_post_auth!($module_name, ctx, $($post_auth_callback),*); + )? + $( $crate::valkey_module_auth!($module_name, ctx, $($auth_callback),*); )? @@ -576,3 +583,62 @@ macro_rules! valkey_module_auth { )* }; } + +/// Registers post-authentication callbacks with the Valkey module +/// +/// These callbacks are invoked after the authentication chain finishes. The macro +/// registers lightweight auth callbacks that schedule a post-notification job so +/// the provided Rust callback runs outside the auth context. The user callback +/// signature should be `fn(&Context, ValkeyString, ValkeyString)` (prev_user, new_user). +#[macro_export] +macro_rules! valkey_module_post_auth { + ($module_name:expr, $ctx:expr, $($post_auth_callback:expr),* $(,)*) => { + $( + { + paste::paste! { + extern "C" fn [<__do_post_auth_ $module_name _ $post_auth_callback>] ( + ctx: *mut $crate::raw::RedisModuleCtx, + username: *mut $crate::raw::RedisModuleString, + _password: *mut $crate::raw::RedisModuleString, + _err: *mut *mut $crate::raw::RedisModuleString, + ) -> std::os::raw::c_int { + let context = $crate::Context::new(ctx); + let ctx_ptr = unsafe { std::ptr::NonNull::new_unchecked(ctx) }; + let username_rs = $crate::ValkeyString::new(Some(ctx_ptr), username); + let username_owned = username_rs.to_string_lossy(); + // Get previous username before auth + let prev_username = match context.get_client_username() { + Ok(u) => u.to_string_lossy(), + Err(_) => "".to_string(), + }; + + // schedule a post-notification job so the callback runs outside + // the auth handler (deferred), and can perform safe operations. + let _ = context.add_post_notification_job(move |ctx| { + let prev = $crate::ValkeyString::create_and_retain(&prev_username); + let newu = $crate::ValkeyString::create_and_retain(&username_owned); + $post_auth_callback(ctx, prev, newu); + }); + + // Do not interfere with the authentication chain. + $crate::AUTH_NOT_HANDLED + } + + #[cfg(not(any( + feature = "min-redis-compatibility-version-7-0", + feature = "min-valkey-compatibility-version-8-0" + )))] + compile_error!("Post-auth callbacks require Redis 7.0 or Valkey 8.0 and above"); + + #[cfg(any( + feature = "min-redis-compatibility-version-7-0", + feature = "min-valkey-compatibility-version-8-0" + ))] + unsafe { + $crate::raw::RedisModule_RegisterAuthCallback.expect("RedisModule_RegisterAuthCallback should exist on Redis/Valkey 7.0 and above")($ctx, Some([<__do_post_auth_ $module_name _ $post_auth_callback>])) + } + } + } + )* + }; +} From e47520dedda76ab07bd3ed08394612cd6aa39ccb Mon Sep 17 00:00:00 2001 From: sachnvmurthy Date: Thu, 27 Nov 2025 15:00:39 +0530 Subject: [PATCH 2/2] fixed min compatibility --- src/macros.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/macros.rs b/src/macros.rs index 5eeee67..b3263f7 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -609,15 +609,15 @@ macro_rules! valkey_module_post_auth { // Get previous username before auth let prev_username = match context.get_client_username() { Ok(u) => u.to_string_lossy(), - Err(_) => "".to_string(), + Err(_) => "default".to_string(), }; // schedule a post-notification job so the callback runs outside // the auth handler (deferred), and can perform safe operations. let _ = context.add_post_notification_job(move |ctx| { let prev = $crate::ValkeyString::create_and_retain(&prev_username); - let newu = $crate::ValkeyString::create_and_retain(&username_owned); - $post_auth_callback(ctx, prev, newu); + let new_user = $crate::ValkeyString::create_and_retain(&username_owned); + $post_auth_callback(ctx, prev, new_user); }); // Do not interfere with the authentication chain. @@ -625,13 +625,13 @@ macro_rules! valkey_module_post_auth { } #[cfg(not(any( - feature = "min-redis-compatibility-version-7-0", + feature = "min-redis-compatibility-version-7-2", feature = "min-valkey-compatibility-version-8-0" )))] - compile_error!("Post-auth callbacks require Redis 7.0 or Valkey 8.0 and above"); + compile_error!("Post-auth callbacks require Redis 7.2 or Valkey 8.0 and above"); #[cfg(any( - feature = "min-redis-compatibility-version-7-0", + feature = "min-redis-compatibility-version-7-2", feature = "min-valkey-compatibility-version-8-0" ))] unsafe {