* Add ability to set specific symbols for different fuzzy flags

* Fix storage of fuzzy flags

diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c
index 5f515862060743ec17e05199132e7eff70a110d3..0f1d47ac7e81b6fff86e6e97451b29f0baa4c922 100644 (file)
--- a/src/fuzzy_storage.c
+++ b/src/fuzzy_storage.c
@@ -339,7 +339,7 @@ read_hashes_file (struct rspamd_worker *wrk)
        return TRUE;
 }
 
-static inline int
+static inline struct rspamd_fuzzy_node *
 check_hash_node (GQueue *hash, fuzzy_hash_t *s, int update_value)
 {
        GList                          *cur;
@@ -358,7 +358,7 @@ check_hash_node (GQueue *hash, fuzzy_hash_t *s, int update_value)
                                if (update_value) {
                                        h->value += update_value;
                                }
-                               return h->value;
+                               return h;
                        }
                }
        }
@@ -373,7 +373,7 @@ check_hash_node (GQueue *hash, fuzzy_hash_t *s, int update_value)
                                msg_info ("new hash weight: %d", h->value);
                                h->value += update_value;
                        }
-                       return h->value;
+                       return h;
                }
                cur = g_list_next (cur);
        }
@@ -392,7 +392,7 @@ check_hash_node (GQueue *hash, fuzzy_hash_t *s, int update_value)
                                g_queue_push_head_link (frequent, cur);
                                msg_info ("moved hash to frequent list");
                        }
-                       return h->value;
+                       return h;
                }
                cur = g_list_next (cur);
        }
@@ -400,13 +400,14 @@ check_hash_node (GQueue *hash, fuzzy_hash_t *s, int update_value)
        }
 #endif
 
-       return 0;
+       return NULL;
 }
 
 static                          int
-process_check_command (struct fuzzy_cmd *cmd)
+process_check_command (struct fuzzy_cmd *cmd, int *flag)
 {
        fuzzy_hash_t                    s;
+       struct rspamd_fuzzy_node       *h;
 
        if (!bloom_check (bf, cmd->hash)) {
                return 0;
@@ -415,7 +416,15 @@ process_check_command (struct fuzzy_cmd *cmd)
        memcpy (s.hash_pipe, cmd->hash, sizeof (s.hash_pipe));
        s.block_size = cmd->blocksize;
 
-       return check_hash_node (hashes[cmd->blocksize % BUCKETS], &s, 0);
+       h = check_hash_node (hashes[cmd->blocksize % BUCKETS], &s, 0);
+
+       if (h == NULL) {
+               return 0;
+       }
+       else {
+               *flag = h->flag;
+               return h->value;
+       }
 }
 
 static                          gboolean
@@ -426,7 +435,7 @@ update_hash (struct fuzzy_cmd *cmd)
        memcpy (s.hash_pipe, cmd->hash, sizeof (s.hash_pipe));
        s.block_size = cmd->blocksize;
 
-       return check_hash_node (hashes[cmd->blocksize % BUCKETS], &s, cmd->value);
+       return check_hash_node (hashes[cmd->blocksize % BUCKETS], &s, cmd->value) != NULL;
 }
 
 static                          gboolean
@@ -448,6 +457,7 @@ process_write_command (struct fuzzy_cmd *cmd)
        h->h.block_size = cmd->blocksize;
        h->time = (uint64_t) time (NULL);
        h->value = cmd->value;
