[Project] Implement settings processing + some neats

6 files changed, 261 insertions, 11 deletions

diff --git a/src/libserver/symcache/symcache_c.cxx b/src/libserver/symcache/symcache_c.cxx
index bdef8b162..df75bf2c8 100644
--- a/src/libserver/symcache/symcache_c.cxx
+++ b/src/libserver/symcache/symcache_c.cxx
@@ -17,12 +17,14 @@
 #include "symcache_internal.hxx"
 #include "symcache_periodic.hxx"
 #include "symcache_item.hxx"
+#include "symcache_runtime.hxx"
 
 /**
  * C API for symcache
  */
 
 #define C_API_SYMCACHE(ptr) (reinterpret_cast<rspamd::symcache::symcache *>(ptr))
+#define C_API_SYMCACHE_RUNTIME(ptr) (reinterpret_cast<rspamd::symcache::symcache_runtime *>(ptr))
 #define C_API_SYMCACHE_ITEM(ptr) (reinterpret_cast<rspamd::symcache::cache_item *>(ptr))
 
 void
@@ -242,4 +244,45 @@ rspamd_symcache_foreach(struct rspamd_symcache *cache,
 	real_cache->symbols_foreach([&](const rspamd::symcache::cache_item* item) {
 		func((struct rspamd_symcache_item *)item, ud);
 	});
+}
+
+void rspamd_symcache_disable_all_symbols (struct rspamd_task *task,
+										  struct rspamd_symcache *_cache,
+										  guint skip_mask)
+{
+	auto *cache_runtime = C_API_SYMCACHE_RUNTIME(task->symcache_runtime);
+
+	cache_runtime->disable_all_symbols(skip_mask);
+}
+
+gboolean
+rspamd_symcache_disable_symbol (struct rspamd_task *task,
+								struct rspamd_symcache *cache,
+								const gchar *symbol)
+{
+	auto *cache_runtime = C_API_SYMCACHE_RUNTIME(task->symcache_runtime);
+	auto *real_cache = C_API_SYMCACHE(cache);
+
+	return cache_runtime->disable_symbol(task, *real_cache, symbol);
+}
+
+gboolean
+rspamd_symcache_enable_symbol (struct rspamd_task *task,
+								struct rspamd_symcache *cache,
+								const gchar *symbol)
+{
+	auto *cache_runtime = C_API_SYMCACHE_RUNTIME(task->symcache_runtime);
+	auto *real_cache = C_API_SYMCACHE(cache);
+
+	return cache_runtime->enable_symbol(task, *real_cache, symbol);
+}
+
+gboolean
+rspamd_symcache_process_settings (struct rspamd_task *task,
+								  struct rspamd_symcache *cache)
+{
+	auto *cache_runtime = C_API_SYMCACHE_RUNTIME(task->symcache_runtime);
+	auto *real_cache = C_API_SYMCACHE(cache);
+
+	return cache_runtime->process_settings(task, *real_cache);
 }
\ No newline at end of file
diff --git a/src/libserver/symcache/symcache_impl.cxx b/src/libserver/symcache/symcache_impl.cxx
index 535fe57b9..99aae3df8 100644
--- a/src/libserver/symcache/symcache_impl.cxx
+++ b/src/libserver/symcache/symcache_impl.cxx
@@ -542,6 +542,13 @@ auto symcache::resort() -> void
 	};
 
 	std::stable_sort(std::begin(ord->d), std::end(ord->d), cache_order_cmp);
+	/* After sorting is done, we can assign all elements in the by_symbol hash */
+	for (auto i = 0; i < ord->size(); i ++) {
+		const auto &it = ord->d[i];
+		ord->by_symbol[it->get_name()] = i;
+		ord->by_cache_id[it->id] = i;
+	}
+	/* Finally set the current order */
 	std::swap(ord, items_by_order);
 }
 
