[Project] Implement selectors registration for regular expressions

3 files changed, 150 insertions, 0 deletions

diff --git a/src/libserver/re_cache.c b/src/libserver/re_cache.c
index d8f7f3d0d..0a70c85b6 100644
--- 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 90acd1501..c14b29ef0 100644
--- 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 382a34a0b..e989bb77f 100644
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -262,6 +262,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),
@@ -3577,6 +3588,89 @@ lua_config_init_subsystem (lua_State *L)
 }
 
 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)
 {
 	LUA_TRACE_POINT;