[Feature] Lua_cdb: Add cdb building interface

diff --git a/src/lua/lua_cdb.c b/src/lua/lua_cdb.c
index 8f13f6ccfabac189e9d027a3936444021baf2c70..91f770cb3fea010f305a08a209f9354b2c87a588 100644 (file)
--- 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);
 }