diff --git a/src/libserver/symcache/symcache_internal.hxx b/src/libserver/symcache/symcache_internal.hxx
index 0862476f6..2c3bfe739 100644
--- a/src/libserver/symcache/symcache_internal.hxx
+++ b/src/libserver/symcache/symcache_internal.hxx
@@ -82,11 +82,19 @@ using cache_item_ptr = std::shared_ptr<cache_item>;
 
 struct order_generation {
 	std::vector<cache_item_ptr> d;
+	/* Mapping from symbol name to the position in the order array */
+	robin_hood::unordered_flat_map<std::string_view, unsigned int> by_symbol;
+	/* Mapping from symbol id to the position in the order array */
+	robin_hood::unordered_flat_map<unsigned int, unsigned int> by_cache_id;
 	unsigned int generation_id;
 
 	explicit order_generation(std::size_t nelts, unsigned id) : generation_id(id) {
 		d.reserve(nelts);
+		by_symbol.reserve(nelts);
+		by_cache_id.reserve(nelts);
 	}
+
+	auto size() const -> auto { return d.size(); }
 };
 
 using order_generation_ptr = std::shared_ptr<order_generation>;
diff --git a/src/libserver/symcache/symcache_runtime.cxx b/src/libserver/symcache/symcache_runtime.cxx
index b1b5f5b99..55a0a62c4 100644
--- a/src/libserver/symcache/symcache_runtime.cxx
+++ b/src/libserver/symcache/symcache_runtime.cxx
@@ -17,6 +17,7 @@
 #include "symcache_internal.hxx"
 #include "symcache_item.hxx"
 #include "symcache_runtime.hxx"
+#include "libutil/cxx/util.hxx"
 
 namespace rspamd::symcache {
 
@@ -28,19 +29,20 @@ constexpr static const auto PROFILE_MESSAGE_SIZE_THRESHOLD = 1024ul * 1024 * 2;
 constexpr static const auto PROFILE_PROBABILITY = 0.01;
 
 auto
-cache_savepoint::create_savepoint(struct rspamd_task *task, symcache &cache) -> cache_savepoint *
+symcache_runtime::create_savepoint(struct rspamd_task *task, symcache &cache) -> symcache_runtime *
 {
-	struct cache_savepoint *checkpoint;
+	struct symcache_runtime *checkpoint;
 
 	cache.maybe_resort();
 
-	checkpoint = (struct cache_savepoint *) rspamd_mempool_alloc0 (task->task_pool,
+	auto &&cur_order = cache.get_cache_order();
+	checkpoint = (struct symcache_runtime *) rspamd_mempool_alloc0 (task->task_pool,
 			sizeof(*checkpoint) +
-			sizeof(struct cache_dynamic_item) * cache.get_items_count());
+			sizeof(struct cache_dynamic_item) * cur_order->size());
 
 	checkpoint->order = cache.get_cache_order();
 	rspamd_mempool_add_destructor(task->task_pool,
-			cache_savepoint::savepoint_dtor, checkpoint);
+			symcache_runtime::savepoint_dtor, checkpoint);
 
 	/* Calculate profile probability */
 	ev_now_update_if_cheap(task->event_loop);
@@ -55,10 +57,168 @@ cache_savepoint::create_savepoint(struct rspamd_task *task, symcache &cache) ->
 		cache.set_last_profile(now);
 	}
 
-	task->checkpoint = (void *)checkpoint;
+	task->symcache_runtime = (void *) checkpoint;
 
 	return checkpoint;
 }
 
