[Feature] Implement mapping of remote flags to local flags

1 files changed, 63 insertions, 0 deletions

diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c
index 532532f16..7f3dd4020 100644
--- a/src/fuzzy_storage.c
+++ b/src/fuzzy_storage.c
@@ -127,6 +127,7 @@ struct rspamd_fuzzy_storage_ctx {
 	GPtrArray *mirrors;
 	const ucl_object_t *update_map;
 	const ucl_object_t *masters_map;
+	GHashTable *master_flags;
 	guint keypair_cache_size;
 	struct event_base *ev_base;
 	gint peer_fd;
@@ -933,6 +934,7 @@ rspamd_fuzzy_mirror_process_update (struct fuzzy_master_update_session *session,
 		finish_processing
 	} state = read_len;
 	GList *updates = NULL, *cur;
+	gpointer flag_ptr;
 
 	if (!rspamd_http_message_get_body (msg, NULL) || !msg->url
 			|| msg->url->len == 0) {
@@ -1058,6 +1060,22 @@ rspamd_fuzzy_mirror_process_update (struct fuzzy_master_update_session *session,
 
 	/* Insert elements to the updates from head */
 	for (cur = updates; cur != NULL; cur = g_list_next (cur)) {
+		pcmd = cur->data;
+
+		if (pcmd->is_shingle) {
+			if ((flag_ptr = g_hash_table_lookup (session->ctx->master_flags,
+					GUINT_TO_POINTER (pcmd->cmd.shingle.basic.flag))) != NULL) {
+				pcmd->cmd.shingle.basic.flag = GPOINTER_TO_UINT (flag_ptr);
+			}
+		}
+		else {
+			if ((flag_ptr = g_hash_table_lookup (session->ctx->master_flags,
+					GUINT_TO_POINTER (pcmd->cmd.normal.flag))) != NULL) {
+				pcmd->cmd.normal.flag = GPOINTER_TO_UINT (flag_ptr);
+			}
+		}
+
+
 		g_queue_push_head (session->ctx->updates_pending, cur->data);
 		cur->data = NULL;
 	}
@@ -1726,6 +1744,41 @@ err:
 }
 
 static gboolean
+fuzzy_storage_parse_master_flags (rspamd_mempool_t *pool,
+	const ucl_object_t *obj,
+	gpointer ud,
+	struct rspamd_rcl_section *section,
+	GError **err)
+{
+	const ucl_object_t *cur;
+	struct rspamd_rcl_struct_parser *pd = ud;
+	struct rspamd_fuzzy_storage_ctx *ctx;
+	ucl_object_iter_t it = NULL;
+	gulong remote_flag;
+	gint64 local_flag;
+
+	ctx = pd->user_struct;
+
+	if (ucl_object_type (obj) != UCL_OBJECT) {
+		g_set_error (err, g_quark_try_string ("fuzzy"), 100,
+				"master_flags option must be an object");
+
+		return FALSE;
+	}
+
+	while ((cur = ucl_iterate_object (obj, &it, true)) != NULL) {
+		if (rspamd_strtoul (cur->key, cur->keylen, &remote_flag) &&
+				ucl_object_toint_safe (cur, &local_flag)) {
+			g_hash_table_insert (ctx->master_flags, GUINT_TO_POINTER (remote_flag),
+					GUINT_TO_POINTER (local_flag));
+		}
+	}
+
+	return TRUE;
+}
+
+
+static gboolean
 fuzzy_parse_keypair (rspamd_mempool_t *pool,
 		const ucl_object_t *obj,
 		gpointer ud,
@@ -1824,6 +1877,7 @@ init_fuzzy (struct rspamd_config *cfg)
 	ctx->keypair_cache_size = DEFAULT_KEYPAIR_CACHE_SIZE;
 	ctx->keys = g_hash_table_new_full (fuzzy_kp_hash, fuzzy_kp_equal,
 			NULL, fuzzy_key_dtor);
+	ctx->master_flags = g_hash_table_new (g_direct_hash, g_direct_equal);
 	ctx->errors_ips = rspamd_lru_hash_new_full (1024,
 			(GDestroyNotify) rspamd_inet_address_destroy, g_free,
 			rspamd_inet_address_hash, rspamd_inet_address_equal);
@@ -1980,6 +2034,15 @@ init_fuzzy (struct rspamd_config *cfg)
 			RSPAMD_CL_FLAG_MULTIPLE,
 			"List of slave hosts");
 
+	rspamd_rcl_register_worker_option (cfg,
+			type,
+			"master_flags",
+			fuzzy_storage_parse_master_flags,
+			ctx,
+			0,
+			0,
+			"Map of flags in form master_flags = { master_flag = local_flag; ... }; ");
+
 	return ctx;
 }