Added handling for unloading an auth module

5 files changed, 73 insertions, 36 deletions

diff --git a/src/module.c b/src/module.c
index 5b4cab69f..17fca5e1d 100644
--- a/src/module.c
+++ b/src/module.c
@@ -350,7 +350,7 @@ unsigned long long ModulesInHooks = 0; /* Total number of modules in hooks
                                           callbacks right now. */
 
 /* Data structures related to the redis module users */
-typedef void (*RedisModuleUserChangedFunc) (void *privdata);
+typedef void (*RedisModuleUserChangedFunc) (RedisModuleCtx ctx, void *privdata);
 
 typedef struct RedisModuleUser {
     user *user; /* Reference to the real redis user */
@@ -358,10 +358,12 @@ typedef struct RedisModuleUser {
 
 typedef struct RedisModuleAuthCtx {
     RedisModuleUserChangedFunc callback; /* Callback when user is changed or 
-                                         /* deleted. */
+                                          * deleted. */
     void *privdata;                      /* Private data for the callback */
-    client *authenticated_client;        /* A reference to client that was
-                                         /* authenticated */
+    client *authenticated_client;        /* A reference to the client that was
+                                          * authenticated */
+    RedisModule *module;                 /* Reference to the module that
+                                          * authenticated the client */
 } RedisModuleAuthCtx;
 
 
@@ -5140,14 +5142,48 @@ int RM_GetTimerInfo(RedisModuleCtx *ctx, RedisModuleTimerID id, uint64_t *remain
 void moduleNotifyUserChanged(client *c) {
     RedisModuleAuthCtx *auth_ctx = (RedisModuleAuthCtx *) c->auth_ctx;
     if (auth_ctx) {
-        auth_ctx->callback(auth_ctx->privdata);
+        RedisModuleCtx ctx = REDISMODULE_CTX_INIT;
+        ctx.module = auth_ctx->module;
+        ctx.client = moduleFreeContextReusedClient;
+
+        auth_ctx->callback(ctx, auth_ctx->privdata);
+
         zfree(auth_ctx);
+        c->auth_ctx = NULL;
+        moduleFreeContext(&ctx);
+    }
+}
+
+
+void revokeClientAuthentication(client *c) {
+    /* Fire the client changed handler now in case we are unloading the module
+     * and need to cleanup.  */
+    moduleNotifyUserChanged(c);
+
+    c->user = DefaultUser;
+    c->authenticated = 0;
+    freeClientAsync(c);
+}
+
+/* Cleanup all clients with an auth_ctx to prevent leaking. */
+void moduleFreeAuthenticatedClients(RedisModule *module) {
+    listIter li;
+    listNode *ln;
+    listRewind(server.clients,&li);
+    while ((ln = listNext(&li)) != NULL) {
+        client *c = listNodeValue(ln);
+        if (!c->auth_ctx) continue;
+
+        RedisModuleAuthCtx *auth_ctx = (RedisModuleAuthCtx *) c->auth_ctx;
+        if (auth_ctx->module == module) { 
+            revokeClientAuthentication(c);
+        }
     }
 }
 
 static int authenticateClientWithUser(RedisModuleCtx *ctx, user *user, RedisModuleAuthCtx *auth_ctx) {
     if (auth_ctx && auth_ctx->authenticated_client) {
-        /* Prevent basic misuse of Auth contexts */
+        /* Prevent basic misuse of the auth context */
         return REDISMODULE_ERR;
     }
     
@@ -5163,6 +5199,7 @@ static int authenticateClientWithUser(RedisModuleCtx *ctx, user *user, RedisModu
     if (auth_ctx) {
         ctx->client->auth_ctx = auth_ctx;
         auth_ctx->authenticated_client = ctx->client;
+        auth_ctx->module = ctx->module;
     }
 
     return REDISMODULE_OK;
@@ -5202,14 +5239,14 @@ int RM_SetModuleUserACL(RedisModuleUser *user, const char* acl) {
 
 /* Authenticate the current contexts user with the provided module user. 
  * Throws an error if the user is disabled or the AuthCtx is misued */
-int RM_AuthenticateClientWithUser(RedisModuleCtx *ctx, RedisModuleUser *module_user, RedisModuleAuthCtx *auth_ctx) {
+int RM_AuthClientWithUser(RedisModuleCtx *ctx, RedisModuleUser *module_user, RedisModuleAuthCtx *auth_ctx) {
     return authenticateClientWithUser(ctx, module_user->user, auth_ctx);
 }
 
 /* Authenticate the current contexts user with the provided redis acl user. 
  * Throws an error if the user is disabled, the user doesn't exit,
  * or the AuthCtx is misused. */
-int RM_AuthenticateClientWithACLUser(RedisModuleCtx *ctx, const char *name, size_t len, RedisModuleAuthCtx *auth_ctx) {
+int RM_AuthClientWithACLUser(RedisModuleCtx *ctx, const char *name, size_t len, RedisModuleAuthCtx *auth_ctx) {
     user *acl_user = ACLGetUserByName(name, len);
 
     if (!acl_user) {
@@ -5231,8 +5268,7 @@ RedisModuleAuthCtx *RM_CreateAuthCtx(RedisModuleUserChangedFunc callback, void *
 /* Revokes the authentication that was granted during the specified
  * AuthCtx. The method is not thread safe. */
 void RM_RevokeAuthentication(RedisModuleAuthCtx *ctx) {
-    /* Revoking authentication frees the client today. */
-    freeClientAsync(ctx->authenticated_client);
+    revokeClientAuthentication(ctx->authenticated_client);
 }
 
 /* --------------------------------------------------------------------------
@@ -6693,7 +6729,7 @@ int moduleUnload(sds name) {
     moduleUnregisterSharedAPI(module);
     moduleUnregisterUsedAPI(module);
     moduleUnregisterFilters(module);
-
+    moduleFreeAuthenticatedClients(module);
     /* Remove any notification subscribers this module might have */
     moduleUnsubscribeNotifications(module);
     moduleUnsubscribeAllServerEvents(module);
@@ -7104,6 +7140,6 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(FreeModuleUser);
     REGISTER_API(RevokeAuthentication);
     REGISTER_API(CreateAuthCtx);
-    REGISTER_API(AuthenticateClientWithACLUser);
-    REGISTER_API(AuthenticateClientWithUser);
+    REGISTER_API(AuthClientWithACLUser);
+    REGISTER_API(AuthClientWithUser);
 }
diff --git a/src/modules/helloacl.c b/src/modules/helloacl.c
index 0eb953329..95dcb09c8 100644
--- a/src/modules/helloacl.c
+++ b/src/modules/helloacl.c
@@ -75,13 +75,14 @@ int ResetCommand_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int
 
 /* Callback handler for user changes, use this to notify a module of 
  * changes to users authenticated by the module */
-void HelloACL_UserChanged(void *privdata) {
+void HelloACL_UserChanged(RedisModuleCtx *ctx, void *privdata) {
     REDISMODULE_NOT_USED(privdata);
+    REDISMODULE_NOT_USED(ctx);
     global_auth_ctx = NULL;
 }
 
 /* HELLOACL.AUTHGLOBAL 
- * Synchronously assigns a module user to the current context. */
+ * Synchronously assigns a module user to the current context client. */
 int AuthGlobalCommand_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     REDISMODULE_NOT_USED(argv);
     REDISMODULE_NOT_USED(argc);
@@ -90,8 +91,8 @@ int AuthGlobalCommand_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
         return RedisModule_ReplyWithError(ctx, "Global user currently used");    
     }
 
-    RedisModuleAuthCtx *auth_ctx = RedisModule_CreateAuthCtx(HelloACL_UserChanged,NULL);
-    RedisModule_AuthenticateClientWithUser(ctx, global, auth_ctx);
+    RedisModuleAuthCtx *auth_ctx = RedisModule_CreateAuthCtx(HelloACL_UserChanged, NULL);
+    RedisModule_AuthClientWithUser(ctx, global, auth_ctx);
     global_auth_ctx = auth_ctx;
 
     return RedisModule_ReplyWithSimpleString(ctx, "OK");
@@ -108,7 +109,7 @@ int HelloACL_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     RedisModuleString *user_string = RedisModule_GetBlockedClientPrivateData(ctx);
     const char *name = RedisModule_StringPtrLen(user_string, &length);
 
-    if (RedisModule_AuthenticateClientWithACLUser(ctx, name, length, NULL) == 
+    if (RedisModule_AuthClientWithACLUser(ctx, name, length, NULL) == 
             REDISMODULE_ERR) {
         return RedisModule_ReplyWithError(ctx, "Invalid Username or password");    
     }
@@ -122,7 +123,7 @@ int HelloACL_Timeout(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     return RedisModule_ReplyWithSimpleString(ctx, "Request timedout");
 }
 
-/* Private data frees data for HELLOACL.AUTHASYNC command. */
+/* FreeData callback frees private data for HELLOACL.AUTHASYNC command. */
 void HelloACL_FreeData(RedisModuleCtx *ctx, void *privdata) {
     REDISMODULE_NOT_USED(ctx);
     RedisModule_FreeString(NULL, privdata);
@@ -141,12 +142,12 @@ void *HelloACL_ThreadMain(void *args) {
 
 /* HELLOACL.AUTHASYNC 
  * Asynchronously assigns an ACL user to the current context. */
-*/
 int AuthAsyncCommand_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     if (argc != 2) return RedisModule_WrongArity(ctx);
 
     pthread_t tid;
-    RedisModuleBlockedClient *bc = RedisModule_BlockClient(ctx, HelloACL_Reply, HelloACL_Timeout, HelloACL_FreeData, TIMEOUT_TIME);
+    RedisModuleBlockedClient *bc = RedisModule_BlockClient(ctx, HelloACL_Reply, 
+        HelloACL_Timeout, HelloACL_FreeData, TIMEOUT_TIME);
     
 
     void **targs = RedisModule_Alloc(sizeof(void*)*2);
@@ -167,23 +168,23 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
     REDISMODULE_NOT_USED(argv);
     REDISMODULE_NOT_USED(argc);
 
-    if (RedisModule_Init(ctx,"helloacl",1,REDISMODULE_APIVER_1)
+    if (RedisModule_Init(ctx, "helloacl", 1, REDISMODULE_APIVER_1)
         == REDISMODULE_ERR) return REDISMODULE_ERR;
 
-    if (RedisModule_CreateCommand(ctx,"helloacl.reset",
-        ResetCommand_RedisCommand,"",0,0,0) == REDISMODULE_ERR)
+    if (RedisModule_CreateCommand(ctx, "helloacl.reset",
+        ResetCommand_RedisCommand, "", 0, 0, 0) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
-    if (RedisModule_CreateCommand(ctx,"helloacl.revoke",
-        RevokeCommand_RedisCommand,"",0,0,0) == REDISMODULE_ERR)
+    if (RedisModule_CreateCommand(ctx, "helloacl.revoke",
+        RevokeCommand_RedisCommand, "", 0, 0, 0) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
-    if (RedisModule_CreateCommand(ctx,"helloacl.authglobal",
-        AuthGlobalCommand_RedisCommand,"no-auth",0,0,0) == REDISMODULE_ERR)
+    if (RedisModule_CreateCommand(ctx, "helloacl.authglobal",
+        AuthGlobalCommand_RedisCommand, "no-auth", 0, 0, 0) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
     if (RedisModule_CreateCommand(ctx,"helloacl.authasync",
-        AuthAsyncCommand_RedisCommand,"no-auth",0,0,0) == REDISMODULE_ERR)
+        AuthAsyncCommand_RedisCommand, "no-auth", 0, 0, 0) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
     global = RedisModule_CreateModuleUser("global");
diff --git a/src/redismodule.h b/src/redismodule.h
index cce3ad036..9b7b1b04a 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -411,7 +411,7 @@ typedef void (*RedisModuleTimerProc)(RedisModuleCtx *ctx, void *data);
 typedef void (*RedisModuleCommandFilterFunc) (RedisModuleCommandFilterCtx *filter);
 typedef void (*RedisModuleForkDoneHandler) (int exitcode, int bysignal, void *user_data);
 typedef void (*RedisModuleInfoFunc)(RedisModuleInfoCtx *ctx, int for_crash_report);
-typedef void (*RedisModuleUserChangedFunc)(void *privdata);
+typedef void (*RedisModuleUserChangedFunc)(RedisModuleCtx *ctx, void *privdata);
 
 #define REDISMODULE_TYPE_METHOD_VERSION 2
 typedef struct RedisModuleTypeMethods {
@@ -639,8 +639,8 @@ int REDISMODULE_API_FUNC(RedisModule_KillForkChild)(int child_pid);
 RedisModuleUser *REDISMODULE_API_FUNC(RedisModule_CreateModuleUser)(const char *name);
 void REDISMODULE_API_FUNC(RedisModule_FreeModuleUser)(RedisModuleUser *user);
 int REDISMODULE_API_FUNC(RedisModule_SetModuleUserACL)(RedisModuleUser *user, const char* acl);
-int REDISMODULE_API_FUNC(RedisModule_AuthenticateClientWithACLUser)(RedisModuleCtx *ctx, const char* name, size_t len, RedisModuleAuthCtx *auth_ctx);
-int REDISMODULE_API_FUNC(RedisModule_AuthenticateClientWithUser)(RedisModuleCtx *ctx, RedisModuleUser *module_user, RedisModuleAuthCtx *auth_ctx);
+int REDISMODULE_API_FUNC(RedisModule_AuthClientWithACLUser)(RedisModuleCtx *ctx, const char* name, size_t len, RedisModuleAuthCtx *auth_ctx);
+int REDISMODULE_API_FUNC(RedisModule_AuthClientWithUser)(RedisModuleCtx *ctx, RedisModuleUser *module_user, RedisModuleAuthCtx *auth_ctx);
 void REDISMODULE_API_FUNC(RedisModule_RevokeAuthentication)(RedisModuleAuthCtx *ctx);
 RedisModuleAuthCtx *REDISMODULE_API_FUNC(RedisModule_CreateAuthCtx)(RedisModuleUserChangedFunc callback, void *privdata);
 #endif
@@ -858,8 +858,8 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(SetModuleUserACL);
     REDISMODULE_GET_API(RevokeAuthentication);
     REDISMODULE_GET_API(CreateAuthCtx);
-    REDISMODULE_GET_API(AuthenticateClientWithACLUser);
-    REDISMODULE_GET_API(AuthenticateClientWithUser);
+    REDISMODULE_GET_API(AuthClientWithACLUser);
+    REDISMODULE_GET_API(AuthClientWithUser);
 #endif
 
     if (RedisModule_IsModuleNameBusy && RedisModule_IsModuleNameBusy(name)) return REDISMODULE_ERR;
diff --git a/src/server.c b/src/server.c
index 5f5e98b15..94c48dc76 100644
--- a/src/server.c
+++ b/src/server.c
@@ -3425,7 +3425,7 @@ int processCommand(client *c) {
                            DefaultUser->flags & USER_FLAG_DISABLED) &&
                         !c->authenticated;
     if (auth_required) {
-        /* AUTH and HELLO and no auth modules are valid even in
+        /* AUTH, HELLO, and no-auth modules are valid even in
          * non-authenticated state. */
         if (!(c->cmd->flags & CMD_NO_AUTH)) {
             flagTransaction(c);
diff --git a/src/server.h b/src/server.h
index 1aa8d5f1f..ba78c2e28 100644
--- a/src/server.h
+++ b/src/server.h
@@ -893,7 +893,7 @@ typedef struct client {
     list *pubsub_patterns;  /* patterns a client is interested in (SUBSCRIBE) */
     sds peerid;             /* Cached peer ID. */
     listNode *client_list_node; /* list node in client list */
-    void *auth_ctx;         /* Structured used to track module authentication */
+    void *auth_ctx;         /* Opaque structure to track module auth */
 
     /* If this client is in tracking mode and this field is non zero,
      * invalidation messages for keys fetched by this client will be send to