+auto
+symcache_runtime::process_settings(struct rspamd_task *task, const symcache &cache) -> bool
+{
+	if (!task->settings) {
+		msg_err_task("`process_settings` is called with no settings");
+		return false;
+	}
+
+	const auto *wl = ucl_object_lookup(task->settings, "whitelist");
+
+	if (wl != nullptr) {
+		msg_info_task("task is whitelisted");
+		task->flags |= RSPAMD_TASK_FLAG_SKIP;
+		return true;
+	}
+
+	auto already_disabled = false;
+
+	auto process_group = [&](const ucl_object_t *gr_obj, auto functor) -> void {
+		ucl_object_iter_t it = nullptr;
+		const ucl_object_t *cur;
+
+		if (gr_obj) {
+			while ((cur = ucl_iterate_object(gr_obj, &it, true)) != nullptr) {
+				if (ucl_object_type(cur) == UCL_STRING) {
+					auto *gr = (struct rspamd_symbols_group *)
+							g_hash_table_lookup(task->cfg->groups,
+									ucl_object_tostring(cur));
+
+					if (gr) {
+						GHashTableIter gr_it;
+						void *k, *v;
+						g_hash_table_iter_init(&gr_it, gr->symbols);
+
+						while (g_hash_table_iter_next(&gr_it, &k, &v)) {
+							functor((const char*)k);
+						}
+					}
+				}
+			}
+		}
+	};
+
+	ucl_object_iter_t it = nullptr;
+	const ucl_object_t *cur;
+
+	const auto *enabled = ucl_object_lookup(task->settings, "symbols_enabled");
+
+	if (enabled) {
+		/* Disable all symbols but selected */
+		disable_all_symbols(SYMBOL_TYPE_EXPLICIT_DISABLE);
+		already_disabled = true;
+		it = nullptr;
+
+		while ((cur = ucl_iterate_object(enabled, &it, true)) != nullptr) {
+			enable_symbol(task, cache,ucl_object_tostring(cur));
+		}
+	}
+
+
+	/* Enable groups of symbols */
+	enabled = ucl_object_lookup(task->settings, "groups_enabled");
+	if (enabled && !already_disabled) {
+		disable_all_symbols(SYMBOL_TYPE_EXPLICIT_DISABLE);
+	}
+	process_group(enabled, [&](const char *sym) {
+		enable_symbol(task, cache, sym);
+	});
+
+	const auto *disabled = ucl_object_lookup(task->settings, "symbols_disabled");
+
+
+	if (disabled) {
+		it = nullptr;
+
+		while ((cur = ucl_iterate_object (disabled, &it, true)) != nullptr) {
+			disable_symbol(task, cache,ucl_object_tostring(cur));
+		}
+	}
+
+	/* Disable groups of symbols */
+	disabled = ucl_object_lookup(task->settings, "groups_disabled");
+	process_group(disabled, [&](const char *sym) {
+		disable_symbol(task, cache, sym);
+	});
+
+	return false;
+}
+
+auto symcache_runtime::disable_all_symbols(int skip_mask) -> void
+{
+	for (auto i = 0; i < order->size(); i ++) {
+		auto *dyn_item = &dynamic_items[i];
+		const auto &item = order->d[i];
+
+		if (!(item->get_flags() & skip_mask)) {
+			dyn_item->finished = true;
+			dyn_item->started = true;
+		}
+	}
+}
+
+auto
+symcache_runtime::disable_symbol(struct rspamd_task *task, const symcache &cache, std::string_view name) -> bool
+{
+	const auto *item = cache.get_item_by_name(name, true);
+
+	if (item != nullptr) {
+
+		auto our_id_maybe = rspamd::find_map(order->by_cache_id, item->id);
+
+		if (our_id_maybe) {
+			auto *dyn_item = &dynamic_items[our_id_maybe.value()];
+			dyn_item->finished = true;
+			dyn_item->started = true;
+			msg_debug_cache_task("disable execution of %s", name.data());
+
+			return true;
+		}
+		else {
+			msg_debug_cache_task("cannot disable %s: id not found %d", name.data(), item->id);
+		}
+	}
+	else {
+		msg_debug_cache_task("cannot disable %s: symbol not found", name.data());
+	}
+
+	return false;
+}
+
+auto
+symcache_runtime::enable_symbol(struct rspamd_task *task, const symcache &cache, std::string_view name) -> bool
+{
+	const auto *item = cache.get_item_by_name(name, true);
+
+	if (item != nullptr) {
+
+		auto our_id_maybe = rspamd::find_map(order->by_cache_id, item->id);
+
+		if (our_id_maybe) {
+			auto *dyn_item = &dynamic_items[our_id_maybe.value()];
+			dyn_item->finished = false;
+			dyn_item->started = false;
+			msg_debug_cache_task("enable execution of %s", name.data());
+
+			return true;
+		}
+		else {
+			msg_debug_cache_task("cannot enable %s: id not found %d", name.data(), item->id);
+		}
+	}
+	else {
+		msg_debug_cache_task("cannot enable %s: symbol not found", name.data());
+	}
+
+	return false;
+}
+
 }
 
