[Project] Implement selectors registration for regular expressions

diff --git a/src/libserver/re_cache.c b/src/libserver/re_cache.c
index d8f7f3d0dd80ce1c68db7c20d07989df1953ea16..0a70c85b6d2b4154dbc6dfda572b8636c42aefde 100644 (file)
--- a/src/libserver/re_cache.c
+++ b/src/libserver/re_cache.c
@@ -23,6 +23,9 @@
 #include "libutil/util.h"
 #include "libutil/regexp.h"
 #include "lua/lua_common.h"
+
+#include "khash.h"
+
 #ifdef WITH_HYPERSCAN
 #include "hs.h"
 #include "unix-std.h"
@@ -70,6 +73,7 @@ static const guchar rspamd_hs_magic[] = {'r', 's', 'h', 's', 'r', 'e', '1', '1'}
                rspamd_hs_magic_vector[] = {'r', 's', 'h', 's', 'r', 'v', '1', '1'};
 #endif
 
+
 struct rspamd_re_class {
        guint64 id;
        enum rspamd_re_type type;
@@ -97,13 +101,18 @@ struct rspamd_re_cache_elt {
        enum rspamd_re_cache_elt_match_type match_type;
 };
 
+KHASH_INIT (lua_selectors_hash, gchar *, int, 1, kh_str_hash_func, kh_str_hash_equal);
+
 struct rspamd_re_cache {
        GHashTable *re_classes;
+
        GPtrArray *re;
+       khash_t (lua_selectors_hash) *selectors;
        ref_entry_t ref;
        guint nre;
        guint max_re_data;
        gchar hash[rspamd_cryptobox_HASHBYTES + 1];
+       lua_State *L;
 #ifdef WITH_HYPERSCAN
        gboolean hyperscan_loaded;
        gboolean disable_hyperscan;
@@ -149,6 +158,8 @@ rspamd_re_cache_destroy (struct rspamd_re_cache *cache)
        GHashTableIter it;
        gpointer k, v;
        struct rspamd_re_class *re_class;
+       gchar *skey;
+       gint sref;
 
        g_assert (cache != NULL);
        g_hash_table_iter_init (&it, cache->re_classes);
@@ -176,6 +187,15 @@ rspamd_re_cache_destroy (struct rspamd_re_cache *cache)
                g_free (re_class);
        }
 
+       if (cache->L) {
+               kh_foreach (cache->selectors, skey, sref, {
+                       luaL_unref (cache->L, LUA_REGISTRYINDEX, sref);
+                       g_free (skey);
+               });
+       }
+
+       kh_destroy (lua_selectors_hash, cache->selectors);
+
        g_hash_table_unref (cache->re_classes);
        g_ptr_array_free (cache->re, TRUE);
        g_free (cache);
@@ -199,6 +219,7 @@ rspamd_re_cache_new (void)
        cache->re_classes = g_hash_table_new (g_int64_hash, g_int64_equal);
        cache->nre = 0;
        cache->re = g_ptr_array_new_full (256, rspamd_re_cache_elt_dtor);
+       cache->selectors = kh_init (lua_selectors_hash);
 #ifdef WITH_HYPERSCAN
        cache->hyperscan_loaded = FALSE;
 #endif
@@ -413,6 +434,8 @@ rspamd_re_cache_init (struct rspamd_re_cache *cache, struct rspamd_config *cfg)
                }
        }
 
+       cache->L = cfg->lua_state;
+
 #ifdef WITH_HYPERSCAN
        const gchar *platform = "generic";
        rspamd_fstring_t *features = rspamd_fstring_new ();
@@ -1979,3 +2002,30 @@ rspamd_re_cache_load_hyperscan (struct rspamd_re_cache *cache,
        return TRUE;
 #endif
 }
+
+void rspamd_re_cache_add_selector (struct rspamd_re_cache *cache,
+                                                                  const gchar *sname,
+                                                                  gint ref)
+{
+       khiter_t k;
+
+       k = kh_get (lua_selectors_hash, cache->selectors, (gchar *)sname);
+
+       if (k == kh_end (cache->selectors)) {
+               gchar *cpy = g_strdup (sname);
+               gint res;
+
+               k = kh_put (lua_selectors_hash, cache->selectors, cpy, &res);
+
+               kh_value (cache->selectors, k) = ref;
+       }
+       else {
+               msg_warn_re_cache ("replacing selector with name %s", sname);
+
+               if (cache->L) {
+                       luaL_unref (cache->L, LUA_REGISTRYINDEX, kh_value (cache->selectors, k));
+               }
+
+               kh_value (cache->selectors, k) = ref;
+       }
+}

