diff --git a/src/knot/conf/module.c b/src/knot/conf/module.c index 32006dc75..2aee20ae0 100644 --- a/src/knot/conf/module.c +++ b/src/knot/conf/module.c @@ -498,3 +498,8 @@ void conf_reset_modules( (void)rcu_xchg_pointer(query_plan, new_plan); } + +redisContext *redis_conn(knotd_mod_t *mod) +{ + return rdb_connect(mod->config, false, " configuration reader"); +} \ No newline at end of file diff --git a/src/knot/conf/module.h b/src/knot/conf/module.h index 05805e33d..2a45e1edd 100644 --- a/src/knot/conf/module.h +++ b/src/knot/conf/module.h @@ -116,3 +116,7 @@ void conf_reset_modules( list_t *query_modules, struct query_plan **query_plan ); + +#include "knot/common/hiredis.h" + +redisContext *redis_conn(knotd_mod_t *mod); \ No newline at end of file diff --git a/src/knot/modules/geoip/geoip.c b/src/knot/modules/geoip/geoip.c index 6f7c826cf..a35027b9b 100644 --- a/src/knot/modules/geoip/geoip.c +++ b/src/knot/modules/geoip/geoip.c @@ -7,6 +7,7 @@ #include #include +#include "knot/common/hiredis.h" #include "knot/conf/schema.h" #include "knot/include/module.h" #include "knot/modules/geoip/geodb.h" @@ -21,6 +22,7 @@ #include "libzscanner/scanner.h" #define MOD_CONFIG_FILE "\x0B""config-file" +#define MOD_CONFIG_RDB "\x0B""rdb-backend" #define MOD_TTL "\x03""ttl" #define MOD_MODE "\x04""mode" #define MOD_DNSSEC "\x06""dnssec" @@ -49,6 +51,7 @@ static const char* mode_key[] = { const yp_item_t geoip_conf[] = { { MOD_CONFIG_FILE, YP_TSTR, YP_VNONE }, + { MOD_CONFIG_RDB, YP_TBOOL, YP_VBOOL = { false } }, { MOD_TTL, YP_TINT, YP_VINT = { 0, UINT32_MAX, 60, YP_STIME } }, { MOD_MODE, YP_TOPT, YP_VOPT = { modes, MODE_SUBNET} }, { MOD_DNSSEC, YP_TBOOL, YP_VNONE }, @@ -69,12 +72,18 @@ static int load_module(check_ctx_t *ctx); int geoip_conf_check(knotd_conf_check_args_t *args) { - knotd_conf_t conf = knotd_conf_check_item(args, MOD_CONFIG_FILE); - if (conf.count == 0) { - args->err_str = "no configuration file specified"; + knotd_conf_t conf_f = knotd_conf_check_item(args, MOD_CONFIG_FILE); + knotd_conf_t conf_db = knotd_conf_check_item(args, MOD_CONFIG_RDB); + if (conf_f.count < 1 && conf_db.single.boolean == false) { + args->err_str = "no configuration file or database specified"; + return KNOT_EINVAL; + } else if (conf_f.count >= 1 && conf_db.single.boolean == true) { + args->err_str = "configuration file rewrites specified configuration database"; return KNOT_EINVAL; } - conf = knotd_conf_check_item(args, MOD_MODE); + knotd_conf_free(&conf_f); + knotd_conf_free(&conf_db); + knotd_conf_t conf = knotd_conf_check_item(args, MOD_MODE); if (conf.count == 1 && conf.single.option == MODE_GEODB) { if (!geodb_available()) { args->err_str = "geodb mode not available"; @@ -683,6 +692,65 @@ cleanup: return ret; } +#include "knot/conf/module.h" +#include + +static int geo_conf_rdb(check_ctx_t *check, conf_mod_id_t *id, geoip_ctx_t *ctx) +{ + redisContext *db = redis_conn(check->mod); + // TODO GEOIP.LOAD should have parameter for MODE (load only GEO/NET/WEIGHT types, not every type/mode) + redisReply *reply = redisCommand(db, "KNOT.GEOIP.LOAD %b", id->data, id->len - 1); + if (reply == NULL) { + return KNOT_ECONN; + } else if (reply->type != REDIS_REPLY_ARRAY) { + freeReplyObject(reply); + return KNOT_EINVAL; + } else if (reply->elements == 0) { + freeReplyObject(reply); + return KNOT_ENOENT; + } + + for (size_t i = 0; i < reply->elements; ++i) { + redisReply *data = reply->element[i]; + geo_view_t *view = calloc(1, sizeof(geo_view_t)); + + int mode = data->element[0]->integer; // TODO unify + redisReply *geo_data = data->element[1]; + knot_dname_t *owner = (knot_dname_t *)data->element[2]->str; + + switch (mode) { + case 1: + if (ctx->mode != MODE_GEODB) { + continue; + } + view->geodata_len[0] = geo_data->len; + view->geodata[0] = strdup(geo_data->str); + // memcpy(view->geodata[0], geo_data->str, geo_data->len); + break; + case 2: + if (ctx->mode != MODE_SUBNET) { + continue; + } + // view->subnet = (struct sockaddr_storage *)remote; + // view->subnet_prefix = (remote->ss_family == AF_INET) ? 32 : 128; + break; + case 3: + if (ctx->mode != MODE_WEIGHTED) { + continue; + } + break; + default: + return KNOT_ENOTSUP; + } + + view->avail++; + view->count++; + add_view_to_trie(owner, view, ctx); + } + + return KNOT_EOK; +} + static void clear_geo_trie(trie_t *trie) { trie_it_t *it = trie_it_begin(trie); @@ -944,6 +1012,9 @@ static knotd_in_state_t geoip_process(knotd_in_state_t state, knot_pkt_t *pkt, return KNOTD_IN_STATE_NODATA; } } +#include "knot/conf/conf.h" +#include "knot/include/module.h" +#include "knot/nameserver/query_module.h" static int load_module(check_ctx_t *check) { @@ -1015,13 +1086,20 @@ static int load_module(check_ctx_t *check) } // Parse geo configuration file. - int ret = geo_conf_yparse(check, ctx); + int ret = KNOT_EOK; + conf = geo_conf(check, MOD_CONFIG_RDB); + if (conf.single.boolean == false) { + ret = geo_conf_yparse(check, ctx); + } if (ret != KNOT_EOK) { free_geoip_ctx(ctx); return ret; } if (mod != NULL) { + if (conf.single.boolean) { + ret = geo_conf_rdb(check, mod->id, ctx); + } // Prepare geo views for faster search. geo_sort_and_link(ctx); diff --git a/src/redis/Makefile.am b/src/redis/Makefile.am index 3754e73ff..bfe31d9d5 100644 --- a/src/redis/Makefile.am +++ b/src/redis/Makefile.am @@ -13,6 +13,7 @@ knot_la_SOURCES = \ knot.h \ libs.h \ type_diff.h \ + type_geoip.h \ type_rrset.h knot_la_CPPFLAGS = $(AM_CPPFLAGS) $(CFLAG_VISIBILITY) $(gnutls_CFLAGS) diff --git a/src/redis/arg.h b/src/redis/arg.h index 920cc18f7..079285e72 100644 --- a/src/redis/arg.h +++ b/src/redis/arg.h @@ -30,6 +30,11 @@ typedef struct { knot_dname_storage_t buff; } arg_dname_t; +typedef struct { + uint8_t len; + const char *str; +} arg_string_t; + #define ARG_OPT_TXT(out, name, dflt, value) { \ out = dflt; \ if (argc > 1) { \ @@ -166,6 +171,37 @@ typedef struct { out.len = ret; \ } +#define ARG_MODULENAME(arg, out, name) { \ + size_t len; \ + const char *ptr = RedisModule_StringPtrLen(arg, &len); \ + /* TODO test alphanum */ \ + if (len == 0 || len > 255) { \ + return RedisModule_ReplyWithError(ctx, RDB_E("invalid " name)); \ + } \ + out.str = ptr; \ + out.len = len; \ +} + +#define ARG_GEO_TYPEVAL_TXT(arg, out, name) { \ + size_t len; \ + const char *ptr = RedisModule_StringPtrLen(arg, &len); \ + if (len > 4 && strncmp(ptr, "geo:", 4) == 0) { \ + ptr += 4; len -= 4; out.type = GEO; \ + } else if (len > 4 && strncmp(ptr, "net:", 4) == 0) { \ + ptr += 4; len -= 4; out.type = NET; \ + } else if (len > 4 && strncmp(ptr, "weight:", 7) == 0) { \ + ptr += 7; len -= 7; out.type = WEIGHT; \ + } else { \ + return RedisModule_ReplyWithError(ctx, RDB_E("malformed " name)); \ + } \ + /* TODO validate ptr */ \ + if (len < 0 || len >= 256) { \ + return RedisModule_ReplyWithError(ctx, RDB_E("invalid " name)); \ + } \ + out.val = (const uint8_t *)ptr; \ + out.val_size = len; \ +} + #define ARG_FLAG(arg, out, flag) { \ size_t len; \ const char *ptr = RedisModule_StringPtrLen(arg, &len); \ diff --git a/src/redis/internal.h b/src/redis/internal.h index e2e940ab2..bbab19474 100644 --- a/src/redis/internal.h +++ b/src/redis/internal.h @@ -30,15 +30,17 @@ #define return_ok throw(KNOT_EOK, NULL) typedef enum { - EVENT = 1, // Keep synchronized with RDB_EVENT_KEY! - ZONES = 2, - ZONE_META = 3, - ZONE = 4, - RRSET = 5, - UPD_META = 6, - UPD_TMP = 7, - UPD = 8, - DIFF = 9, + EVENT = 1, // Keep synchronized with RDB_EVENT_KEY! + ZONES = 2, + ZONE_META = 3, + ZONE = 4, + RRSET = 5, + UPD_META = 6, + UPD_TMP = 7, + UPD = 8, + DIFF = 9, + GEOIP_MOD = 10, + GEOIP_RRSET = 11, } rdb_type_t; typedef struct { @@ -51,6 +53,19 @@ typedef struct { uint8_t lock[TXN_MAX_COUNT]; } zone_meta_storage_t; +typedef enum { + EMPTY = 0, + GEO, + NET, + WEIGHT +} geoip_meta_type_t; + +typedef struct { + uint8_t type; + uint8_t val_size; + const uint8_t *val; +} geoip_typeval_t; + typedef struct { uint16_t counter; uint16_t lock[TXN_MAX_COUNT]; @@ -75,6 +90,7 @@ typedef RedisModuleKey *rrset_k; typedef RedisModuleKey *diff_k; typedef RedisModuleKey *upd_meta_k; typedef RedisModuleKey *zone_meta_k; +typedef RedisModuleKey *geoip_k; typedef RedisModuleKey *index_k; static void *redismodule_alloc(void *ptr, size_t bytes) @@ -114,6 +130,49 @@ static RedisModuleString *meta_keyname_construct(const uint8_t prefix, RedisModu return RedisModule_CreateString(ctx, buf, wire_ctx_offset(&w)); } +static index_k get_geoip_index(RedisModuleCtx *ctx, const arg_string_t *module, + int rights) +{ + char buf[RDB_PREFIX_LEN + 1 + 1 + 255]; + + wire_ctx_t w = wire_ctx_init((uint8_t *)buf, sizeof(buf)); + wire_ctx_write(&w, RDB_PREFIX, RDB_PREFIX_LEN); + wire_ctx_write_u8(&w, GEOIP_MOD); + wire_ctx_write_u8(&w, module->len); + wire_ctx_write(&w, module->str, module->len); + + RedisModule_Assert(w.error == KNOT_EOK); + + RedisModuleString *keyname = RedisModule_CreateString(ctx, buf, wire_ctx_offset(&w)); + index_k key = RedisModule_OpenKey(ctx, keyname, rights); + RedisModule_FreeString(ctx, keyname); + return key; +} + +static RedisModuleString *geoip_keyname_construct(RedisModuleCtx *ctx, + const arg_string_t *module, + const arg_dname_t *owner, + geoip_typeval_t *geo, + uint16_t rtype) +{ + char buf[RDB_PREFIX_LEN + 1 + 1 + 255 + 1 + KNOT_DNAME_MAXLEN + 1 + 255]; + + wire_ctx_t w = wire_ctx_init((uint8_t *)buf, sizeof(buf)); + wire_ctx_write(&w, RDB_PREFIX, RDB_PREFIX_LEN); + wire_ctx_write_u8(&w, GEOIP_RRSET); + wire_ctx_write_u8(&w, module->len); + wire_ctx_write(&w, module->str, module->len); + wire_ctx_write_u8(&w, owner->len); + wire_ctx_write(&w, owner->data, owner->len); + wire_ctx_write_u8(&w, geo->type); + wire_ctx_write_u8(&w, geo->val_size); + wire_ctx_write(&w, geo->val, geo->val_size); + wire_ctx_write_u16(&w, rtype); + RedisModule_Assert(w.error == KNOT_EOK); + + return RedisModule_CreateString(ctx, buf, wire_ctx_offset(&w)); +} + static RedisModuleString *rrset_keyname_construct(RedisModuleCtx *ctx, const arg_dname_t *origin, const rdb_txn_t *txn, const arg_dname_t *owner, uint16_t rtype) @@ -2475,3 +2534,93 @@ static void upd_load(RedisModuleCtx *ctx, const arg_dname_t *origin, rdb_txn_t * } RedisModule_ReplySetArrayLength(ctx, counter); } + +static void geoip_load(RedisModuleCtx *ctx, const arg_string_t *name, + dump_mode_t mode) +{ + index_k geoip_index = get_geoip_index(ctx, name, REDISMODULE_READ | REDISMODULE_WRITE); + if (geoip_index == NULL) { + RedisModule_ReplyWithEmptyArray(ctx); + return; + } + + RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_ARRAY_LEN); + size_t len = 0; + foreach_in_zset(geoip_index) { + RedisModuleString *el = RedisModule_ZsetRangeCurrentElement(geoip_index, NULL); + + size_t el_len = 0; + uint8_t *wire = (uint8_t *)RedisModule_StringPtrLen(el, &el_len); + + wire_ctx_t w = wire_ctx_init(wire, el_len); + wire_ctx_skip(&w, RDB_PREFIX_LEN + 1); + wire_ctx_skip(&w, wire_ctx_read_u8(&w)); + uint8_t owner_len = wire_ctx_read_u8(&w); + knot_dname_t *owner = w.position; + wire_ctx_skip(&w, owner_len); + uint8_t geo_type = wire_ctx_read_u8(&w); + uint8_t geo_len = wire_ctx_read_u8(&w); + knot_dname_t *geo_val = w.position; + wire_ctx_skip(&w, geo_len); + uint16_t rtype = wire_ctx_read_u16(&w); + RedisModule_Assert(w.error == KNOT_EOK); + + geoip_k rrset_key = RedisModule_OpenKey(ctx, el, REDISMODULE_READ); + geoip_v *rrset = RedisModule_ModuleTypeGetValue(rrset_key); + + RedisModule_ReplyWithArray(ctx, 5); + RedisModule_ReplyWithLongLong(ctx, geo_type); + RedisModule_ReplyWithStringBuffer(ctx, (const char *)geo_val, geo_len); + RedisModule_ReplyWithStringBuffer(ctx, (const char *)owner, owner_len); + RedisModule_ReplyWithLongLong(ctx, rtype); + RedisModule_ReplyWithArray(ctx, rrset->count); + for (int i = 0; i < rrset->size; ++i) { + RedisModule_ReplyWithStringBuffer(ctx, (const char *)rrset->rdata[i].data, rrset->rdata[i].len); + } + + RedisModule_FreeString(ctx, el); + ++len; + } + RedisModule_ReplySetArrayLength(ctx, len); +} + +static void geoip_add(RedisModuleCtx *ctx, const arg_string_t *name, + const arg_dname_t *owner, geoip_typeval_t *type, + uint16_t rtype, uint8_t *rdata, size_t rdata_len) +{ + index_k geoip_index = get_geoip_index(ctx, name, REDISMODULE_READ | REDISMODULE_WRITE); + if (geoip_index == NULL) { + RedisModule_ReplyWithEmptyArray(ctx); + return; + } + + RedisModuleString *geoip_keyname = geoip_keyname_construct(ctx, name, owner, type, rtype); + geoip_k geoip_key = RedisModule_OpenKey(ctx, geoip_keyname, REDISMODULE_READ | REDISMODULE_WRITE); + geoip_v *geoip = NULL; + if (RedisModule_KeyType(geoip_key) == REDISMODULE_KEYTYPE_EMPTY) { + geoip = RedisModule_Alloc(sizeof(geoip_v)); + knot_rdataset_init(geoip); + RedisModule_ModuleTypeSetValue(geoip_key, rdb_geoip_t, geoip); + } else if (RedisModule_ModuleTypeGetType(geoip_key) == rdb_geoip_t) { + geoip = RedisModule_ModuleTypeGetValue(geoip_key); + } else { + RedisModule_CloseKey(geoip_key); + RedisModule_CloseKey(geoip_index); + RedisModule_ReplyWithError(ctx, RDB_EMALF); + return; + } + + uint8_t storage[sizeof(uint16_t) + rdata_len]; + knot_rdata_t *rd = (knot_rdata_t *)storage; + rd->len = rdata_len; + memcpy(rd->data, rdata, rdata_len); + + knot_rdataset_add(geoip, rd, &mm); + + RedisModule_CloseKey(geoip_key); + + RedisModule_ZsetAdd(geoip_index, 0.0, geoip_keyname, NULL); + RedisModule_CloseKey(geoip_index); + + RedisModule_ReplyWithSimpleString(ctx, RDB_RETURN_OK); +} diff --git a/src/redis/knot.c b/src/redis/knot.c index 5f09e93f6..e4de836f8 100644 --- a/src/redis/knot.c +++ b/src/redis/knot.c @@ -16,6 +16,7 @@ #include "redis/libs.h" #include "redis/arg.h" #include "redis/type_diff.h" +#include "redis/type_geoip.h" #include "redis/type_rrset.h" #include "redis/internal.h" @@ -65,6 +66,20 @@ static RedisModuleCommandArg zone_load_txt_info_args[] = { { 0 } }; +static RedisModuleCommandArg geoip_load_txt_info_args[] = { + {"module", REDISMODULE_ARG_TYPE_STRING, -1, NULL, NULL, NULL, REDISMODULE_CMD_ARG_NONE}, + { 0 } +}; + +static RedisModuleCommandArg geoip_store_txt_info_args[] = { + {"module", REDISMODULE_ARG_TYPE_STRING, -1, NULL, NULL, NULL, REDISMODULE_CMD_ARG_NONE}, + {"owner", REDISMODULE_ARG_TYPE_STRING, -1, NULL, NULL, NULL, REDISMODULE_CMD_ARG_NONE}, + {"geo", REDISMODULE_ARG_TYPE_STRING, -1, NULL, NULL, NULL, REDISMODULE_CMD_ARG_NONE}, + {"rtype", REDISMODULE_ARG_TYPE_STRING, -1, NULL, NULL, NULL, REDISMODULE_CMD_ARG_NONE}, + {"data", REDISMODULE_ARG_TYPE_STRING, -1, NULL, NULL, NULL, REDISMODULE_CMD_ARG_NONE}, + { 0 } +}; + static RedisModuleCommandArg zone_list_txt_info_args[] = { {"opt", REDISMODULE_ARG_TYPE_PURE_TOKEN, -1, "--instances", NULL, NULL, REDISMODULE_CMD_ARG_OPTIONAL}, { 0 } @@ -211,6 +226,24 @@ static const RedisModuleCommandInfo upd_load_txt_info = { .args = upd_load_txt_info_args, }; +static const RedisModuleCommandInfo geoip_load_txt_info = { + .version = REDISMODULE_COMMAND_INFO_VERSION, + .summary = "Load geoip configuration", + .complexity = "O(u), where u is the number of records in the retrieved updates", + .since = "7.0.0", + .arity = 2, + .args = geoip_load_txt_info_args, +}; + +static const RedisModuleCommandInfo geoip_store_txt_info = { + .version = REDISMODULE_COMMAND_INFO_VERSION, + .summary = "Load geoip configuration", + .complexity = "O(u), where u is the number of records in the retrieved updates", + .since = "7.0.0", + .arity = 6, + .args = geoip_store_txt_info_args, +}; + static int zone_begin_txt(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { arg_dname_t origin; @@ -844,6 +877,48 @@ static int upd_load_bin(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) return REDISMODULE_OK; } +static int geoip_load_txt(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) +{ + if (argc < 2) { + return RedisModule_WrongArity(ctx); + } + + arg_string_t module_name; + ARG_MODULENAME(argv[1], module_name, "module name"); + + geoip_load(ctx, &module_name, DUMP_TXT); + + return REDISMODULE_OK; +} + +static int geoip_store_txt(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) +{ + if (argc < 6) { + return RedisModule_WrongArity(ctx); + } + + arg_string_t module_name; + ARG_MODULENAME(argv[1], module_name, "module name"); + + arg_dname_t owner; + ARG_DNAME_TXT(argv[2], owner, NULL, "record owner"); + + geoip_typeval_t tv; + ARG_GEO_TYPEVAL_TXT(argv[3], tv, "geoip typeval") + + uint16_t rtype; + ARG_RTYPE_TXT(argv[4], rtype); + + uint8_t *rdata; + size_t rdata_len; + ARG_DATA(argv[5], rdata_len, rdata, "rdata"); + + geoip_add(ctx, &module_name, &owner, &tv, rtype, rdata, rdata_len); + + return REDISMODULE_OK; +} + + #define LOAD_ERROR(ctx, msg) { \ RedisModule_Log(ctx, REDISMODULE_LOGLEVEL_WARNING, RDB_E(msg)); \ RedisModule_ReplyWithError(ctx, RDB_E(msg)); \ @@ -893,6 +968,11 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) LOAD_ERROR(ctx, "failed to load type " RRSET_NAME); } + rdb_geoip_t = RedisModule_CreateDataType(ctx, GEOIP_NAME, GEOIP_ENCODING_VERSION, &geoip_tm); + if (rdb_geoip_t == NULL) { + LOAD_ERROR(ctx, "failed to load type " GEOIP_NAME); + } + rdb_diff_t = RedisModule_CreateDataType(ctx, DIFF_NAME, DIFF_ENCODING_VERSION, &diff_tm); if (rdb_diff_t == NULL) { LOAD_ERROR(ctx, "failed to load type " DIFF_NAME); @@ -914,6 +994,8 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) register_command_txt("KNOT.UPD.ABORT", upd_abort_txt, "write") || register_command_txt("KNOT.UPD.DIFF", upd_diff_txt, "readonly") || register_command_txt("KNOT.UPD.LOAD", upd_load_txt, "readonly") || + register_command_txt("KNOT.GEOIP.LOAD", geoip_load_txt, "readonly") || + register_command_txt("KNOT.GEOIP.STORE", geoip_store_txt, "write fast") || register_command_bin(RDB_CMD_ZONE_EXISTS, zone_exists_bin, "readonly") || register_command_bin(RDB_CMD_ZONE_BEGIN, zone_begin_bin, "write") || register_command_bin(RDB_CMD_ZONE_STORE, zone_store_bin, "write") || @@ -930,6 +1012,7 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) register_command_bin(RDB_CMD_UPD_DIFF, upd_diff_bin, "readonly") || register_command_bin(RDB_CMD_UPD_LOAD, upd_load_bin, "readonly") || register_command_bin("KNOT_BIN.AOF.RRSET", rrset_aof_rewrite, "write") || // Add "internal" with newer Redis. + register_command_bin("KNOT_BIN.AOF.GEOIP", geoip_aof_rewrite, "write") || // Add "internal" with newer Redis. register_command_bin("KNOT_BIN.AOF.DIFF", diff_aof_rewrite, "write")) // Add "internal" with newer Redis. { LOAD_ERROR(ctx, "failed to load commands"); diff --git a/src/redis/type_geoip.h b/src/redis/type_geoip.h new file mode 100644 index 000000000..745ea9d0d --- /dev/null +++ b/src/redis/type_geoip.h @@ -0,0 +1,129 @@ +/* Copyright (C) CZ.NIC, z.s.p.o. and contributors + * SPDX-License-Identifier: LGPL-2.1-or-later + * For more information, see + */ + +#pragma once + +#define GEOIP_ENCODING_VERSION 1 +#define GEOIP_NAME "KnotGeoIP" + +typedef knot_rdataset_t geoip_v; + +static RedisModuleType *rdb_geoip_t; + +static void *rdb_geoip_load(RedisModuleIO *io, int encver) +{ + if (encver != GEOIP_ENCODING_VERSION) { + RedisModule_LogIOError(io, REDISMODULE_LOGLEVEL_WARNING, RDB_ECOMPAT); + return NULL; + } + + geoip_v *rrset = RedisModule_Alloc(sizeof(geoip_v)); + if (rrset == NULL) { + RedisModule_LogIOError(io, REDISMODULE_LOGLEVEL_WARNING, RDB_EALLOC); + return NULL; + } + size_t len = 0; + rrset->count = RedisModule_LoadUnsigned(io); + rrset->rdata = (knot_rdata_t *)RedisModule_LoadStringBuffer(io, &len); + if (len > UINT32_MAX) { + RedisModule_LogIOError(io, REDISMODULE_LOGLEVEL_WARNING, RDB_EMALF); + RedisModule_Free(rrset->rdata); + RedisModule_Free(rrset); + return NULL; + } + rrset->size = len; + + return rrset; +} + +static void rdb_geoip_save(RedisModuleIO *io, void *value) +{ + geoip_v *rrset = (geoip_v *)value; + + RedisModule_SaveUnsigned(io, rrset->count); + RedisModule_SaveStringBuffer(io, (const char *)rrset->rdata, rrset->size); +} + +static size_t rdb_geoip_mem_usage(const void *value) +{ + const geoip_v *rrset = (const geoip_v *)value; + if (value == NULL) { + return 0UL; + } + return sizeof(*rrset) + rrset->size; +} + +static void rdb_geoip_rewrite(RedisModuleIO *io, RedisModuleString *key, void *value) +{ + size_t key_strlen = 0; + const geoip_v *rrset = (const geoip_v *)value; + const uint8_t *key_str = (const uint8_t *)RedisModule_StringPtrLen(key, &key_strlen); + RedisModule_EmitAOF(io, "KNOT_BIN.AOF.GEOPI", "blb", + key_str, key_strlen, + (long long)rrset->count, + rrset->rdata, rrset->size); +} + +static void rdb_geoip_free(void *value) +{ + geoip_v *rrset = (geoip_v *)value; + RedisModule_Free(rrset->rdata); + RedisModule_Free(rrset); +} + +static int geoip_aof_rewrite(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) +{ + // TODO + if (argc != 4) { + return RedisModule_WrongArity(ctx); + } + + geoip_v *rrset = RedisModule_Calloc(1, sizeof(geoip_v)); + if (rrset == NULL) { + return RedisModule_ReplyWithError(ctx, RDB_EALLOC); + } + + RedisModuleKey *rrset_key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE); + + + long long count_val = 0; + int ret = RedisModule_StringToLongLong(argv[2], &count_val); + if (ret != REDISMODULE_OK) { + RedisModule_CloseKey(rrset_key); + return RedisModule_ReplyWithError(ctx, RDB_EMALF); + } else if (count_val < 0 || count_val > UINT16_MAX) { + RedisModule_CloseKey(rrset_key); + return RedisModule_ReplyWithError(ctx, RDB_EMALF); + } + rrset->count = count_val; + + size_t rdataset_strlen = 0; + const char *rdataset_str = RedisModule_StringPtrLen(argv[3], &rdataset_strlen); + if (rdataset_strlen > UINT32_MAX) { + RedisModule_CloseKey(rrset_key); + return RedisModule_ReplyWithError(ctx, RDB_EMALF); + } else if (rdataset_strlen != 0) { + rrset->rdata = RedisModule_Alloc(rdataset_strlen); + rrset->size = rdataset_strlen; + memcpy(rrset->rdata, rdataset_str, rdataset_strlen); + } else { + rrset->rdata = NULL; + rrset->size = 0; + } + + RedisModule_ModuleTypeSetValue(rrset_key, rdb_geoip_t, rrset); + RedisModule_CloseKey(rrset_key); + + return RedisModule_ReplyWithNull(ctx); +} + +RedisModuleTypeMethods geoip_tm = { + .version = REDISMODULE_TYPE_METHOD_VERSION, + .rdb_load = rdb_geoip_load, + .rdb_save = rdb_geoip_save, + .mem_usage = rdb_geoip_mem_usage, + .aof_rewrite = rdb_geoip_rewrite, + .free = rdb_geoip_free +};