1 #define REDISMODULE_EXPERIMENTAL_API
2 #include "redismodule.h"
3 
4 #include <string.h>
5 
6 static RedisModuleString *log_key_name;
7 
8 static const char log_command_name[] = "commandfilter.log";
9 static const char ping_command_name[] = "commandfilter.ping";
10 static const char unregister_command_name[] = "commandfilter.unregister";
11 static int in_log_command = 0;
12 
13 static RedisModuleCommandFilter *filter = NULL;
14 
CommandFilter_UnregisterCommand(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)15 int CommandFilter_UnregisterCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
16 {
17     (void) argc;
18     (void) argv;
19 
20     RedisModule_ReplyWithLongLong(ctx,
21             RedisModule_UnregisterCommandFilter(ctx, filter));
22 
23     return REDISMODULE_OK;
24 }
25 
CommandFilter_PingCommand(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)26 int CommandFilter_PingCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
27 {
28     (void) argc;
29     (void) argv;
30 
31     RedisModuleCallReply *reply = RedisModule_Call(ctx, "ping", "c", "@log");
32     if (reply) {
33         RedisModule_ReplyWithCallReply(ctx, reply);
34         RedisModule_FreeCallReply(reply);
35     } else {
36         RedisModule_ReplyWithSimpleString(ctx, "Unknown command or invalid arguments");
37     }
38 
39     return REDISMODULE_OK;
40 }
41 
CommandFilter_LogCommand(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)42 int CommandFilter_LogCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
43 {
44     RedisModuleString *s = RedisModule_CreateString(ctx, "", 0);
45 
46     int i;
47     for (i = 1; i < argc; i++) {
48         size_t arglen;
49         const char *arg = RedisModule_StringPtrLen(argv[i], &arglen);
50 
51         if (i > 1) RedisModule_StringAppendBuffer(ctx, s, " ", 1);
52         RedisModule_StringAppendBuffer(ctx, s, arg, arglen);
53     }
54 
55     RedisModuleKey *log = RedisModule_OpenKey(ctx, log_key_name, REDISMODULE_WRITE|REDISMODULE_READ);
56     RedisModule_ListPush(log, REDISMODULE_LIST_HEAD, s);
57     RedisModule_CloseKey(log);
58     RedisModule_FreeString(ctx, s);
59 
60     in_log_command = 1;
61 
62     size_t cmdlen;
63     const char *cmdname = RedisModule_StringPtrLen(argv[1], &cmdlen);
64     RedisModuleCallReply *reply = RedisModule_Call(ctx, cmdname, "v", &argv[2], argc - 2);
65     if (reply) {
66         RedisModule_ReplyWithCallReply(ctx, reply);
67         RedisModule_FreeCallReply(reply);
68     } else {
69         RedisModule_ReplyWithSimpleString(ctx, "Unknown command or invalid arguments");
70     }
71 
72     in_log_command = 0;
73 
74     return REDISMODULE_OK;
75 }
76 
CommandFilter_CommandFilter(RedisModuleCommandFilterCtx * filter)77 void CommandFilter_CommandFilter(RedisModuleCommandFilterCtx *filter)
78 {
79     if (in_log_command) return;  /* don't process our own RM_Call() from CommandFilter_LogCommand() */
80 
81     /* Fun manipulations:
82      * - Remove @delme
83      * - Replace @replaceme
84      * - Append @insertbefore or @insertafter
85      * - Prefix with Log command if @log encounterd
86      */
87     int log = 0;
88     int pos = 0;
89     while (pos < RedisModule_CommandFilterArgsCount(filter)) {
90         const RedisModuleString *arg = RedisModule_CommandFilterArgGet(filter, pos);
91         size_t arg_len;
92         const char *arg_str = RedisModule_StringPtrLen(arg, &arg_len);
93 
94         if (arg_len == 6 && !memcmp(arg_str, "@delme", 6)) {
95             RedisModule_CommandFilterArgDelete(filter, pos);
96             continue;
97         }
98         if (arg_len == 10 && !memcmp(arg_str, "@replaceme", 10)) {
99             RedisModule_CommandFilterArgReplace(filter, pos,
100                     RedisModule_CreateString(NULL, "--replaced--", 12));
101         } else if (arg_len == 13 && !memcmp(arg_str, "@insertbefore", 13)) {
102             RedisModule_CommandFilterArgInsert(filter, pos,
103                     RedisModule_CreateString(NULL, "--inserted-before--", 19));
104             pos++;
105         } else if (arg_len == 12 && !memcmp(arg_str, "@insertafter", 12)) {
106             RedisModule_CommandFilterArgInsert(filter, pos + 1,
107                     RedisModule_CreateString(NULL, "--inserted-after--", 18));
108             pos++;
109         } else if (arg_len == 4 && !memcmp(arg_str, "@log", 4)) {
110             log = 1;
111         }
112         pos++;
113     }
114 
115     if (log) RedisModule_CommandFilterArgInsert(filter, 0,
116             RedisModule_CreateString(NULL, log_command_name, sizeof(log_command_name)-1));
117 }
118 
RedisModule_OnLoad(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)119 int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
120     if (RedisModule_Init(ctx,"commandfilter",1,REDISMODULE_APIVER_1)
121             == REDISMODULE_ERR) return REDISMODULE_ERR;
122 
123     if (argc != 2) {
124         RedisModule_Log(ctx, "warning", "Log key name not specified");
125         return REDISMODULE_ERR;
126     }
127 
128     long long noself = 0;
129     log_key_name = RedisModule_CreateStringFromString(ctx, argv[0]);
130     RedisModule_StringToLongLong(argv[1], &noself);
131 
132     if (RedisModule_CreateCommand(ctx,log_command_name,
133                 CommandFilter_LogCommand,"write deny-oom",1,1,1) == REDISMODULE_ERR)
134             return REDISMODULE_ERR;
135 
136     if (RedisModule_CreateCommand(ctx,ping_command_name,
137                 CommandFilter_PingCommand,"deny-oom",1,1,1) == REDISMODULE_ERR)
138             return REDISMODULE_ERR;
139 
140     if (RedisModule_CreateCommand(ctx,unregister_command_name,
141                 CommandFilter_UnregisterCommand,"write deny-oom",1,1,1) == REDISMODULE_ERR)
142             return REDISMODULE_ERR;
143 
144     if ((filter = RedisModule_RegisterCommandFilter(ctx, CommandFilter_CommandFilter,
145                     noself ? REDISMODULE_CMDFILTER_NOSELF : 0))
146             == NULL) return REDISMODULE_ERR;
147 
148     return REDISMODULE_OK;
149 }
150