diff --git a/src/libserver/re_cache.h b/src/libserver/re_cache.h
index 90acd150146cca8ebfd480731bce06ac139336e7..c14b29ef0d8a76f80b9bec70e57bb307cde691d9 100644 (file)
--- a/src/libserver/re_cache.h
+++ b/src/libserver/re_cache.h
@@ -171,4 +171,10 @@ gboolean rspamd_re_cache_is_valid_hyperscan_file (struct rspamd_re_cache *cache,
  */
 gboolean rspamd_re_cache_load_hyperscan (struct rspamd_re_cache *cache,
                const char *cache_dir);
+
+/**
+ * Registers lua selector in the cache
+ */
+void rspamd_re_cache_add_selector (struct rspamd_re_cache *cache,
+               const gchar *sname, gint ref);
 #endif

diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c
index 382a34a0b4a682b1836cccfdc1466bc03631df8e..e989bb77f59ab5c4053fdd20d5449fca838ccfa9 100644 (file)
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -261,6 +261,16 @@ rspamd_config:register_dependency('SYMBOL_FROM', 'SYMBOL_TO')
  */
 LUA_FUNCTION_DEF (config, register_dependency);
 
+/**
+ * @method rspamd_config:register_re_selector(name, selector_str)
+ * Registers selector with the specific name to use in regular expressions in form
+ * name=/re/$ or name=/re/{selector}
+ * @param {string} name name of the selector
+ * @param {selector_str} selector string
+ * @return true if selector has been registered
+ */
+LUA_FUNCTION_DEF (config, register_re_selector);
+
 /**
  * @method rspamd_config:set_symbol({table})
  * Sets the value of a specified symbol in a metric. This function accepts table with the following elements:
@@ -771,6 +781,7 @@ static const struct luaL_reg configlib_m[] = {
        LUA_INTERFACE_DEF (config, register_regexp),
        LUA_INTERFACE_DEF (config, replace_regexp),
        LUA_INTERFACE_DEF (config, register_worker_script),
+       LUA_INTERFACE_DEF (config, register_re_selector),
        LUA_INTERFACE_DEF (config, add_on_load),
        LUA_INTERFACE_DEF (config, add_periodic),
        LUA_INTERFACE_DEF (config, get_symbols_count),
@@ -3576,6 +3587,89 @@ lua_config_init_subsystem (lua_State *L)
        return 0;
 }
 
+static gint
+lua_config_register_re_selector (lua_State *L)
+{
+       LUA_TRACE_POINT;
+       struct rspamd_config *cfg = lua_check_config (L, 1);
+       const gchar *name = luaL_checkstring (L, 2);
+       const gchar *selector_str = luaL_checkstring (L, 3);
+       gint top = lua_gettop (L);
+       bool res = false;
+
+       if (cfg && name && selector_str) {
+               if (luaL_dostring (L, "return require \"lua_selectors\"") != 0) {
+                       msg_warn_config ("cannot require lua_selectors: %s",
+                                       lua_tostring (L, -1));
+               }
+               else {
+                       if (lua_type (L, -1) != LUA_TTABLE) {
+                               msg_warn_config ("lua selectors must return "
+                                                                "table and not %s",
+                                               lua_typename (L, lua_type (L, -1)));
+                       }
+                       else {
+                               lua_pushstring (L, "create_selector_closure");
+                               lua_gettable (L, -2);
+
+                               if (lua_type (L, -1) != LUA_TFUNCTION) {
+                                       msg_warn_config ("create_selector_closure must return "
+                                                                        "function and not %s",
+                                                       lua_typename (L, lua_type (L, -1)));
+                               }
+                               else {
+                                       gint err_idx, ret;
+                                       GString *tb;
+                                       struct rspamd_config **pcfg;
+
+                                       lua_pushcfunction (L, &rspamd_lua_traceback);
+                                       err_idx = lua_gettop (L);
+
+                                       /* Push function */
+                                       lua_pushvalue (L, -2);
+
+                                       pcfg = lua_newuserdata (L, sizeof (*pcfg));
+                                       rspamd_lua_setclass (L, "rspamd{config}", -1);
+                                       *pcfg = cfg;
+                                       lua_pushstring (L, selector_str);
+
+                                       if ((ret = lua_pcall (L, 2, 1, err_idx)) != 0) {
+                                               tb = lua_touserdata (L, -1);
+                                               msg_err_config ("call to create_selector_closure lua "
+                                                                               "script failed (%d): %v", ret, tb);
+
+                                               if (tb) {
+                                                       g_string_free (tb, TRUE);
+                                               }
+                                       }
+                                       else {
+                                               if (lua_type (L, -1) != LUA_TFUNCTION) {
+                                                       msg_warn_config ("create_selector_closure "
+                                                                                        "invocation must return "
+                                                                                        "function and not %s",
+                                                                       lua_typename (L, lua_type (L, -1)));
+                                               }
+                                               else {
+                                                       ret = luaL_ref (L, LUA_REGISTRYINDEX);
+                                                       rspamd_re_cache_add_selector (cfg->re_cache,
+                                                                       name, ret);
+                                                       res = true;
+                                               }
+                                       }
+                               }
+                       }
+               }
+       }
+       else {
+               return luaL_error (L, "invalid arguments");
+       }
+
+       lua_settop (L, top);
+       lua_pushboolean (L, res);
+
+       return 1;
+}
+
 static gint
 lua_config_get_tld_path (lua_State *L)
 {