* Add trie interface to lua api

4 files changed, 151 insertions, 0 deletions

diff --git a/src/expressions.c b/src/expressions.c
index d6e6647f1..5df36711b 100644
--- a/src/expressions.c
+++ b/src/expressions.c
@@ -962,7 +962,11 @@ rspamd_parts_distance (struct worker_task * task, GList * args, void *unused)
 		if (p1->parent && p1->parent == p2->parent) {
 			parent = p1->parent;
 			ct = g_mime_object_get_content_type (parent);
+#ifndef GMIME24
 			if (ct == NULL || ! g_mime_content_type_is_type (ct, "multipart", "alternative")) {
+#else
+			if (ct == NULL || ! g_mime_content_type_is_type ((GMimeContentType *)ct, "multipart", "alternative")) {
+#endif
 				debug_task ("two parts are not belong to multipart/alternative container, skip check");
 				return FALSE;
 			}
diff --git a/src/lua/lua_common.c b/src/lua/lua_common.c
index c37588bb9..adf26e686 100644
--- a/src/lua/lua_common.c
+++ b/src/lua/lua_common.c
@@ -224,6 +224,7 @@ init_lua (struct config_file *cfg)
 	(void)luaopen_config (L);
 	(void)luaopen_radix (L);
 	(void)luaopen_hash_table (L);
+	(void)luaopen_trie (L);
 	(void)luaopen_task (L);
 	(void)luaopen_textpart (L);
 	(void)luaopen_image (L);
diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index 18fa6481d..105fb7ca2 100644
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -27,6 +27,7 @@ int luaopen_config (lua_State *L);
 int luaopen_metric (lua_State *L);
 int luaopen_radix (lua_State *L);
 int luaopen_hash_table (lua_State *L);
+int luaopen_trie (lua_State * L);
 int luaopen_textpart (lua_State *L);
 int luaopen_image (lua_State *L);
 int luaopen_classifier (lua_State *L);
diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c
index 7ae7d3e0e..a247bbbf7 100644
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -26,7 +26,9 @@
 #include "lua_common.h"
 #include "../expressions.h"
 #include "../map.h"
+#include "../message.h"
 #include "../radix.h"
+#include "../trie.h"
 #include "../classifiers/classifiers.h"
 
 /* Config file methods */
@@ -71,6 +73,21 @@ static const struct luaL_reg    hashlib_m[] = {
 	{NULL, NULL}
 };
 
+/* Suffix trie */
+LUA_FUNCTION_DEF (trie, create);
+LUA_FUNCTION_DEF (trie, add_pattern);
+LUA_FUNCTION_DEF (trie, search_text);
+LUA_FUNCTION_DEF (trie, search_task);
+
+static const struct luaL_reg    trielib_m[] = {
+	LUA_INTERFACE_DEF (trie, create),
+	LUA_INTERFACE_DEF (trie, add_pattern),
+	LUA_INTERFACE_DEF (trie, search_task),
+	LUA_INTERFACE_DEF (trie, search_task),
+	{"__tostring", lua_class_tostring},
+	{NULL, NULL}
+};
+
 static struct config_file      *
 lua_check_config (lua_State * L)
 {
@@ -95,6 +112,14 @@ lua_check_hash_table (lua_State * L)
 	return **((GHashTable ***)ud);
 }
 
+static rspamd_trie_t          *
+lua_check_trie (lua_State * L)
+{
+	void                           *ud = luaL_checkudata (L, 1, "rspamd{trie}");
+	luaL_argcheck (L, ud != NULL, 1, "'trie' expected");
+	return **((rspamd_trie_t ***)ud);
+}
+
 /*** Config functions ***/
 static int
 lua_config_get_module_opt (lua_State * L)
@@ -478,6 +503,117 @@ lua_hash_table_get_key (lua_State * L)
 	return 1;
 }
 
+/* Trie functions */
+static int
+lua_trie_create (lua_State *L)
+{
+	rspamd_trie_t                 *trie, **ptrie;
+	gboolean                       icase = FALSE;
+
+	if (lua_gettop (L) == 1) {
+		icase = lua_toboolean (L, 1);
+	}
+
+	trie = rspamd_trie_create (icase);
+
+	ptrie = lua_newuserdata (L, sizeof (rspamd_trie_t *));
+	lua_setclass (L, "rspamd{trie}", -1);
+	*ptrie = trie;
+
+	return 1;
+}
+
+static int
+lua_trie_add_pattern (lua_State *L)
+{
+	rspamd_trie_t                 *trie = lua_check_trie (L);
+	const gchar                   *pattern;
+	gint                           id;
+
+	if (trie) {
+		pattern  = luaL_checkstring (L, 2);
+		id = luaL_checknumber (L, 3);
+
+		if (pattern != NULL) {
+			rspamd_trie_insert (trie, pattern, id);
+			lua_pushboolean (L, 1);
+		}
+	}
+
+	lua_pushboolean (L, 0);
+
+	return 1;
+}
+
+static int
+lua_trie_search_text (lua_State *L)
+{
+	rspamd_trie_t                 *trie = lua_check_trie (L);
+	const gchar                   *text, *pos;
+	gint                           id, i = 1;
+	gsize                          len;
+
+	if (trie) {
+		text = luaL_checkstring (L, 2);
+		len = strlen (text);
+		if (text) {
+			lua_newtable (L);
+			pos = text;
+			while (pos < text + len && (pos = rspamd_trie_lookup (trie, pos, len, &id)) != NULL) {
+				lua_pushinteger (L, i);
+				lua_pushinteger (L, id);
+				lua_settable (L, -3);
+				i ++;
+			}
+			return 1;
+		}
+	}
+
+	lua_pushnil (L);
+	return 1;
+}
+
+static int
+lua_trie_search_task (lua_State *L)
+{
+	rspamd_trie_t                 *trie = lua_check_trie (L);
+	struct worker_task            *task;
+	struct mime_text_part         *part;
+	GList                         *cur;
+	const gchar                   *pos, *end;
+	gint                           id, i = 1;
+	void                           *ud;
+
+	if (trie) {
+		ud = luaL_checkudata (L, 2, "rspamd{task}");
+		luaL_argcheck (L, ud != NULL, 1, "'task' expected");
+		task = *((struct worker_task **)ud);
+		if (task) {
+			lua_newtable (L);
+			cur = task->text_parts;
+			while (cur) {
+				part = cur->data;
+				if (!part->is_empty && part->content != NULL) {
+					pos = (const gchar *)part->content->data;
+					end = pos + part->content->len;
+					while (pos < end && (pos = rspamd_trie_lookup (trie, pos, part->content->len, &id)) != NULL) {
+						lua_pushinteger (L, i);
+						lua_pushinteger (L, id);
+						lua_settable (L, -3);
+						i ++;
+					}
+				}
+				cur = g_list_next (cur);
+			}
+			return 1;
+		}
+	}
+
+	lua_pushnil (L);
+	return 1;
+}
+/* Init functions */
+
 int
 luaopen_config (lua_State * L)
 {
@@ -504,3 +640,12 @@ luaopen_hash_table (lua_State * L)
 
 	return 1;
 }
+
+int
+luaopen_trie (lua_State * L)
+{
+	lua_newclass (L, "rspamd{trie}", trielib_m);
+	luaL_openlib (L, "rspamd_trie", null_reg, 0);
+
+	return 1;
+}