futriix/tests/modules/commandfilter.c
Dmitry Polyakovsky f0e1edc273
Updated modules examples and tests to use valkey vs redis (#349)
Scope of the changes:
- updated example modules to reference Valkey vs Redis in variable names
- updated tests to use valkeymodule.h
- updated vars in tests/modules to use valkey vs redis in variable names

Summary of the testing:
- ran make for all modules, loaded them into valkey-server and tested
commands
- ran make for test/modules
- ran make test for the entire codebase

---------

Signed-off-by: Dmitry Polyakovsky <dmitry.polyakovky@oracle.com>
Co-authored-by: Dmitry Polyakovsky <dmitry.polyakovky@oracle.com>
2024-04-23 17:55:44 +02:00

252 lines
8.8 KiB
C

#include "valkeymodule.h"
#include <string.h>
#include <strings.h>
static ValkeyModuleString *log_key_name;
static const char log_command_name[] = "commandfilter.log";
static const char ping_command_name[] = "commandfilter.ping";
static const char retained_command_name[] = "commandfilter.retained";
static const char unregister_command_name[] = "commandfilter.unregister";
static const char unfiltered_clientid_name[] = "unfilter_clientid";
static int in_log_command = 0;
unsigned long long unfiltered_clientid = 0;
static ValkeyModuleCommandFilter *filter, *filter1;
static ValkeyModuleString *retained;
int CommandFilter_UnregisterCommand(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
{
(void) argc;
(void) argv;
ValkeyModule_ReplyWithLongLong(ctx,
ValkeyModule_UnregisterCommandFilter(ctx, filter));
return VALKEYMODULE_OK;
}
int CommandFilter_PingCommand(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
{
(void) argc;
(void) argv;
ValkeyModuleCallReply *reply = ValkeyModule_Call(ctx, "ping", "c", "@log");
if (reply) {
ValkeyModule_ReplyWithCallReply(ctx, reply);
ValkeyModule_FreeCallReply(reply);
} else {
ValkeyModule_ReplyWithSimpleString(ctx, "Unknown command or invalid arguments");
}
return VALKEYMODULE_OK;
}
int CommandFilter_Retained(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
{
(void) argc;
(void) argv;
if (retained) {
ValkeyModule_ReplyWithString(ctx, retained);
} else {
ValkeyModule_ReplyWithNull(ctx);
}
return VALKEYMODULE_OK;
}
int CommandFilter_LogCommand(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
{
ValkeyModuleString *s = ValkeyModule_CreateString(ctx, "", 0);
int i;
for (i = 1; i < argc; i++) {
size_t arglen;
const char *arg = ValkeyModule_StringPtrLen(argv[i], &arglen);
if (i > 1) ValkeyModule_StringAppendBuffer(ctx, s, " ", 1);
ValkeyModule_StringAppendBuffer(ctx, s, arg, arglen);
}
ValkeyModuleKey *log = ValkeyModule_OpenKey(ctx, log_key_name, VALKEYMODULE_WRITE|VALKEYMODULE_READ);
ValkeyModule_ListPush(log, VALKEYMODULE_LIST_HEAD, s);
ValkeyModule_CloseKey(log);
ValkeyModule_FreeString(ctx, s);
in_log_command = 1;
size_t cmdlen;
const char *cmdname = ValkeyModule_StringPtrLen(argv[1], &cmdlen);
ValkeyModuleCallReply *reply = ValkeyModule_Call(ctx, cmdname, "v", &argv[2], argc - 2);
if (reply) {
ValkeyModule_ReplyWithCallReply(ctx, reply);
ValkeyModule_FreeCallReply(reply);
} else {
ValkeyModule_ReplyWithSimpleString(ctx, "Unknown command or invalid arguments");
}
in_log_command = 0;
return VALKEYMODULE_OK;
}
int CommandFilter_UnfilteredClientId(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
{
if (argc < 2)
return ValkeyModule_WrongArity(ctx);
long long id;
if (ValkeyModule_StringToLongLong(argv[1], &id) != VALKEYMODULE_OK) {
ValkeyModule_ReplyWithError(ctx, "invalid client id");
return VALKEYMODULE_OK;
}
if (id < 0) {
ValkeyModule_ReplyWithError(ctx, "invalid client id");
return VALKEYMODULE_OK;
}
unfiltered_clientid = id;
ValkeyModule_ReplyWithSimpleString(ctx, "OK");
return VALKEYMODULE_OK;
}
/* Filter to protect against Bug #11894 reappearing
*
* ensures that the filter is only run the first time through, and not on reprocessing
*/
void CommandFilter_BlmoveSwap(ValkeyModuleCommandFilterCtx *filter)
{
if (ValkeyModule_CommandFilterArgsCount(filter) != 6)
return;
ValkeyModuleString *arg = ValkeyModule_CommandFilterArgGet(filter, 0);
size_t arg_len;
const char *arg_str = ValkeyModule_StringPtrLen(arg, &arg_len);
if (arg_len != 6 || strncmp(arg_str, "blmove", 6))
return;
/*
* Swapping directional args (right/left) from source and destination.
* need to hold here, can't push into the ArgReplace func, as it will cause other to freed -> use after free
*/
ValkeyModuleString *dir1 = ValkeyModule_HoldString(NULL, ValkeyModule_CommandFilterArgGet(filter, 3));
ValkeyModuleString *dir2 = ValkeyModule_HoldString(NULL, ValkeyModule_CommandFilterArgGet(filter, 4));
ValkeyModule_CommandFilterArgReplace(filter, 3, dir2);
ValkeyModule_CommandFilterArgReplace(filter, 4, dir1);
}
void CommandFilter_CommandFilter(ValkeyModuleCommandFilterCtx *filter)
{
unsigned long long id = ValkeyModule_CommandFilterGetClientId(filter);
if (id == unfiltered_clientid) return;
if (in_log_command) return; /* don't process our own RM_Call() from CommandFilter_LogCommand() */
/* Fun manipulations:
* - Remove @delme
* - Replace @replaceme
* - Append @insertbefore or @insertafter
* - Prefix with Log command if @log encountered
*/
int log = 0;
int pos = 0;
while (pos < ValkeyModule_CommandFilterArgsCount(filter)) {
const ValkeyModuleString *arg = ValkeyModule_CommandFilterArgGet(filter, pos);
size_t arg_len;
const char *arg_str = ValkeyModule_StringPtrLen(arg, &arg_len);
if (arg_len == 6 && !memcmp(arg_str, "@delme", 6)) {
ValkeyModule_CommandFilterArgDelete(filter, pos);
continue;
}
if (arg_len == 10 && !memcmp(arg_str, "@replaceme", 10)) {
ValkeyModule_CommandFilterArgReplace(filter, pos,
ValkeyModule_CreateString(NULL, "--replaced--", 12));
} else if (arg_len == 13 && !memcmp(arg_str, "@insertbefore", 13)) {
ValkeyModule_CommandFilterArgInsert(filter, pos,
ValkeyModule_CreateString(NULL, "--inserted-before--", 19));
pos++;
} else if (arg_len == 12 && !memcmp(arg_str, "@insertafter", 12)) {
ValkeyModule_CommandFilterArgInsert(filter, pos + 1,
ValkeyModule_CreateString(NULL, "--inserted-after--", 18));
pos++;
} else if (arg_len == 7 && !memcmp(arg_str, "@retain", 7)) {
if (retained) ValkeyModule_FreeString(NULL, retained);
retained = ValkeyModule_CommandFilterArgGet(filter, pos + 1);
ValkeyModule_RetainString(NULL, retained);
pos++;
} else if (arg_len == 4 && !memcmp(arg_str, "@log", 4)) {
log = 1;
}
pos++;
}
if (log) ValkeyModule_CommandFilterArgInsert(filter, 0,
ValkeyModule_CreateString(NULL, log_command_name, sizeof(log_command_name)-1));
}
int ValkeyModule_OnLoad(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
if (ValkeyModule_Init(ctx,"commandfilter",1,VALKEYMODULE_APIVER_1)
== VALKEYMODULE_ERR) return VALKEYMODULE_ERR;
if (argc != 2 && argc != 3) {
ValkeyModule_Log(ctx, "warning", "Log key name not specified");
return VALKEYMODULE_ERR;
}
long long noself = 0;
log_key_name = ValkeyModule_CreateStringFromString(ctx, argv[0]);
ValkeyModule_StringToLongLong(argv[1], &noself);
retained = NULL;
if (ValkeyModule_CreateCommand(ctx,log_command_name,
CommandFilter_LogCommand,"write deny-oom",1,1,1) == VALKEYMODULE_ERR)
return VALKEYMODULE_ERR;
if (ValkeyModule_CreateCommand(ctx,ping_command_name,
CommandFilter_PingCommand,"deny-oom",1,1,1) == VALKEYMODULE_ERR)
return VALKEYMODULE_ERR;
if (ValkeyModule_CreateCommand(ctx,retained_command_name,
CommandFilter_Retained,"readonly",1,1,1) == VALKEYMODULE_ERR)
return VALKEYMODULE_ERR;
if (ValkeyModule_CreateCommand(ctx,unregister_command_name,
CommandFilter_UnregisterCommand,"write deny-oom",1,1,1) == VALKEYMODULE_ERR)
return VALKEYMODULE_ERR;
if (ValkeyModule_CreateCommand(ctx, unfiltered_clientid_name,
CommandFilter_UnfilteredClientId, "admin", 1,1,1) == VALKEYMODULE_ERR)
return VALKEYMODULE_ERR;
if ((filter = ValkeyModule_RegisterCommandFilter(ctx, CommandFilter_CommandFilter,
noself ? VALKEYMODULE_CMDFILTER_NOSELF : 0))
== NULL) return VALKEYMODULE_ERR;
if ((filter1 = ValkeyModule_RegisterCommandFilter(ctx, CommandFilter_BlmoveSwap, 0)) == NULL)
return VALKEYMODULE_ERR;
if (argc == 3) {
const char *ptr = ValkeyModule_StringPtrLen(argv[2], NULL);
if (!strcasecmp(ptr, "noload")) {
/* This is a hint that we return ERR at the last moment of OnLoad. */
ValkeyModule_FreeString(ctx, log_key_name);
if (retained) ValkeyModule_FreeString(NULL, retained);
return VALKEYMODULE_ERR;
}
}
return VALKEYMODULE_OK;
}
int ValkeyModule_OnUnload(ValkeyModuleCtx *ctx) {
ValkeyModule_FreeString(ctx, log_key_name);
if (retained) ValkeyModule_FreeString(NULL, retained);
return VALKEYMODULE_OK;
}