[Feature] Maps: Add on_load support

diff --git a/src/libserver/maps/map.c b/src/libserver/maps/map.c
index bb1f1f3fc06b12aac6d96b8469ffcd04b68c1957..04557f0fef9e77178e0378c373e12e9a2c01609f 100644 (file)
--- a/src/libserver/maps/map.c
+++ b/src/libserver/maps/map.c
@@ -995,6 +995,10 @@ rspamd_map_periodic_dtor (struct map_periodic_cbdata *periodic)
        if (periodic->need_modify || periodic->cbdata.errored) {
                /* Need to notify the real data structure */
                periodic->map->fin_callback (&periodic->cbdata, periodic->map->user_data);
+
+               if (map->on_load_function) {
+                       map->on_load_function(map, map->on_load_ud);
+               }
        }
        else {
                /* Not modified */
@@ -2300,6 +2304,10 @@ rspamd_map_preload (struct rspamd_config *cfg)
 
                        if (succeed) {
                                map->fin_callback (&fake_cbd.cbdata, map->user_data);
+
+                               if (map->on_load_function) {
+                                       map->on_load_function(map, map->on_load_ud);
+                               }
                        }
                        else {
                                msg_info_map ("preload of %s failed", map->name);
@@ -2336,6 +2344,10 @@ rspamd_map_remove_all (struct rspamd_config *cfg)
                        *map->user_data = NULL;
                }
 
+               if (map->on_load_ud_dtor) {
+                       map->on_load_ud_dtor(map->on_load_ud);
+               }
+
                for (i = 0; i < map->backends->len; i ++) {
                        bk = g_ptr_array_index (map->backends, i);
 
@@ -3106,3 +3118,14 @@ rspamd_map_traverse (struct rspamd_map *map, rspamd_map_traverse_cb cb,
                map->traverse_function (*map->user_data, cb, cbdata, reset_hits);
        }
 }
+
+void
+rspamd_map_set_on_load_function (struct rspamd_map *map, rspamd_map_on_load_function cb,
+       gpointer cbdata, GDestroyNotify dtor)
+{
+       if (map) {
+               map->on_load_function = cb;
+               map->on_load_ud = cbdata;
+               map->on_load_ud_dtor = dtor;
+       }
+}

diff --git a/src/libserver/maps/map.h b/src/libserver/maps/map.h
index 6d77454fb3152d7af05e9f9209f95afddbbd10b8..ac2edc82a88ea8d736bbc39613ef257ca6ea8fe1 100644 (file)
--- a/src/libserver/maps/map.h
+++ b/src/libserver/maps/map.h
@@ -21,6 +21,12 @@ extern "C" {
 struct map_cb_data;
 struct rspamd_worker;
 
+/**
+ * Common map object
+ */
+struct rspamd_config;
+struct rspamd_map;
+
 /**
  * Callback types
  */
@@ -37,12 +43,7 @@ typedef gboolean (*rspamd_map_traverse_cb) (gconstpointer key,
 typedef void (*rspamd_map_traverse_function) (void *data,
                                                                                          rspamd_map_traverse_cb cb,
                                                                                          gpointer cbdata, gboolean reset_hits);
-
-/**
- * Common map object
- */
-struct rspamd_config;
-struct rspamd_map;
+typedef void (*rspamd_map_on_load_function) (struct rspamd_map *map, gpointer ud);
 
 /**
  * Callback data for async load
@@ -151,6 +152,15 @@ rspamd_map_traverse_function rspamd_map_get_traverse_function (struct rspamd_map
 void rspamd_map_traverse (struct rspamd_map *map, rspamd_map_traverse_cb cb,
                                                  gpointer cbdata, gboolean reset_hits);
 
+/**
+ * Set map on load callback
+ * @param map
+ * @param cb
+ * @param cbdata
+ */
+void rspamd_map_set_on_load_function (struct rspamd_map *map, rspamd_map_on_load_function cb,
+                                                                gpointer cbdata, GDestroyNotify dtor);
+
 #ifdef  __cplusplus
 }
 #endif

diff --git a/src/libserver/maps/map_private.h b/src/libserver/maps/map_private.h
index 74b2ea0426eeff832ba7967d3ba7a2c04d9fb656..bbbac0cd6f4116484c2b285a81ae4345f53cf499 100644 (file)
--- a/src/libserver/maps/map_private.h
+++ b/src/libserver/maps/map_private.h
@@ -151,6 +151,9 @@ struct rspamd_map {
        rspamd_map_tmp_dtor tmp_dtor;
        gpointer tmp_dtor_data;
        rspamd_map_traverse_function traverse_function;
+       rspamd_map_on_load_function on_load_function;
+       gpointer on_load_ud;
+       GDestroyNotify on_load_ud_dtor;
        gpointer lua_map;
        gsize nelts;
        guint64 digest;

diff --git a/src/lua/lua_map.c b/src/lua/lua_map.c
index fe01c303112d8f9418b5f788cc2f0275bce337d3..29e2053f2747598979ca55e62ceee27b4f4d5d11 100644 (file)
--- a/src/lua/lua_map.c
+++ b/src/lua/lua_map.c
@@ -135,6 +135,13 @@ LUA_FUNCTION_DEF (map, get_stats);
  */
 LUA_FUNCTION_DEF (map, foreach);
 
+/***
+ * @method map:on_load(callback)
+ * Sets a callback for a map that is called when map is loaded
+ * @param {function} callback callback function, that accepts no arguments (pass maps in a closure if needed)
+ */
+LUA_FUNCTION_DEF (map, on_load);
+
 /***
  * @method map:get_data_digest()
  * Get data digest for specific map
@@ -159,6 +166,7 @@ static const struct luaL_reg maplib_m[] = {
        LUA_INTERFACE_DEF (map, get_uri),
        LUA_INTERFACE_DEF (map, get_stats),
        LUA_INTERFACE_DEF (map, foreach),
+       LUA_INTERFACE_DEF (map, on_load),
        LUA_INTERFACE_DEF (map, get_data_digest),
        LUA_INTERFACE_DEF (map, get_nelts),
        {"__tostring", rspamd_lua_class_tostring},
@@ -1131,6 +1139,7 @@ lua_map_foreach (lua_State * L)
                cbdata.L = L;
                lua_pushvalue (L, 2); /* func */
                cbdata.cbref = lua_gettop (L);
+               cbdata.use_text = use_text;
 
                if (map->map->traverse_function) {
                        rspamd_map_traverse (map->map, lua_map_foreach_cb, &cbdata, FALSE);
@@ -1364,6 +1373,58 @@ lua_map_get_uri (lua_State *L)
        return map->map->backends->len;
 }
 
+struct lua_map_on_load_cbdata {
+       lua_State *L;
+       gint ref;
+};
+
+static void
+lua_map_on_load_dtor (gpointer p)
+{
+       struct lua_map_on_load_cbdata *cbd = p;
+
+       luaL_unref (cbd->L, LUA_REGISTRYINDEX, cbd->ref);
+       g_free (cbd);
+}
+
+static void
+lua_map_on_load_handler (struct rspamd_map *map, gpointer ud)
+{
+       struct lua_map_on_load_cbdata *cbd = ud;
+       lua_State *L = cbd->L;
+
+       lua_rawgeti (L, LUA_REGISTRYINDEX, cbd->ref);
+
+       if (lua_pcall(L, 0, 0, 0) != 0) {
+               msg_err_map ("call to on_load function failed: %s", lua_tostring (L, -1));
+       }
+}
+
+static gint
+lua_map_on_load (lua_State *L)
+{
+       LUA_TRACE_POINT;
+       struct rspamd_lua_map *map = lua_check_map (L, 1);
+
+       if (map == NULL) {
+               return luaL_error (L, "invalid arguments");
+       }
+
+       if (lua_type (L, 2) == LUA_TFUNCTION) {
+               lua_pushvalue (L, 2);
+               struct lua_map_on_load_cbdata *cbd = g_malloc (sizeof (struct lua_map_on_load_cbdata));
+               cbd->L = L;
+               cbd->ref = luaL_ref (L, LUA_REGISTRYINDEX);
+
+               rspamd_map_set_on_load_function(map->map, lua_map_on_load_handler, cbd, lua_map_on_load_dtor);
+       }
+       else {
+               return luaL_error (L, "invalid callback");
+       }
+
+       return 0;
+}
+
 void
 luaopen_map (lua_State * L)
 {