CommandFilter API: REDISMODULE_CMDFILTER_NOSELF.

Add a flag to automatically protect filters from being called
recursively by their own module.

4 files changed, 62 insertions, 15 deletions

diff --git a/src/module.c b/src/module.c
index 52ee977ca..e6fcd1765 100644
--- a/src/module.c
+++ b/src/module.c
@@ -50,6 +50,7 @@ struct RedisModule {
     list *usedby;   /* List of modules using APIs from this one. */
     list *using;    /* List of modules we use some APIs of. */
     list *filters;  /* List of filters the module has registered. */
+    int in_call;    /* RM_Call() nesting level */
 };
 typedef struct RedisModule RedisModule;
 
@@ -283,6 +284,8 @@ typedef struct RedisModuleCommandFilter {
     RedisModule *module;
     /* Filter callback function */
     RedisModuleCommandFilterFunc callback;
+    /* REDISMODULE_CMDFILTER_* flags */
+    int flags;
 } RedisModuleCommandFilter;
 
 /* Registered filters */
@@ -2754,6 +2757,8 @@ RedisModuleCallReply *RM_Call(RedisModuleCtx *ctx, const char *cmdname, const ch
     c->db = ctx->client->db;
     c->argv = argv;
     c->argc = argc;
+    if (ctx->module) ctx->module->in_call++;
+
     /* We handle the above format error only when the client is setup so that
      * we can free it normally. */
     if (argv == NULL) goto cleanup;
@@ -2820,6 +2825,7 @@ RedisModuleCallReply *RM_Call(RedisModuleCtx *ctx, const char *cmdname, const ch
     autoMemoryAdd(ctx,REDISMODULE_AM_REPLY,reply);
 
 cleanup:
+    if (ctx->module) ctx->module->in_call--;
     freeClient(c);
     return reply;
 }
@@ -4859,17 +4865,27 @@ int moduleUnregisterFilters(RedisModule *module) {
  *
  * Note that in the above use case, if `MODULE.SET` itself uses
  * `RedisModule_Call()` the filter will be applied on that call as well.  If
- * that is not desired, the module itself is responsible for maintaining a flag
- * to identify and avoid this form of re-entrancy.
+ * that is not desired, the `REDISMODULE_CMDFILTER_NOSELF` flag can be set when
+ * registering the filter.
+ *
+ * The `REDISMODULE_CMDFILTER_NOSELF` flag prevents execution flows that
+ * originate from the module's own `RM_Call()` from reaching the filter.  This
+ * flag is effective for all execution flows, including nested ones, as long as
+ * the execution begins from the module's command context or a thread-safe
+ * context that is associated with a blocking command.
+ *
+ * Detached thread-safe contexts are *not* associated with the module and cannot
+ * be protected by this flag.
  *
  * If multiple filters are registered (by the same or different modules), they
  * are executed in the order of registration.
  */
 
-RedisModuleCommandFilter *RM_RegisterCommandFilter(RedisModuleCtx *ctx, RedisModuleCommandFilterFunc callback) {
+RedisModuleCommandFilter *RM_RegisterCommandFilter(RedisModuleCtx *ctx, RedisModuleCommandFilterFunc callback, int flags) {
     RedisModuleCommandFilter *filter = zmalloc(sizeof(*filter));
     filter->module = ctx->module;
     filter->callback = callback;
+    filter->flags = flags;
 
     listAddNodeTail(moduleCommandFilters, filter);
     listAddNodeTail(ctx->module->filters, filter);
@@ -4910,6 +4926,12 @@ void moduleCallCommandFilters(client *c) {
     while((ln = listNext(&li))) {
         RedisModuleCommandFilter *f = ln->value;
 
+        /* Skip filter if REDISMODULE_CMDFILTER_NOSELF is set and module is
+         * currently processing a command.
+         */
+        if ((f->flags & REDISMODULE_CMDFILTER_NOSELF) && f->module->in_call) continue;
+
+        /* Call filter */
         f->callback(&filter);
     }
 
diff --git a/src/modules/hellofilter.c b/src/modules/hellofilter.c
index 9cd440df2..448e12983 100644
--- a/src/modules/hellofilter.c
+++ b/src/modules/hellofilter.c
@@ -8,7 +8,7 @@ static RedisModuleString *log_key_name;
 static const char log_command_name[] = "hellofilter.log";
 static const char ping_command_name[] = "hellofilter.ping";
 static const char unregister_command_name[] = "hellofilter.unregister";
-static int in_module = 0;
+static int in_log_command = 0;
 
 static RedisModuleCommandFilter *filter = NULL;
 
@@ -57,7 +57,7 @@ int HelloFilter_LogCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int ar
     RedisModule_CloseKey(log);
     RedisModule_FreeString(ctx, s);
 
-    in_module = 1;
+    in_log_command = 1;
 
     size_t cmdlen;
     const char *cmdname = RedisModule_StringPtrLen(argv[1], &cmdlen);
@@ -69,14 +69,14 @@ int HelloFilter_LogCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int ar
         RedisModule_ReplyWithSimpleString(ctx, "Unknown command or invalid arguments");
     }
 
-    in_module = 0;
+    in_log_command = 0;
 
     return REDISMODULE_OK;
 }
 
 void HelloFilter_CommandFilter(RedisModuleCommandFilterCtx *filter)
 {
-    if (in_module) return;  /* don't process our own RM_Call() */
+    if (in_log_command) return;  /* don't process our own RM_Call() from HelloFilter_LogCommand() */
 
     /* Fun manipulations:
      * - Remove @delme
@@ -120,12 +120,14 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
     if (RedisModule_Init(ctx,"hellofilter",1,REDISMODULE_APIVER_1)
             == REDISMODULE_ERR) return REDISMODULE_ERR;
 
-    if (argc != 1) {
+    if (argc != 2) {
         RedisModule_Log(ctx, "warning", "Log key name not specified");
         return REDISMODULE_ERR;
     }
 
+    long long noself = 0;
     log_key_name = RedisModule_CreateStringFromString(ctx, argv[0]);
+    RedisModule_StringToLongLong(argv[1], &noself);
 
     if (RedisModule_CreateCommand(ctx,log_command_name,
                 HelloFilter_LogCommand,"write deny-oom",1,1,1) == REDISMODULE_ERR)
@@ -139,7 +141,8 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
                 HelloFilter_UnregisterCommand,"write deny-oom",1,1,1) == REDISMODULE_ERR)
             return REDISMODULE_ERR;
 
-    if ((filter = RedisModule_RegisterCommandFilter(ctx, HelloFilter_CommandFilter))
+    if ((filter = RedisModule_RegisterCommandFilter(ctx, HelloFilter_CommandFilter, 
+                    noself ? REDISMODULE_CMDFILTER_NOSELF : 0))
             == NULL) return REDISMODULE_ERR;
 
     return REDISMODULE_OK;
diff --git a/src/redismodule.h b/src/redismodule.h
index 3314278a1..86d190a27 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -132,6 +132,11 @@
  * of timers that are going to expire, sorted by expire time. */
 typedef uint64_t RedisModuleTimerID;
 
+/* CommandFilter Flags */
+
+/* Do filter RedisModule_Call() commands initiated by module itself. */
+#define REDISMODULE_CMDFILTER_NOSELF    (1<<0)
+
 /* ------------------------- End of common defines ------------------------ */
 
 #ifndef REDISMODULE_CORE
@@ -341,7 +346,7 @@ void REDISMODULE_API_FUNC(RedisModule_SetDisconnectCallback)(RedisModuleBlockedC
 void REDISMODULE_API_FUNC(RedisModule_SetClusterFlags)(RedisModuleCtx *ctx, uint64_t flags);
 int REDISMODULE_API_FUNC(RedisModule_ExportSharedAPI)(RedisModuleCtx *ctx, const char *apiname, void *func);
 void *REDISMODULE_API_FUNC(RedisModule_GetSharedAPI)(RedisModuleCtx *ctx, const char *apiname);
-RedisModuleCommandFilter *REDISMODULE_API_FUNC(RedisModule_RegisterCommandFilter)(RedisModuleCtx *ctx, RedisModuleCommandFilterFunc cb);
+RedisModuleCommandFilter *REDISMODULE_API_FUNC(RedisModule_RegisterCommandFilter)(RedisModuleCtx *ctx, RedisModuleCommandFilterFunc cb, int flags);
 int REDISMODULE_API_FUNC(RedisModule_UnregisterCommandFilter)(RedisModuleCtx *ctx, RedisModuleCommandFilter *filter);
 int REDISMODULE_API_FUNC(RedisModule_CommandFilterArgsCount)(RedisModuleCommandFilterCtx *fctx);
 const RedisModuleString *REDISMODULE_API_FUNC(RedisModule_CommandFilterArgGet)(RedisModuleCommandFilterCtx *fctx, int pos);
diff --git a/tests/modules/commandfilter.tcl b/tests/modules/commandfilter.tcl
index 8645d8279..1e5c41d2b 100644
--- a/tests/modules/commandfilter.tcl
+++ b/tests/modules/commandfilter.tcl
@@ -1,7 +1,7 @@
 set testmodule [file normalize src/modules/hellofilter.so]
 
 start_server {tags {"modules"}} {
-    r module load $testmodule log-key
+    r module load $testmodule log-key 0
 
     test {Command Filter handles redirected commands} {
         r set mykey @log
@@ -50,18 +50,35 @@ start_server {tags {"modules"}} {
         r lrange log-key 0 -1
     } {}
 
-    r module load $testmodule log-key-2
+    r module load $testmodule log-key 0
 
     test {Command Filter unregister works as expected} {
         # Validate reloading succeeded
+        r del log-key
         r set mykey @log
-        assert_equal "{set mykey @log}" [r lrange log-key-2 0 -1]
+        assert_equal "{set mykey @log}" [r lrange log-key 0 -1]
 
         # Unregister
         r hellofilter.unregister
-        r del log-key-2
+        r del log-key
 
         r set mykey @log
-        r lrange log-key-2 0 -1
+        r lrange log-key 0 -1
     } {}
+
+    r module unload hellofilter
+    r module load $testmodule log-key 1
+
+    test {Command Filter REDISMODULE_CMDFILTER_NOSELF works as expected} {
+        r set mykey @log
+        assert_equal "{set mykey @log}" [r lrange log-key 0 -1]
+
+        r del log-key
+        r hellofilter.ping
+        assert_equal {} [r lrange log-key 0 -1]
+
+        r eval "redis.call('hellofilter.ping')" 0
+        assert_equal {} [r lrange log-key 0 -1]
+    }
+
 }