[Feature] Implement mapping of remote flags to local flags

diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c
index 532532f16bb658bdd3a7d27ab42e985ef5813b52..7f3dd4020686dfc9ea894fa2c44c6ce839ad3201 100644 (file)
--- 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;
        }
@@ -1725,6 +1743,41 @@ err:
        return FALSE;
 }
 
+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,
@@ -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;
 }