diff --git a/src/libserver/symcache/symcache_runtime.hxx b/src/libserver/symcache/symcache_runtime.hxx
index 4d43a7015..78fc3bf7c 100644
--- a/src/libserver/symcache/symcache_runtime.hxx
+++ b/src/libserver/symcache/symcache_runtime.hxx
@@ -44,7 +44,7 @@ struct cache_dynamic_item {
 static_assert(sizeof(cache_dynamic_item) == sizeof(std::uint64_t));
 static_assert(std::is_trivial_v<cache_dynamic_item>);
 
-class cache_savepoint {
+class symcache_runtime {
 	unsigned items_inflight;
 	bool profile;
 	bool has_slow;
@@ -59,16 +59,48 @@ class cache_savepoint {
 	/* Dynamically expanded as needed */
 	struct cache_dynamic_item dynamic_items[];
 	/* We allocate this structure merely in memory pool, so destructor is absent */
-	~cache_savepoint() = delete;
+	~symcache_runtime() = delete;
 	/* Dropper for a shared ownership */
 	static auto savepoint_dtor(void *ptr) -> void {
-		auto *real_savepoint = (cache_savepoint *)ptr;
+		auto *real_savepoint = (symcache_runtime *)ptr;
 
 		/* Drop shared ownership */
 		real_savepoint->order.reset();
 	}
 public:
-	static auto create_savepoint(struct rspamd_task *task, symcache &cache) -> cache_savepoint *;
+	/**
+	 * Creates a cache runtime using task mempool
+	 * @param task
+	 * @param cache
+	 * @return
+	 */
+	static auto create_savepoint(struct rspamd_task *task, symcache &cache) -> symcache_runtime *;
+	/**
+	 * Process task settings
+	 * @param task
+	 * @return
+	 */
+	auto process_settings(struct rspamd_task *task, const symcache &cache) -> bool;
+
+	/**
+	 * Disable all symbols but not touching ones that are in the specific mask
+	 * @param skip_mask
+	 */
+	auto disable_all_symbols(int skip_mask) -> void;
+
+	/**
+	 * Disable a symbol (or it's parent)
+	 * @param name
+	 * @return
+	 */
+	auto disable_symbol(struct rspamd_task *task, const symcache &cache, std::string_view name) -> bool;
+
+	/**
+	 * Enable a symbol (or it's parent)
+	 * @param name
+	 * @return
+	 */
+	auto enable_symbol(struct rspamd_task *task, const symcache &cache, std::string_view name) -> bool;
 };
 
 
diff --git a/src/libserver/task.h b/src/libserver/task.h
index 9c9db32e9..af9a2aa4f 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -207,7 +207,7 @@ struct rspamd_task {
 	struct ev_timer timeout_ev;                        /**< Global task timeout							*/
 	struct ev_io guard_ev;                            /**< Event for input sanity guard 					*/
 
-	gpointer checkpoint;                            /**< Opaque checkpoint data							*/
+	gpointer symcache_runtime;                            /**< Opaque checkpoint data							*/
 	ucl_object_t *settings;                            /**< Settings applied to task						*/
 	struct rspamd_config_settings_elt *settings_elt;    /**< preprocessed settings id elt				*/