[Feature] Allow augmentations set in Lua API

6 files changed, 111 insertions, 19 deletions

diff --git a/src/libserver/rspamd_symcache.h b/src/libserver/rspamd_symcache.h
index 99158fbc3..95ca614af 100644
--- a/src/libserver/rspamd_symcache.h
+++ b/src/libserver/rspamd_symcache.h
@@ -127,6 +127,16 @@ gint rspamd_symcache_add_symbol (struct rspamd_symcache *cache,
 								 gint parent);
 
 /**
+ * Adds augmentation to the symbol
+ * @param cache
+ * @param sym_id
+ * @param augmentation
+ * @return
+ */
+bool rspamd_symcache_add_symbol_augmentation(struct rspamd_symcache *cache,
+		int sym_id, const char *augmentation);
+
+/**
  * Add callback to be executed whenever symbol has peak value
  * @param cache
  * @param cbref
diff --git a/src/libserver/symcache/symcache_c.cxx b/src/libserver/symcache/symcache_c.cxx
index fafc89d19..4d1f28800 100644
--- a/src/libserver/symcache/symcache_c.cxx
+++ b/src/libserver/symcache/symcache_c.cxx
@@ -84,6 +84,29 @@ rspamd_symcache_add_symbol(struct rspamd_symcache *cache,
 	}
 }
 
+bool
+rspamd_symcache_add_symbol_augmentation(struct rspamd_symcache *cache,
+											 int sym_id, const char *augmentation)
+{
+	auto *real_cache = C_API_SYMCACHE(cache);
+	auto log_tag = [&]() { return real_cache->log_tag(); };
+
+	if (augmentation == nullptr) {
+		msg_err_cache("null augmentation is not allowed for item %d", sym_id);
+		return false;
+	}
+
+
+	auto *item = real_cache->get_item_by_id_mut(sym_id, false);
+
+	if (item == nullptr) {
+		msg_err_cache("item %d is not found", sym_id);
+		return false;
+	}
+
+	return item->add_augmentation(*real_cache, augmentation);
+}
+
 void
 rspamd_symcache_set_peak_callback(struct rspamd_symcache *cache, gint cbref)
 {
diff --git a/src/libserver/symcache/symcache_impl.cxx b/src/libserver/symcache/symcache_impl.cxx
index e557f6212..fb3d2d965 100644
--- a/src/libserver/symcache/symcache_impl.cxx
+++ b/src/libserver/symcache/symcache_impl.cxx
@@ -369,6 +369,29 @@ auto symcache::get_item_by_id(int id, bool resolve_parent) const -> const cache_
 	return ret.get();
 }
 
+auto symcache::get_item_by_id_mut(int id, bool resolve_parent) const -> cache_item *
+{
+	if (id < 0 || id >= items_by_id.size()) {
+		msg_err_cache("internal error: requested item with id %d, when we have just %d items in the cache",
+				id, (int) items_by_id.size());
+		return nullptr;
+	}
+
+	auto &ret = items_by_id[id];
+
+	if (!ret) {
+		msg_err_cache("internal error: requested item with id %d but it is empty; qed",
+				id);
+		return nullptr;
+	}
+
+	if (resolve_parent && ret->is_virtual()) {
+		return (cache_item *)ret->get_parent(*this);
+	}
+
+	return ret.get();
+}
+
 auto symcache::get_item_by_name(std::string_view name, bool resolve_parent) const -> const cache_item *
 {
 	auto it = items_by_symbol.find(name);
diff --git a/src/libserver/symcache/symcache_internal.hxx b/src/libserver/symcache/symcache_internal.hxx
index 6a96eb547..9bce37532 100644
--- a/src/libserver/symcache/symcache_internal.hxx
+++ b/src/libserver/symcache/symcache_internal.hxx
@@ -223,6 +223,7 @@ public:
 	 * @return
 	 */
 	auto get_item_by_id(int id, bool resolve_parent) const -> const cache_item *;
