commit cb40d91: [Project] Add C method to process ANN tokens vector

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2019-07-06 12:29:25 +0100
URL: https://github.com/rspamd/rspamd/commit/cb40d91448fee3cd8069a155e42effd39c8b0f92

[Project] Add C method to process ANN tokens vector

---
 src/libserver/rspamd_symcache.c | 10 +++++++++
 src/libserver/rspamd_symcache.h |  6 +++++
 src/lua/lua_task.c              | 50 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/src/libserver/rspamd_symcache.c b/src/libserver/rspamd_symcache.c
index d6fa726c6..3de339654 100644
--- a/src/libserver/rspamd_symcache.c
+++ b/src/libserver/rspamd_symcache.c
@@ -3390,4 +3390,14 @@ rspamd_symcache_process_settings_elt (struct rspamd_symcache *cache,
 			}
 		}
 	}
+}
+
+enum rspamd_symbol_type
+rspamd_symcache_item_flags (struct rspamd_symcache_item *item)
+{
+	if (item) {
+		return item->type;
+	}
+
+	return 0;
 }
\ No newline at end of file
diff --git a/src/libserver/rspamd_symcache.h b/src/libserver/rspamd_symcache.h
index 7bfb1d093..640a1c898 100644
--- a/src/libserver/rspamd_symcache.h
+++ b/src/libserver/rspamd_symcache.h
@@ -493,4 +493,10 @@ gboolean rspamd_symcache_is_item_allowed (struct rspamd_task *task,
 										  struct rspamd_symcache_item *item,
 										  gboolean exec_only);
 
+/**
+ * Returns symbcache item flags
+ * @param item
+ * @return
+ */
+enum rspamd_symbol_type rspamd_symcache_item_flags (struct rspamd_symcache_item *item);
 #endif
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index f9033eef8..99faa4b3e 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -640,6 +640,16 @@ LUA_FUNCTION_DEF (task, get_symbols_numeric);
  */
 LUA_FUNCTION_DEF (task, get_symbols_tokens);
 
+/***
+ * @method task:process_ann_tokens(symbols, ann_tokens, offset)
+ * Processes ann tokens
+ * @param {table|string} symbols list of symbols in this profile
+ * @param {table|number} ann_tokens list of tokens (including metatokens)
+ * @param {integer} offset offset for symbols token (#metatokens)
+ * @return nothing
+ */
+LUA_FUNCTION_DEF (task, process_ann_tokens);
+
 /***
  * @method task:has_symbol(name)
  * Fast path to check if a specified symbol is in the task's results
@@ -1118,6 +1128,7 @@ static const struct luaL_reg tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, get_symbols_all),
 	LUA_INTERFACE_DEF (task, get_symbols_numeric),
 	LUA_INTERFACE_DEF (task, get_symbols_tokens),
+	LUA_INTERFACE_DEF (task, process_ann_tokens),
 	LUA_INTERFACE_DEF (task, has_symbol),
 	LUA_INTERFACE_DEF (task, enable_symbol),
 	LUA_INTERFACE_DEF (task, disable_symbol),
@@ -4236,6 +4247,45 @@ lua_task_get_symbols_tokens (lua_State *L)
 	return 1;
 }
 
+static gint
+lua_task_process_ann_tokens (lua_State *L)
+{
+	LUA_TRACE_POINT;
+	struct rspamd_task *task = lua_check_task (L, 1);
+	gint offset = luaL_checkinteger (L, 4);
+
+	if (task && lua_istable (L, 2) && lua_istable (L, 3)) {
+		guint symlen = rspamd_lua_table_size (L, 2);
+
+		for (guint i = 1; i <= symlen; i ++, offset ++) {
+			const gchar *sym;
+			struct rspamd_symbol_result *sres;
+
+			lua_rawgeti (L, 2, i);
+			sym = lua_tostring (L, -1);
+
+			sres = rspamd_task_find_symbol_result (task, sym);
+
+			if (sres && !(sres->flags & RSPAMD_SYMBOL_RESULT_IGNORED)) {
+				if (!isnan (sres->score) &&
+					!(rspamd_symcache_item_flags (sres->sym->cache_item) &
+					  SYMBOL_TYPE_NOSTAT)) {
+
+					lua_pushnumber (L, tanh (sres->score));
+					lua_rawseti (L, 3, offset + 1);
+				}
+			}
+
+			lua_pop (L, 1); /* Symbol name */
+		}
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 0;
+}
+
 enum lua_date_type {
 	DATE_CONNECT = 0,
 	DATE_MESSAGE,