Create Interface for Modules to Register a Custom Auth Module Callback

[KarthikSubbarao]

Create Interface for Modules to Register a Custom Auth Module Callback.

In Valkey (even on 7.2), we have the ability to register custom authentication callbacks from Modules using the ValkeyModule_RegisterAuthCallback Module API.

https://valkey.io/topics/modules-api-ref/#ValkeyModule_RegisterAuthCallback

For Rust Modules, we can abstract away a lot of interactions by implementing a macro that allows both blocking and non blocking authentication implementations in the module.

High level valkeymodule-rs Interface updates:

• Have a single new macro (valkey_module_auth!) to allow registering authentication callbacks. This macro should also be embedded inside the valkey_module! macro such that it can be optionally used by Valkey Rust Modules.
• valkey_module_auth! should accept both a callback argument (the function signature should be safe for Rust Modules to implement, ie - no unsafe operations) and should allow modules to implement both blocking and non blocking callbacks.
• Extend the BlockedClient interface in blocked.rs to allow a new mode of blocking clients - block on authentication by using the ValkeyModule_BlockClientOnAuth
• Extend the BlockedClient interface in blocked.rs to allow setting and getting private info. This would be used in blocking authentication to store the result of the authentication and use it when the client is unblocked.

See the C Module example of this here: https://github.com/valkey-io/valkey/blob/unstable/tests/modules/auth.c#L83-#L221

[VoletiRam]

I will be working on this issue. Please assign it to me.

Here is the detailed approach to support custom authentication callbacks in modules. Please check and suggest if something needs to be added/changed.

We intend to create a macro called valkey_module_auth! that will accept an array of callbacks. Inside the macro, we will register each authentication callback by looping through the array and using a safe wrapper of ValkeyModule_RegisterAuthCallback to register them with the core server. Here is how the interface will look:

pub type AuthCallback = fn(ctx: &Context, username: &str, password: &str) -> Result<bool, String>;

pub fn register_authcallback(ctx: &Context, callback: AuthCallback) -> raw::Status {

The valkey_module_auth! macro:

// The auth macro
#[macro_export]
macro_rules! valkey_module_auth {
    (callbacks: [$($callback:expr),* $(,)?]) => {
        pub fn register_auth_callbacks(ctx: &Context) -> Result<(), String> {
            $(
                match register_authcallback(ctx, $callback) {

Integrating the auth callback registration with the existing valkey_module! macro as an optional feature:

#[macro_export]
macro_rules! valkey_module {
    (
        name: $module_name:expr,
        version: $module_version:expr,
        allocator: ($allocator_type:ty, $allocator_init:expr),
        $(auth: {
            callbacks: [$($callback:expr),* $(,)*]
        },)?
        commands: [/* ... */]
    ) => {
        // ... existing implementation ...

        // Add auth callback registration
        $(
            valkey_module_auth! {
                callbacks: [$($callback),*]
            }
            
            if let Err(e) = register_auth_callbacks(&context) {
                context.log_warning(&format!("Failed to register auth callbacks: {}", e));
                return raw::Status::Err as c_int;
            }
        )?

To allow both blocking and non-blocking authentication in callbacks, we will provide a block_client_on_auth wrapper for ValkeyModule_BlockClientOnAuth that accepts a reply authentication callback from module authors and registers it with the server.
The interface for block_client_on_auth will look like this:

pub fn block_client_on_auth(ctx: &Context, auth_callback: AuthCallback) -> BlockedClient {
    

We will also provide set_private_data and get_private_data APIs in Blocked.rs to store and retrieve the authentication results at the module level.

The APIs

// Store private data at module level
    pub fn set_private_data<T: Send + Sync + 'static>(&mut self, data: T) {
        self.private_data = Some(Box::new(data));
    }

    // Retrieve private data at module level
    pub fn get_private_data<T: 'static>(&self) -> Option<&T> {
        self.private_data
            .as_ref()
            .and_then(|data| data.downcast_ref::<T>())
    }
}

Example for end user

// This is the callback registered with the module
fn my_module_auth(ctx: &Context, username: &str, password: &str) -> Result<bool, String> {
    // Create blocked client with the reply callback that will be called later by core
    let mut blocked_client = block_client_on_auth(ctx, my_auth_reply);
  
    // Do their async auth processing
    // This could be calling external service, database, etc.
    std::thread::spawn(move || {
        // Their auth logic here
        // When done, they would call something like unblock_client()
        // Core will then call my_auth_reply
    });
    
        // Store the auth result here and use it when they unblock_client() to authenticate client
    blocked_client.set_private_data(MyAuthResult {
        success: false,
        details: format!("Processing auth for user: {}", username),
    });

    // Return true to indicate we're handling the auth
    Ok(true)
}

// Reply callback with same safe signature
fn my_auth_reply(ctx: &Context, username: &str, password: &str) -> Result<bool, String> {
    // Their unblock logic here, called by core later
    Ok(true) // or Ok(false)/Err based on their auth result
}

// Module registration
valkey_module! {
    name: "auth_example",
    version: 1,
    allocator: (ValkeyAlloc, ValkeyAlloc),
    auth: {
        callbacks: [
            my_module_auth,
        ]
    },