+	auto get_item_by_id_mut(int id, bool resolve_parent) const -> cache_item *;
 	/**
 	 * Get an item by it's name
 	 * @param name
diff --git a/src/libserver/symcache/symcache_item.cxx b/src/libserver/symcache/symcache_item.cxx
index 091e6cbf9..1bac08495 100644
--- a/src/libserver/symcache/symcache_item.cxx
+++ b/src/libserver/symcache/symcache_item.cxx
@@ -370,7 +370,12 @@ cache_item::add_augmentation(const symcache &cache, std::string_view augmentatio
 
 	augmentations.insert(std::string(augmentation));
 
-	return known_augmentations.contains(augmentation);
+	auto ret = known_augmentations.contains(augmentation);
+
+	msg_debug_cache("added %s augmentation %s for symbol %s",
+			ret ? "known" : "unknown", augmentation.data(), symbol.data());
+
+	return ret;
 }
 
 auto
diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c
index 747f7b51e..b2333eab1 100644
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -2050,31 +2050,46 @@ lua_config_register_symbol (lua_State * L)
 				allowed_ids, forbidden_ids,
 				FALSE);
 
-		if (!isnan (score) || group) {
-			if (one_shot) {
-				nshots = 1;
-			}
+		if (ret != -1) {
+			if (!isnan(score) || group) {
+				if (one_shot) {
+					nshots = 1;
+				}
 
-			rspamd_config_add_symbol (cfg, name,
-					score, description, group, flags,
-					0, nshots);
+				rspamd_config_add_symbol(cfg, name,
+						score, description, group, flags,
+						0, nshots);
 
-			lua_pushstring (L, "groups");
-			lua_gettable (L, 2);
+				lua_pushstring(L, "groups");
+				lua_gettable(L, 2);
 
-			if (lua_istable (L, -1)) {
-				for (lua_pushnil (L); lua_next (L, -2); lua_pop (L, 1)) {
-					if (lua_isstring (L, -1)) {
-						rspamd_config_add_symbol_group (cfg, name,
-								lua_tostring (L, -1));
-					}
-					else {
-						return luaL_error (L, "invalid groups element");
+				if (lua_istable (L, -1)) {
+					for (lua_pushnil(L); lua_next(L, -2); lua_pop (L, 1)) {
+						if (lua_isstring(L, -1)) {
+							rspamd_config_add_symbol_group(cfg, name,
+									lua_tostring (L, -1));
+						}
+						else {
+							return luaL_error(L, "invalid groups element");
+						}
 					}
 				}
+
+				lua_pop (L, 1);
+			}
+
+			lua_pushstring (L, "augmentations");
+			lua_gettable (L, 2);
+
+			if (lua_type (L, -1) == LUA_TTABLE) {
+
+				for (lua_pushnil (L); lua_next (L, 2); lua_pop (L, 1)) {
+					rspamd_symcache_add_symbol_augmentation(cfg->cache, ret,
+							lua_tostring(L, -1));
+				}
 			}
 
-			lua_pop (L, 1);
+			lua_pop (L, 1); /* Table itself */
 		}
 	}
 	else {
@@ -2751,10 +2766,25 @@ lua_config_newindex (lua_State *L)
 					g_assert (name != NULL);
 					rspamd_symcache_add_condition_delayed (cfg->cache,
 							name, L, condref);
+
 				}
 				else {
 					lua_pop (L, 1);
 				}
+
+				/* Check for augmentations */
+				lua_pushstring (L, "augmentations");
+				lua_gettable (L, -2);
+
+				if (lua_type (L, -1) == LUA_TTABLE) {
+
+					for (lua_pushnil(L); lua_next(L, 2); lua_pop (L, 1)) {
+						rspamd_symcache_add_symbol_augmentation(cfg->cache, id,
+								lua_tostring(L, -1));
+					}
+				}
+
+				lua_pop (L, 1);
 			}
 
 			/*