[Feature] Add ability to deny specific fuzzy flags by default

1 files changed, 72 insertions, 3 deletions

diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c
index 3d2d7497f..7d19b6d2d 100644
--- a/src/fuzzy_storage.c
+++ b/src/fuzzy_storage.c
@@ -106,6 +106,8 @@ struct rspamd_leaky_bucket_elt {
 };
 
 static const guint64 rspamd_fuzzy_storage_magic = 0x291a3253eb1b3ea5ULL;
+KHASH_SET_INIT_INT(fuzzy_key_forbidden_ids);
+
 
 struct rspamd_fuzzy_storage_ctx {
 	guint64 magic;
@@ -168,6 +170,7 @@ struct rspamd_fuzzy_storage_ctx {
 	gint lua_pre_handler_cbref;
 	gint lua_post_handler_cbref;
 	gint lua_blacklist_cbref;
+	khash_t(fuzzy_key_forbidden_ids) *default_forbidden_ids;
 };
 
 enum fuzzy_cmd_type {
@@ -202,8 +205,6 @@ struct fuzzy_peer_request {
 	struct fuzzy_peer_cmd cmd;
 };
 
-KHASH_SET_INIT_INT(fuzzy_key_forbidden_ids);
-
 struct fuzzy_key {
 	struct rspamd_cryptobox_keypair *key;
 	struct rspamd_cryptobox_pubkey *pk;
@@ -782,6 +783,20 @@ rspamd_fuzzy_make_reply (struct rspamd_fuzzy_cmd *cmd,
 			session->reply.rep.v1.value = 0;
 		}
 
+		{
+			khiter_t k;
+
+			k = kh_get(fuzzy_key_forbidden_ids, session->ctx->default_forbidden_ids, session->reply.rep.v1.flag);
+
+			if (k != kh_end(session->ctx->default_forbidden_ids)) {
+				/* Hash is from a forbidden flag by default */
+				session->reply.rep.ts = 0;
+				session->reply.rep.v1.prob = 0.0f;
+				session->reply.rep.v1.value = 0;
+				session->reply.rep.v1.flag = 0;
+			}
+		}
+
 		if (flags & RSPAMD_FUZZY_REPLY_ENCRYPTED) {
 
 			if (session->reply.rep.v1.prob > 0 && session->key && session->key->forbidden_ids) {
@@ -792,7 +807,7 @@ rspamd_fuzzy_make_reply (struct rspamd_fuzzy_cmd *cmd,
 				if (k != kh_end (session->key->forbidden_ids)) {
 					/* Hash is from a forbidden flag for this key */
 					session->reply.rep.ts = 0;
-					session->reply.rep.v1.prob = 0.0;
+					session->reply.rep.v1.prob = 0.0f;
 					session->reply.rep.v1.value = 0;
 					session->reply.rep.v1.flag = 0;
 				}
@@ -2188,6 +2203,46 @@ rspamd_fuzzy_storage_stat (struct rspamd_main *rspamd_main,
 }
 
 static gboolean
+	fuzzy_parse_forbidden_ids (rspamd_mempool_t *pool,
+		const ucl_object_t *obj,
+		gpointer ud,
+		struct rspamd_rcl_section *section,
+		GError **err)
+{
+	struct rspamd_rcl_struct_parser *pd = (struct rspamd_rcl_struct_parser *)ud;
+	struct rspamd_fuzzy_storage_ctx *ctx;
+
+	ctx = (struct rspamd_fuzzy_storage_ctx *)pd->user_struct;
+
+	if (ucl_object_type (obj) == UCL_ARRAY) {
+		const ucl_object_t *cur;
+		ucl_object_iter_t it = NULL;
+		guint64 id;
+
+		while ((cur = ucl_object_iterate (obj, &it, true)) != NULL) {
+			if (ucl_object_toint_safe (cur, &id)) {
+				int r;
+
+				kh_put(fuzzy_key_forbidden_ids, ctx->default_forbidden_ids, id, &r);
+			}
+			else {
+				return FALSE;
+			}
+		}
+
+		return TRUE;
+	}
+	else if (ucl_object_type (obj) == UCL_INT) {
+		int r;
+		kh_put(fuzzy_key_forbidden_ids, ctx->default_forbidden_ids, ucl_object_toint (obj), &r);
+
+		return TRUE;
+	}
+
+	return FALSE;
+}
+
+static gboolean
 fuzzy_parse_keypair (rspamd_mempool_t *pool,
 		const ucl_object_t *obj,
 		gpointer ud,
@@ -2326,6 +2381,7 @@ init_fuzzy (struct rspamd_config *cfg)
 	ctx->leaky_bucket_burst = NAN;
 	ctx->leaky_bucket_rate = NAN;
 	ctx->delay = NAN;
+	ctx->default_forbidden_ids = kh_init(fuzzy_key_forbidden_ids);
 
 	rspamd_rcl_register_worker_option (cfg,
 			type,
@@ -2397,6 +2453,15 @@ init_fuzzy (struct rspamd_config *cfg)
 
 	rspamd_rcl_register_worker_option (cfg,
 			type,
+			"forbidden_ids",
+			fuzzy_parse_forbidden_ids,
+			ctx,
+			0,
+			0,
+			"Deny specific flags by default");
+
+	rspamd_rcl_register_worker_option (cfg,
+			type,
 			"keypair_cache_size",
 			rspamd_rcl_parse_struct_integer,
 			ctx,
@@ -2872,6 +2937,10 @@ start_fuzzy (struct rspamd_worker *worker)
 		luaL_unref (ctx->cfg->lua_state, LUA_REGISTRYINDEX, ctx->lua_blacklist_cbref);
 	}
 
+	if (ctx->default_forbidden_ids) {
+		kh_destroy(fuzzy_key_forbidden_ids, ctx->default_forbidden_ids);
+	}
+
 	REF_RELEASE (ctx->cfg);
 	rspamd_log_close (worker->srv->logger);
 	rspamd_unset_crash_handler (worker->srv);