commit 05457bf: [Feature] Lua_cdb: Add cdb building interface

Vsevolod Stakhov vsevolod at highsecure.ru
Tue Oct 19 19:56:05 UTC 2021


Author: Vsevolod Stakhov
Date: 2021-10-19 17:49:34 +0100
URL: https://github.com/rspamd/rspamd/commit/05457bf4200250d8262bd1de87b0a33545cf70a4

[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);
 }


More information about the Commits mailing list