redict/tests/modules/datatype.c
Roshan Khatri 6948dacaf6
Module commands to have ACL categories. (#11708)
This allows modules to register commands to existing ACL categories and blocks the creation of [sub]commands, datatypes and registering the configs outside of the OnLoad function.

For allowing modules to register commands to existing ACL categories,
This PR implements a new API int RM_SetCommandACLCategories() which takes a pointer to a RedisModuleCommand and a C string aclflags containing the set of space separated ACL categories.
Example, 'write slow' marks the command as part of the write and slow ACL categories.

The C string aclflags is tokenized by implementing a helper function categoryFlagsFromString(). Theses tokens are matched and the corresponding ACL categories flags are set by a helper function matchAclCategoriesFlags. The helper function categoryFlagsFromString() returns the corresponding categories_flags or returns -1 if some token not processed correctly.

If the module contains commands which are registered to existing ACL categories, the number of [sub]commands are tracked by num_commands_with_acl_categories in struct RedisModule. Further, the allowed command bit-map of the existing users are recomputed from the command_rules list, by implementing a function called ACLRecomputeCommandBitsFromCommandRulesAllUsers() for the existing users to have access to the module commands on runtime.

## Breaking change
This change requires that registering commands and subcommands only occur during a modules "OnLoad" function, in order to allow efficient recompilation of ACL bits. We also chose to block registering configs and types, since we believe it's only valid for those to be created during onLoad. We check for this onload flag in struct RedisModule to check if the call is made from the OnLoad function.

Co-authored-by: Madelyn Olson <madelyneolson@gmail.com>
2023-03-21 10:07:11 -07:00

315 lines
10 KiB
C

/* This module current tests a small subset but should be extended in the future
* for general ModuleDataType coverage.
*/
/* define macros for having usleep */
#define _BSD_SOURCE
#define _DEFAULT_SOURCE
#include <unistd.h>
#include "redismodule.h"
static RedisModuleType *datatype = NULL;
static int load_encver = 0;
/* used to test processing events during slow loading */
static volatile int slow_loading = 0;
static volatile int is_in_slow_loading = 0;
#define DATATYPE_ENC_VER 1
typedef struct {
long long intval;
RedisModuleString *strval;
} DataType;
static void *datatype_load(RedisModuleIO *io, int encver) {
load_encver = encver;
int intval = RedisModule_LoadSigned(io);
if (RedisModule_IsIOError(io)) return NULL;
RedisModuleString *strval = RedisModule_LoadString(io);
if (RedisModule_IsIOError(io)) return NULL;
DataType *dt = (DataType *) RedisModule_Alloc(sizeof(DataType));
dt->intval = intval;
dt->strval = strval;
if (slow_loading) {
RedisModuleCtx *ctx = RedisModule_GetContextFromIO(io);
is_in_slow_loading = 1;
while (slow_loading) {
RedisModule_Yield(ctx, REDISMODULE_YIELD_FLAG_CLIENTS, "Slow module operation");
usleep(1000);
}
is_in_slow_loading = 0;
}
return dt;
}
static void datatype_save(RedisModuleIO *io, void *value) {
DataType *dt = (DataType *) value;
RedisModule_SaveSigned(io, dt->intval);
RedisModule_SaveString(io, dt->strval);
}
static void datatype_free(void *value) {
if (value) {
DataType *dt = (DataType *) value;
if (dt->strval) RedisModule_FreeString(NULL, dt->strval);
RedisModule_Free(dt);
}
}
static void *datatype_copy(RedisModuleString *fromkey, RedisModuleString *tokey, const void *value) {
const DataType *old = value;
/* Answers to ultimate questions cannot be copied! */
if (old->intval == 42)
return NULL;
DataType *new = (DataType *) RedisModule_Alloc(sizeof(DataType));
new->intval = old->intval;
new->strval = RedisModule_CreateStringFromString(NULL, old->strval);
/* Breaking the rules here! We return a copy that also includes traces
* of fromkey/tokey to confirm we get what we expect.
*/
size_t len;
const char *str = RedisModule_StringPtrLen(fromkey, &len);
RedisModule_StringAppendBuffer(NULL, new->strval, "/", 1);
RedisModule_StringAppendBuffer(NULL, new->strval, str, len);
RedisModule_StringAppendBuffer(NULL, new->strval, "/", 1);
str = RedisModule_StringPtrLen(tokey, &len);
RedisModule_StringAppendBuffer(NULL, new->strval, str, len);
return new;
}
static int datatype_set(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
if (argc != 4) {
RedisModule_WrongArity(ctx);
return REDISMODULE_OK;
}
long long intval;
if (RedisModule_StringToLongLong(argv[2], &intval) != REDISMODULE_OK) {
RedisModule_ReplyWithError(ctx, "Invalid integer value");
return REDISMODULE_OK;
}
RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
DataType *dt = RedisModule_Calloc(sizeof(DataType), 1);
dt->intval = intval;
dt->strval = argv[3];
RedisModule_RetainString(ctx, dt->strval);
RedisModule_ModuleTypeSetValue(key, datatype, dt);
RedisModule_CloseKey(key);
RedisModule_ReplyWithSimpleString(ctx, "OK");
return REDISMODULE_OK;
}
static int datatype_restore(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
if (argc != 4) {
RedisModule_WrongArity(ctx);
return REDISMODULE_OK;
}
long long encver;
if (RedisModule_StringToLongLong(argv[3], &encver) != REDISMODULE_OK) {
RedisModule_ReplyWithError(ctx, "Invalid integer value");
return REDISMODULE_OK;
}
DataType *dt = RedisModule_LoadDataTypeFromStringEncver(argv[2], datatype, encver);
if (!dt) {
RedisModule_ReplyWithError(ctx, "Invalid data");
return REDISMODULE_OK;
}
RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
RedisModule_ModuleTypeSetValue(key, datatype, dt);
RedisModule_CloseKey(key);
RedisModule_ReplyWithLongLong(ctx, load_encver);
return REDISMODULE_OK;
}
static int datatype_get(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
if (argc != 2) {
RedisModule_WrongArity(ctx);
return REDISMODULE_OK;
}
RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ);
DataType *dt = RedisModule_ModuleTypeGetValue(key);
RedisModule_CloseKey(key);
if (!dt) {
RedisModule_ReplyWithNullArray(ctx);
} else {
RedisModule_ReplyWithArray(ctx, 2);
RedisModule_ReplyWithLongLong(ctx, dt->intval);
RedisModule_ReplyWithString(ctx, dt->strval);
}
return REDISMODULE_OK;
}
static int datatype_dump(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
if (argc != 2) {
RedisModule_WrongArity(ctx);
return REDISMODULE_OK;
}
RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ);
DataType *dt = RedisModule_ModuleTypeGetValue(key);
RedisModule_CloseKey(key);
RedisModuleString *reply = RedisModule_SaveDataTypeToString(ctx, dt, datatype);
if (!reply) {
RedisModule_ReplyWithError(ctx, "Failed to save");
return REDISMODULE_OK;
}
RedisModule_ReplyWithString(ctx, reply);
RedisModule_FreeString(ctx, reply);
return REDISMODULE_OK;
}
static int datatype_swap(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
if (argc != 3) {
RedisModule_WrongArity(ctx);
return REDISMODULE_OK;
}
RedisModuleKey *a = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
RedisModuleKey *b = RedisModule_OpenKey(ctx, argv[2], REDISMODULE_WRITE);
void *val = RedisModule_ModuleTypeGetValue(a);
int error = (RedisModule_ModuleTypeReplaceValue(b, datatype, val, &val) == REDISMODULE_ERR ||
RedisModule_ModuleTypeReplaceValue(a, datatype, val, NULL) == REDISMODULE_ERR);
if (!error)
RedisModule_ReplyWithSimpleString(ctx, "OK");
else
RedisModule_ReplyWithError(ctx, "ERR failed");
RedisModule_CloseKey(a);
RedisModule_CloseKey(b);
return REDISMODULE_OK;
}
/* used to enable or disable slow loading */
static int datatype_slow_loading(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
if (argc != 2) {
RedisModule_WrongArity(ctx);
return REDISMODULE_OK;
}
long long ll;
if (RedisModule_StringToLongLong(argv[1], &ll) != REDISMODULE_OK) {
RedisModule_ReplyWithError(ctx, "Invalid integer value");
return REDISMODULE_OK;
}
slow_loading = ll;
RedisModule_ReplyWithSimpleString(ctx, "OK");
return REDISMODULE_OK;
}
/* used to test if we reached the slow loading code */
static int datatype_is_in_slow_loading(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
REDISMODULE_NOT_USED(argv);
if (argc != 1) {
RedisModule_WrongArity(ctx);
return REDISMODULE_OK;
}
RedisModule_ReplyWithLongLong(ctx, is_in_slow_loading);
return REDISMODULE_OK;
}
int createDataTypeBlockCheck(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
REDISMODULE_NOT_USED(argv);
REDISMODULE_NOT_USED(argc);
static RedisModuleType *datatype_outside_onload = NULL;
RedisModuleTypeMethods datatype_methods = {
.version = REDISMODULE_TYPE_METHOD_VERSION,
.rdb_load = datatype_load,
.rdb_save = datatype_save,
.free = datatype_free,
.copy = datatype_copy
};
datatype_outside_onload = RedisModule_CreateDataType(ctx, "test_dt_outside_onload", 1, &datatype_methods);
/* This validates that it's not possible to create datatype outside OnLoad,
* thus returns an error if it succeeds. */
if (datatype_outside_onload == NULL) {
RedisModule_ReplyWithSimpleString(ctx, "OK");
} else {
RedisModule_ReplyWithError(ctx, "UNEXPECTEDOK");
}
return REDISMODULE_OK;
}
int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
REDISMODULE_NOT_USED(argv);
REDISMODULE_NOT_USED(argc);
if (RedisModule_Init(ctx,"datatype",DATATYPE_ENC_VER,REDISMODULE_APIVER_1) == REDISMODULE_ERR)
return REDISMODULE_ERR;
/* Creates a command which creates a datatype outside OnLoad() function. */
if (RedisModule_CreateCommand(ctx,"block.create.datatype.outside.onload", createDataTypeBlockCheck, "write", 0, 0, 0) == REDISMODULE_ERR)
return REDISMODULE_ERR;
RedisModule_SetModuleOptions(ctx, REDISMODULE_OPTIONS_HANDLE_IO_ERRORS);
RedisModuleTypeMethods datatype_methods = {
.version = REDISMODULE_TYPE_METHOD_VERSION,
.rdb_load = datatype_load,
.rdb_save = datatype_save,
.free = datatype_free,
.copy = datatype_copy
};
datatype = RedisModule_CreateDataType(ctx, "test___dt", 1, &datatype_methods);
if (datatype == NULL)
return REDISMODULE_ERR;
if (RedisModule_CreateCommand(ctx,"datatype.set", datatype_set,
"write deny-oom", 1, 1, 1) == REDISMODULE_ERR)
return REDISMODULE_ERR;
if (RedisModule_CreateCommand(ctx,"datatype.get", datatype_get,"",1,1,1) == REDISMODULE_ERR)
return REDISMODULE_ERR;
if (RedisModule_CreateCommand(ctx,"datatype.restore", datatype_restore,
"write deny-oom", 1, 1, 1) == REDISMODULE_ERR)
return REDISMODULE_ERR;
if (RedisModule_CreateCommand(ctx,"datatype.dump", datatype_dump,"",1,1,1) == REDISMODULE_ERR)
return REDISMODULE_ERR;
if (RedisModule_CreateCommand(ctx, "datatype.swap", datatype_swap,
"write", 1, 1, 1) == REDISMODULE_ERR)
return REDISMODULE_ERR;
if (RedisModule_CreateCommand(ctx, "datatype.slow_loading", datatype_slow_loading,
"allow-loading", 0, 0, 0) == REDISMODULE_ERR)
return REDISMODULE_ERR;
if (RedisModule_CreateCommand(ctx, "datatype.is_in_slow_loading", datatype_is_in_slow_loading,
"allow-loading", 0, 0, 0) == REDISMODULE_ERR)
return REDISMODULE_ERR;
return REDISMODULE_OK;
}