Add command filtering argument handling API.

3 files changed, 132 insertions, 13 deletions

diff --git a/src/module.c b/src/module.c
index 1780342ed..741c546b7 100644
--- a/src/module.c
+++ b/src/module.c
@@ -291,6 +291,10 @@ typedef struct RedisModuleCommandFilter {
 /* Registered filters */
 static list *moduleCommandFilters;
 
+typedef struct RedisModuleCommandFilterCtx {
+    RedisModuleString **argv;
+    int argc;
+} RedisModuleCommandFilterCtx;
 
 /* --------------------------------------------------------------------------
  * Prototypes
@@ -4842,6 +4846,78 @@ void moduleCallCommandFilters(client *c) {
     c->argc = cmd.argc;
 }
 
+/* Return the number of arguments a filtered command has.  The number of
+ * arguments include the command itself.
+ */
+int RM_CommandFilterArgsCount(RedisModuleCommandFilterCtx *filter)
+{
+    return filter->argc;
+}
+
+/* Return the specified command argument.  The first argument (position 0) is
+ * the command itself, and the rest are user-provided args.
+ */
+const RedisModuleString *RM_CommandFilterArgGet(RedisModuleCommandFilterCtx *filter, int pos)
+{
+    if (pos < 0 || pos >= filter->argc) return NULL;
+    return filter->argv[pos];
+}
+
+/* Modify the filtered command by inserting a new argument at the specified
+ * position.  The specified RedisModuleString argument may be used by Redis
+ * after the filter context is destroyed, so it must not be auto-memory
+ * allocated, freed or used elsewhere.
+ */
+
+int RM_CommandFilterArgInsert(RedisModuleCommandFilterCtx *filter, int pos, RedisModuleString *arg)
+{
+    int i;
+
+    if (pos < 0 || pos > filter->argc) return REDISMODULE_ERR;
+
+    filter->argv = zrealloc(filter->argv, (filter->argc+1)*sizeof(RedisModuleString *));
+    for (i = filter->argc; i > pos; i--) {
+        filter->argv[i] = filter->argv[i-1];
+    }
+    filter->argv[pos] = arg;
+    filter->argc++;
+
+    return REDISMODULE_OK;
+}
+
+/* Modify the filtered command by replacing an existing argument with a new one.
+ * The specified RedisModuleString argument may be used by Redis after the
+ * filter context is destroyed, so it must not be auto-memory allocated, freed
+ * or used elsewhere.
+ */
+
+int RM_CommandFilterArgReplace(RedisModuleCommandFilterCtx *filter, int pos, RedisModuleString *arg)
+{
+    if (pos < 0 || pos >= filter->argc) return REDISMODULE_ERR;
+    
+    decrRefCount(filter->argv[pos]);
+    filter->argv[pos] = arg;
+
+    return REDISMODULE_OK;
+}
+
+/* Modify the filtered command by deleting an argument at the specified
+ * position.
+ */
+int RM_CommandFilterArgDelete(RedisModuleCommandFilterCtx *filter, int pos)
+{
+    int i;
+    if (pos < 0 || pos >= filter->argc) return REDISMODULE_ERR;
+
+    decrRefCount(filter->argv[pos]);
+    for (i = pos; i < filter->argc-1; i++) {
+        filter->argv[i] = filter->argv[i+1];
+    }
+    filter->argc--;
+    
+    return REDISMODULE_OK;
+}
+
 /* --------------------------------------------------------------------------
  * Modules API internals
  * -------------------------------------------------------------------------- */
@@ -5295,4 +5371,9 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(ExportSharedAPI);
     REGISTER_API(GetSharedAPI);
     REGISTER_API(RegisterCommandFilter);
+    REGISTER_API(CommandFilterArgsCount);
+    REGISTER_API(CommandFilterArgGet);
+    REGISTER_API(CommandFilterArgInsert);
+    REGISTER_API(CommandFilterArgReplace);
+    REGISTER_API(CommandFilterArgDelete);
 }
diff --git a/src/modules/hellofilter.c b/src/modules/hellofilter.c
index c9e33158f..84eb02c30 100644
--- a/src/modules/hellofilter.c
+++ b/src/modules/hellofilter.c
@@ -1,6 +1,8 @@
 #define REDISMODULE_EXPERIMENTAL_API
 #include "../redismodule.h"
 
+#include <string.h>
+
 static RedisModuleString *log_key_name;
 
 static const char log_command_name[] = "hellofilter.log";
@@ -35,16 +37,46 @@ int HelloFilter_LogCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int ar
     return REDISMODULE_OK;
 }
 
