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,
        ]
    },
    

[VoletiRam]

Please review when you have time @KarthikSubbarao @dmitrypol

[KarthikSubbarao]

Thanks for looking into this @VoletiRam . Providing some feedback on the interface below..

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

@VoletiRam - We do not need the inner "callbacks", right? Similar to how commands accept arguments, we can have it this way:

    auth: {
        [
            my_module_auth,
        ]
    }

See: https://github.com/valkey-io/valkeymodule-rs/blob/main/examples/hello.rs#L30

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

let mut blocked_client = block_client_on_auth(ctx, my_auth_reply);

We can follow the convention of normal blocking clients and update this to be:

ctx.block_client_on_auth(my_auth_reply)

See: https://github.com/valkey-io/valkeymodule-rs/blob/main/examples/block.rs#L9

Instead of an isolated function, you can impl this function on the Context structure in this file: https://github.com/valkey-io/valkeymodule-rs/blob/main/src/context/blocked.rs#L21

// 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)

}

A few comments about the BlockedClient:

1. The block of code below "Store the auth result here" here will be called from a background thread / blocking operation. So, it would be called from within the spawn thread block and once the blocking operation is done, the client will be unblocked. You can update your example code to do the same ((1) set the private data/auth result inside the spawn block and (2) move the BlockedClient structure into this spawn block so that it gets dropped there once we go out of context in the spawn block).
2. The priv data is provided through the ValkeyModule_UnblockClient call, so the core essentially owns this data once the Unblock client API is called. With this being the case, it would be good to follow this pattern. When the BlockedClient is dropped, you can transfer ownership of the private data to the core and clear any reference to this from the BlockedClient structure. We can remove the get_private_data function because this private data will be provided once the auth reply callback is invoked from the core.