diff --git a/src/module.c b/src/module.c index a54bd1ad8..4ff865d2c 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 */ @@ -2756,6 +2759,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; @@ -2822,6 +2827,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; } @@ -4857,17 +4863,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); @@ -4908,6 +4924,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 37b7d0d59..e567743a4 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 @@ -340,7 +345,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] + } + }