+       h->flag = cmd->flag;
 #ifdef WITH_JUDY
        if (use_judy) {
                pvalue = JudySLIns (&jtree, h->h.hash_pipe, PJE0);
@@ -562,13 +572,13 @@ else {                                                                                                                                                                                    \
 static void
 process_fuzzy_command (struct fuzzy_session *session)
 {
-       int r;
+       int r, flag = 0;
        char buf[64];
 
        switch (session->cmd.cmd) {
        case FUZZY_CHECK:
-               if ((r = process_check_command (&session->cmd))) {
-                       r = snprintf (buf, sizeof (buf), "OK %d" CRLF, r);
+               if ((r = process_check_command (&session->cmd, &flag))) {
+                       r = snprintf (buf, sizeof (buf), "OK %d %d" CRLF, r, flag);
                        if (sendto (session->fd, buf, r, 0, (struct sockaddr *)&session->sa, session->salen) == -1) {
                                msg_err ("error while writing reply: %s", strerror (errno));
                        }

diff --git a/src/plugins/fuzzy_check.c b/src/plugins/fuzzy_check.c
index 1b3f386172bf3572e487510fd4bd01857815ddf3..279d97394addff5e419bd606418eb237dbb5011e 100644 (file)
--- a/src/plugins/fuzzy_check.c
+++ b/src/plugins/fuzzy_check.c
@@ -71,6 +71,7 @@ struct fuzzy_ctx {
        double                          max_score;
        uint32_t                        min_hash_len;
        radix_tree_t                   *whitelist;
+       GHashTable                     *flags;
 };
 
 struct fuzzy_client_session {
@@ -104,6 +105,37 @@ static void                     fuzzy_symbol_callback (struct worker_task *task,
 static void                     fuzzy_add_handler (char **args, struct controller_session *session);
 static void                     fuzzy_delete_handler (char **args, struct controller_session *session);
 
+/* Flags string is in format <numeric_flag>:<SYMBOL>[, <numeric_flag>:<SYMBOL>...] */
+static void
+parse_flags_string (char *str)
+{
+       char                          **strvec, *p, *item, *err_str;
+       int                             num, i, flag;
+       
+       strvec = g_strsplit (str, ", ;", 0);
+       num = g_strv_length (strvec);
+
+       for (i = 0; i < num; i ++) {
+               item = strvec[i];
+               if ((p = strchr (item, ':')) != NULL) {
+                       *p = '\0';
+                       p ++;
+                       /* Now in p we have name of symbol and in item we have its number */
+                       errno = 0;
+                       flag = strtol (item, &err_str, 10);
+                       if (errno != 0 || (err_str && *err_str != '\0')) {
+                               msg_info ("cannot parse flag %s: %s", item, strerror (errno));
+                       }
+                       else {
+                               /* Add flag to hash table */
+                               g_hash_table_insert (fuzzy_module_ctx->flags, GINT_TO_POINTER(flag), memory_pool_strdup (fuzzy_module_ctx->fuzzy_pool, p));
+                       }
+               }
+       }
+
+       g_strfreev (strvec);
+}
+
 static void
 parse_servers_string (char *str)
 {
@@ -193,6 +225,7 @@ fuzzy_check_module_init (struct config_file *cfg, struct module_ctx **ctx)
        fuzzy_module_ctx->fuzzy_pool = memory_pool_new (memory_pool_get_size ());
        fuzzy_module_ctx->servers = NULL;
        fuzzy_module_ctx->servers_num = 0;
+       fuzzy_module_ctx->flags = g_hash_table_new (g_int_hash, g_int_equal);
 
        *ctx = (struct module_ctx *)fuzzy_module_ctx;
 
@@ -243,10 +276,12 @@ fuzzy_check_module_config (struct config_file *cfg)
                fuzzy_module_ctx->whitelist = NULL;
        }
 
-
        if ((value = get_module_opt (cfg, "fuzzy_check", "servers")) != NULL) {
                parse_servers_string (value);
        }
+       if ((value = get_module_opt (cfg, "fuzzy_check", "flags")) != NULL) {
+               parse_flags_string (value);
+       }
 
        metric = g_hash_table_lookup (cfg->metrics, fuzzy_module_ctx->metric);
        if (metric == NULL) {
@@ -287,7 +322,8 @@ fuzzy_check_module_reconfig (struct config_file *cfg)
        fuzzy_module_ctx->servers = NULL;
        fuzzy_module_ctx->servers_num = 0;
        fuzzy_module_ctx->fuzzy_pool = memory_pool_new (memory_pool_get_size ());
-
+       
+       g_hash_table_remove_all (fuzzy_module_ctx->flags);
        return fuzzy_check_module_config (cfg);
 }
 
@@ -313,7 +349,7 @@ fuzzy_io_callback (int fd, short what, void *arg)
 {
        struct fuzzy_client_session    *session = arg;
        struct fuzzy_cmd                cmd;
-       char                            buf[62], *err_str;
+       char                            buf[62], *err_str, *symbol;
        int                             value = 0, flag = 0, r;
        double                          nval;
 
@@ -346,8 +382,13 @@ fuzzy_io_callback (int fd, short what, void *arg)
                        }
                        *err_str = '\0';
                        nval = fuzzy_normalize (value);
+                       /* Get symbol by flag */
+                       if ((symbol = g_hash_table_lookup (fuzzy_module_ctx->flags, GINT_TO_POINTER (flag))) == NULL) {
+                               /* Default symbol */
+                               symbol = fuzzy_module_ctx->symbol;
+                       }
                        snprintf (buf, sizeof (buf), "%d: %d / %.2f", flag, value, nval);
-                       insert_result (session->task, fuzzy_module_ctx->metric, fuzzy_module_ctx->symbol, nval, g_list_prepend (NULL, 
+                       insert_result (session->task, fuzzy_module_ctx->metric, symbol, nval, g_list_prepend (NULL, 
                                                memory_pool_strdup (session->task->task_pool, buf)));
                }
                goto ok;