-void HelloFilter_CommandFilter(RedisModuleCtx *ctx, RedisModuleFilteredCommand *cmd)
+void HelloFilter_CommandFilter(RedisModuleCtx *ctx, RedisModuleCommandFilterCtx *filter)
 {
-    cmd->argv = RedisModule_Realloc(cmd->argv, (cmd->argc+1)*sizeof(RedisModuleString *));
-    int i;
+    (void) ctx;
 
-    for (i = cmd->argc; i > 0; i--) {
-        cmd->argv[i] = cmd->argv[i-1];
+    /* Fun manipulations:
+     * - Remove @delme
+     * - Replace @replaceme
+     * - Append @insertbefore or @insertafter
+     * - Prefix with Log command if @log encounterd
+     */
+    int log = 0;
+    int pos = 0;
+    while (pos < RedisModule_CommandFilterArgsCount(filter)) {
+        const RedisModuleString *arg = RedisModule_CommandFilterArgGet(filter, pos);
+        size_t arg_len;
+        const char *arg_str = RedisModule_StringPtrLen(arg, &arg_len);
+
+        if (arg_len == 6 && !memcmp(arg_str, "@delme", 6)) {
+            RedisModule_CommandFilterArgDelete(filter, pos);
+            continue;
+        } 
+        if (arg_len == 10 && !memcmp(arg_str, "@replaceme", 10)) {
+            RedisModule_CommandFilterArgReplace(filter, pos,
+                    RedisModule_CreateString(NULL, "--replaced--", 12));
+        } else if (arg_len == 13 && !memcmp(arg_str, "@insertbefore", 13)) {
+            RedisModule_CommandFilterArgInsert(filter, pos,
+                    RedisModule_CreateString(NULL, "--inserted-before--", 19));
+            pos++;
+        } else if (arg_len == 12 && !memcmp(arg_str, "@insertafter", 12)) {
+            RedisModule_CommandFilterArgInsert(filter, pos + 1,
+                    RedisModule_CreateString(NULL, "--inserted-after--", 18));
+            pos++;
+        } else if (arg_len == 4 && !memcmp(arg_str, "@log", 4)) {
+            log = 1;
+        }
+        pos++;
     }
-    cmd->argv[0] = RedisModule_CreateString(ctx, log_command_name, sizeof(log_command_name)-1);
-    cmd->argc++;
+
+    if (log) RedisModule_CommandFilterArgInsert(filter, 0,
+            RedisModule_CreateString(NULL, log_command_name, sizeof(log_command_name)-1));
 }
 
 int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
diff --git a/src/redismodule.h b/src/redismodule.h
index 54ce99d96..426a6df69 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -150,6 +150,7 @@ typedef struct RedisModuleBlockedClient RedisModuleBlockedClient;
 typedef struct RedisModuleClusterInfo RedisModuleClusterInfo;
 typedef struct RedisModuleDict RedisModuleDict;
 typedef struct RedisModuleDictIter RedisModuleDictIter;
+typedef struct RedisModuleCommandFilterCtx RedisModuleCommandFilterCtx;
 
 typedef int (*RedisModuleCmdFunc)(RedisModuleCtx *ctx, RedisModuleString **argv, int argc);
 typedef void (*RedisModuleDisconnectFunc)(RedisModuleCtx *ctx, RedisModuleBlockedClient *bc);
@@ -162,12 +163,7 @@ typedef void (*RedisModuleTypeDigestFunc)(RedisModuleDigest *digest, void *value
 typedef void (*RedisModuleTypeFreeFunc)(void *value);
 typedef void (*RedisModuleClusterMessageReceiver)(RedisModuleCtx *ctx, const char *sender_id, uint8_t type, const unsigned char *payload, uint32_t len);
 typedef void (*RedisModuleTimerProc)(RedisModuleCtx *ctx, void *data);
-
-typedef struct RedisModuleFilteredCommand {
-    RedisModuleString **argv;
-    int argc;
-} RedisModuleFilteredCommand;
-typedef void (*RedisModuleCommandFilterFunc) (RedisModuleCtx *ctx, RedisModuleFilteredCommand *cmd);
+typedef void (*RedisModuleCommandFilterFunc) (RedisModuleCtx *ctx, RedisModuleCommandFilterCtx *filter);
 
 #define REDISMODULE_TYPE_METHOD_VERSION 1
 typedef struct RedisModuleTypeMethods {
@@ -344,6 +340,11 @@ void REDISMODULE_API_FUNC(RedisModule_SetClusterFlags)(RedisModuleCtx *ctx, uint
 int REDISMODULE_API_FUNC(RedisModule_ExportSharedAPI)(RedisModuleCtx *ctx, const char *apiname, void *func);
 void *REDISMODULE_API_FUNC(RedisModule_GetSharedAPI)(RedisModuleCtx *ctx, const char *apiname);
 int REDISMODULE_API_FUNC(RedisModule_RegisterCommandFilter)(RedisModuleCtx *ctx, RedisModuleCommandFilterFunc cb);
+int REDISMODULE_API_FUNC(RedisModule_CommandFilterArgsCount)(RedisModuleCommandFilterCtx *filter);
+const RedisModuleString *REDISMODULE_API_FUNC(RedisModule_CommandFilterArgGet)(RedisModuleCommandFilterCtx *filter, int pos);
+int REDISMODULE_API_FUNC(RedisModule_CommandFilterArgInsert)(RedisModuleCommandFilterCtx *filter, int pos, RedisModuleString *arg);
+int REDISMODULE_API_FUNC(RedisModule_CommandFilterArgReplace)(RedisModuleCommandFilterCtx *filter, int pos, RedisModuleString *arg);
+int REDISMODULE_API_FUNC(RedisModule_CommandFilterArgDelete)(RedisModuleCommandFilterCtx *filter, int pos);
 #endif
 
 /* This is included inline inside each Redis module. */
@@ -507,6 +508,11 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(ExportSharedAPI);
     REDISMODULE_GET_API(GetSharedAPI);
     REDISMODULE_GET_API(RegisterCommandFilter);
+    REDISMODULE_GET_API(CommandFilterArgsCount);
+    REDISMODULE_GET_API(CommandFilterArgGet);
+    REDISMODULE_GET_API(CommandFilterArgInsert);
+    REDISMODULE_GET_API(CommandFilterArgReplace);
+    REDISMODULE_GET_API(CommandFilterArgDelete);
 #endif
 
     if (RedisModule_IsModuleNameBusy && RedisModule_IsModuleNameBusy(name)) return REDISMODULE_ERR;