From 05457bf4200250d8262bd1de87b0a33545cf70a4 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Tue, 19 Oct 2021 17:49:34 +0100 Subject: [PATCH] [Feature] Lua_cdb: Add cdb building interface --- src/lua/lua_cdb.c | 170 +++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 154 insertions(+), 16 deletions(-) diff --git a/src/lua/lua_cdb.c b/src/lua/lua_cdb.c index 8f13f6ccf..91f770cb3 100644 --- a/src/lua/lua_cdb.c +++ b/src/lua/lua_cdb.c @@ -23,34 +23,65 @@ LUA_FUNCTION_DEF (cdb, lookup); LUA_FUNCTION_DEF (cdb, get_name); LUA_FUNCTION_DEF (cdb, destroy); +LUA_FUNCTION_DEF (cdb, build); +LUA_FUNCTION_DEF (cdb_builder, add); +LUA_FUNCTION_DEF (cdb_builder, finalize); +LUA_FUNCTION_DEF (cdb_builder, dtor); + static const struct luaL_reg cdblib_m[] = { LUA_INTERFACE_DEF (cdb, lookup), + {"find", lua_cdb_lookup}, LUA_INTERFACE_DEF (cdb, get_name), {"__tostring", rspamd_lua_class_tostring}, {"__gc", lua_cdb_destroy}, {NULL, NULL} }; + +static const struct luaL_reg cdbbuilderlib_m[] = { + LUA_INTERFACE_DEF (cdb_builder, add), + LUA_INTERFACE_DEF (cdb_builder, finalize), + {"__tostring", rspamd_lua_class_tostring}, + {"__gc", lua_cdb_builder_dtor}, + {NULL, NULL} +}; + static const struct luaL_reg cdblib_f[] = { LUA_INTERFACE_DEF (cdb, create), + {"open", lua_cdb_create}, + {"build", lua_cdb_build}, {NULL, NULL} }; static struct cdb * -lua_check_cdb (lua_State * L) +lua_check_cdb (lua_State * L, int pos) { - void *ud = rspamd_lua_check_udata (L, 1, "rspamd{cdb}"); + void *ud = rspamd_lua_check_udata (L, pos, "rspamd{cdb}"); - luaL_argcheck (L, ud != NULL, 1, "'cdb' expected"); + luaL_argcheck (L, ud != NULL, pos, "'cdb' expected"); return ud ? *((struct cdb **)ud) : NULL; } +static struct cdb_make * +lua_check_cdb_builder (lua_State * L, int pos) +{ + void *ud = rspamd_lua_check_udata (L, pos, "rspamd{cdb_builder}"); + + luaL_argcheck (L, ud != NULL, pos, "'cdb_builder' expected"); + return ud ? ((struct cdb_make *)ud) : NULL; +} + static gint lua_cdb_create (lua_State *L) { struct cdb *cdb, **pcdb; const gchar *filename; gint fd; - struct ev_loop *ev_base = lua_check_ev_base (L, 2); + + struct ev_loop *ev_base = NULL; + + if (lua_type(L, 2) == LUA_TUSERDATA) { + ev_base = lua_check_ev_base(L, 2); + } filename = luaL_checkstring (L, 1); /* If file begins with cdb://, just skip it */ @@ -63,7 +94,7 @@ lua_cdb_create (lua_State *L) lua_pushnil (L); } else { - cdb = g_malloc (sizeof (struct cdb)); + cdb = g_malloc0 (sizeof (struct cdb)); cdb->filename = g_strdup (filename); if (cdb_init (cdb, fd) == -1) { g_free (cdb->filename); @@ -87,7 +118,9 @@ lua_cdb_create (lua_State *L) } } #endif - cdb_add_timer (cdb, ev_base, CDB_REFRESH_TIME); + if (ev_base) { + cdb_add_timer(cdb, ev_base, CDB_REFRESH_TIME); + } pcdb = lua_newuserdata (L, sizeof (struct cdb *)); rspamd_lua_setclass (L, "rspamd{cdb}", -1); *pcdb = cdb; @@ -100,7 +133,7 @@ lua_cdb_create (lua_State *L) static gint lua_cdb_get_name (lua_State *L) { - struct cdb *cdb = lua_check_cdb (L); + struct cdb *cdb = lua_check_cdb (L, 1); if (!cdb) { lua_error (L); @@ -113,19 +146,18 @@ lua_cdb_get_name (lua_State *L) static gint lua_cdb_lookup (lua_State *L) { - struct cdb *cdb = lua_check_cdb (L); - const gchar *what; + struct cdb *cdb = lua_check_cdb (L, 1); + gsize klen; + const gchar *what = luaL_checklstring(L, 2, &klen); gchar *value; gsize vlen; gint64 vpos; - if (!cdb) { - lua_error (L); - return 1; + if (!cdb || what == NULL) { + return lua_error (L); } - what = luaL_checkstring (L, 2); - if (cdb_find (cdb, what, strlen (what)) > 0) { + if (cdb_find (cdb, what, klen) > 0) { /* Extract and push value to lua as string */ vpos = cdb_datapos (cdb); vlen = cdb_datalen (cdb); @@ -144,11 +176,13 @@ lua_cdb_lookup (lua_State *L) static gint lua_cdb_destroy (lua_State *L) { - struct cdb *cdb = lua_check_cdb (L); + struct cdb *cdb = lua_check_cdb (L, 1); if (cdb) { cdb_free (cdb); - (void)close (cdb->cdb_fd); + if (cdb->cdb_fd != -1) { + (void) close(cdb->cdb_fd); + } g_free (cdb->filename); g_free (cdb); } @@ -156,6 +190,108 @@ lua_cdb_destroy (lua_State *L) return 0; } +static gint +lua_cdb_build (lua_State *L) +{ + const char *filename = luaL_checkstring (L, 1); + int fd, mode = 00755; + + if (filename == NULL) { + return luaL_error (L, "invalid arguments, filename expected"); + } + + /* If file begins with cdb://, just skip it */ + if (g_ascii_strncasecmp (filename, "cdb://", sizeof ("cdb://") - 1) == 0) { + filename += sizeof ("cdb://") - 1; + } + + if (lua_isnumber (L, 2)) { + mode = lua_tointeger (L, 2); + } + + fd = rspamd_file_xopen (filename, O_RDWR | O_CREAT | O_TRUNC, mode, 0); + + if (fd == -1) { + lua_pushnil (L); + lua_pushfstring (L, "cannot open cdb: %s, %s", filename, strerror (errno)); + + return 2; + } + + struct cdb_make *cdbm = lua_newuserdata (L, sizeof(struct cdb_make)); + + g_assert (cdb_make_start(cdbm, fd) == 0); + rspamd_lua_setclass (L, "rspamd{cdb_builder}", -1); + + return 1; +} + +static gint +lua_cdb_builder_add (lua_State *L) +{ + struct cdb_make *cdbm = lua_check_cdb_builder(L, 1); + gsize data_sz, key_sz; + const char *key = lua_tolstring (L, 2, &key_sz); + const char *data = lua_tolstring (L, 3, &data_sz); + + if (cdbm == NULL || key == NULL || data == NULL || cdbm->cdb_fd == -1) { + return luaL_error(L, "invalid arguments"); + } + + if (cdb_make_add (cdbm, key, key_sz, data, data_sz) == -1) { + lua_pushvalue(L, 1); + lua_pushfstring(L, "cannot push value to cdb: %s", strerror(errno)); + + return 2; + } + + /* Allow chaining */ + lua_pushvalue(L, 1); + return 1; +} + +static gint +lua_cdb_builder_finalize (lua_State *L) +{ + struct cdb_make *cdbm = lua_check_cdb_builder(L, 1); + + if (cdbm == NULL || cdbm->cdb_fd == -1) { + return luaL_error(L, "invalid arguments"); + } + + if (cdb_make_finish (cdbm) == -1) { + lua_pushvalue(L, 1); + lua_pushfstring(L, "cannot finish value to cdb: %s", strerror(errno)); + + return 2; + } + + close (cdbm->cdb_fd); + cdbm->cdb_fd = -1; /* To distinguish finalized object */ + + /* Allow chaining */ + lua_pushvalue (L, 1); + return 1; +} + +static gint +lua_cdb_builder_dtor (lua_State *L) +{ + struct cdb_make *cdbm = lua_check_cdb_builder(L, 1); + + if (cdbm == NULL) { + return luaL_error(L, "invalid arguments"); + } + + if (cdbm->cdb_fd != -1) { + cdb_make_finish (cdbm); + close (cdbm->cdb_fd); + cdbm->cdb_fd = -1; /* Finalized object */ + } + + return 0; +} + static gint lua_load_cdb (lua_State *L) { @@ -170,5 +306,7 @@ luaopen_cdb (lua_State * L) { rspamd_lua_new_class (L, "rspamd{cdb}", cdblib_m); lua_pop (L, 1); + rspamd_lua_new_class (L, "rspamd{cdb_builder}", cdbbuilderlib_m); + lua_pop (L, 1); rspamd_lua_add_preload (L, "rspamd_cdb", lua_load_cdb); } -- 2.39.5