[Feature] Allow to query storage about number of fuzzy hashes stored

diff --git a/src/client/rspamc.c b/src/client/rspamc.c
index 68c33472d97753ed0da2439813fd034baf02b68b..6c99f99844d627a6141500f945b2d97df190ef69 100644 (file)
--- a/src/client/rspamc.c
+++ b/src/client/rspamc.c
@@ -910,10 +910,23 @@ rspamc_stat_output (FILE *out, ucl_object_t *obj)
        rspamd_printf_gstring (out_str, "Oversized chunks: %L\n",
                ucl_object_toint (ucl_object_lookup (obj, "chunks_oversized")));
        /* Fuzzy */
-       rspamd_printf_gstring (out_str, "Fuzzy hashes stored: %L\n",
-               ucl_object_toint (ucl_object_lookup (obj, "fuzzy_stored")));
-       rspamd_printf_gstring (out_str, "Fuzzy hashes expired: %L\n",
-               ucl_object_toint (ucl_object_lookup (obj, "fuzzy_expired")));
+
+       st = ucl_object_lookup (obj, "fuzzy_hashes");
+       if (st) {
+               ucl_object_iter_t it = NULL;
+               const ucl_object_t *cur;
+               gint64 stored = 0;
+
+               while ((cur = ucl_iterate_object (st, &it, true)) != NULL) {
+                       rspamd_printf_gstring (out_str, "Fuzzy hashes in storage \"%s\": %L\n",
+                                       ucl_object_key (cur),
+                                       ucl_object_toint (cur));
+                       stored += ucl_object_toint (cur);
+               }
+
+               rspamd_printf_gstring (out_str, "Fuzzy hashes stored: %L\n",
+                               stored);
+       }
 
        st = ucl_object_lookup (obj, "fuzzy_checked");
        if (st != NULL && ucl_object_type (st) == UCL_ARRAY) {

diff --git a/src/controller.c b/src/controller.c
index 1021939aea9ed656e04767b78901c73c3923bfdc..085eda5d4a9e74acede7d34f20070f2f93af8590 100644 (file)
--- a/src/controller.c
+++ b/src/controller.c
@@ -24,6 +24,7 @@
 #include "libserver/worker_util.h"
 #include "cryptobox.h"
 #include "ottery.h"
+#include "fuzzy_storage.h"
 #include "libutil/rrd.h"
 #include "unix-std.h"
 #include "utlist.h"
@@ -103,6 +104,8 @@ const struct timeval rrd_update_time = {
 
 const guint64 rspamd_controller_ctx_magic = 0xf72697805e6941faULL;
 
+extern void fuzzy_stat_command (struct rspamd_task *task);
+
 gpointer init_controller_worker (struct rspamd_config *cfg);
 void start_controller_worker (struct rspamd_worker *worker);
 
@@ -1925,7 +1928,9 @@ rspamd_controller_stat_fin_task (void *ud)
 {
        struct rspamd_stat_cbdata *cbdata = ud;
        struct rspamd_http_connection_entry *conn_ent;
-       ucl_object_t *top;
+       ucl_object_t *top, *ar;
+       GList *fuzzy_elts, *cur;
+       struct rspamd_fuzzy_stat_entry *entry;
 
        conn_ent = cbdata->conn_ent;
        top = cbdata->top;
@@ -1937,6 +1942,23 @@ rspamd_controller_stat_fin_task (void *ud)
                ucl_object_insert_key (top, cbdata->stat, "statfiles", 0, false);
        }
 
+       fuzzy_elts = rspamd_mempool_get_variable (cbdata->task->task_pool, "fuzzy_stat");
+
+       if (fuzzy_elts) {
+               ar = ucl_object_typed_new (UCL_OBJECT);
+
+               for (cur = fuzzy_elts; cur != NULL; cur = g_list_next (cur)) {
+                       entry = cur->data;
+
+                       if (entry->name) {
+                               ucl_object_insert_key (ar, ucl_object_fromint (entry->fuzzy_cnt),
+                                               entry->name, 0, true);
+                       }
+               }
+
+               ucl_object_insert_key (top, ar, "fuzzy_hashes", 0, false);
+       }
+
        rspamd_controller_send_ucl (conn_ent, top);
 
 
@@ -2064,6 +2086,8 @@ rspamd_controller_handle_stat_common (
                rspamd_mempool_stat_reset ();
        }
 
+       fuzzy_stat_command (task);
+
        /* Now write statistics for each statfile */
        rspamd_stat_statistics (task, session->ctx->cfg, &cbdata->learned,
                        &cbdata->stat);

diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c
index d82ef342c79b8a699afe655fc849e66cd10095e1..247653c1de1b4acb65f79bd61e626cb0ec0b9f3f 100644 (file)
--- a/src/fuzzy_storage.c
+++ b/src/fuzzy_storage.c
@@ -669,6 +669,11 @@ rspamd_fuzzy_process_command (struct fuzzy_session *session)
                result = rspamd_fuzzy_backend_check (session->ctx->backend, cmd,
                                session->ctx->expire);
        }
+       else if (cmd->cmd == FUZZY_STAT) {
+               result.prob = 1.0;
+               result.value = 0;
+               result.flag = rspamd_fuzzy_backend_count (session->ctx->backend);
+       }
        else {
                if (rspamd_fuzzy_check_client (session)) {
 

diff --git a/src/fuzzy_storage.h b/src/fuzzy_storage.h
index aca6ab9525eea9fe87eee1aedebc0a446ca2e744..c23672ed89528563ecbfd9ba3885a22c71a999c9 100644 (file)
--- a/src/fuzzy_storage.h
+++ b/src/fuzzy_storage.h
@@ -13,6 +13,7 @@
 #define FUZZY_CHECK 0
 #define FUZZY_WRITE 1
 #define FUZZY_DEL 2
+#define FUZZY_STAT 3
 
 
 /**
@@ -79,4 +80,9 @@ RSPAMD_PACKED(rspamd_fuzzy_encrypted_reply) {
 
 static const guchar fuzzy_encrypted_magic[4] = {'r', 's', 'f', 'e'};
 
+struct rspamd_fuzzy_stat_entry {
+       const gchar *name;
+       guint32 fuzzy_cnt;
+};
+
 #endif

diff --git a/src/plugins/fuzzy_check.c b/src/plugins/fuzzy_check.c
index d2b0137f2b4e36b73c681f9168a8ff9e292adc8e..f775e9e091da6d4207f418f65d9d9e1fa2f3e928 100644 (file)
--- a/src/plugins/fuzzy_check.c
+++ b/src/plugins/fuzzy_check.c
@@ -1117,6 +1117,48 @@ fuzzy_cmd_from_task_meta (struct fuzzy_rule *rule,
        return io;
 }
 
+static struct fuzzy_cmd_io *
+fuzzy_cmd_stat (struct fuzzy_rule *rule,
+               int c,
+               gint flag,
+               guint32 weight,
+               rspamd_mempool_t *pool)
+{
+       struct rspamd_fuzzy_cmd *cmd;
+       struct rspamd_fuzzy_encrypted_cmd *enccmd;
+       struct fuzzy_cmd_io *io;
+
+       if (rule->peer_key) {
+               enccmd = rspamd_mempool_alloc0 (pool, sizeof (*enccmd));
+               cmd = &enccmd->cmd;
+       }
+       else {
+               cmd = rspamd_mempool_alloc0 (pool, sizeof (*cmd));
+       }
+
+       cmd->cmd = c;
+       cmd->version = RSPAMD_FUZZY_PLUGIN_VERSION;
+       cmd->shingles_count = 0;
+       cmd->tag = ottery_rand_uint32 ();
+
+       io = rspamd_mempool_alloc (pool, sizeof (*io));
+       io->flags = 0;
+       io->tag = cmd->tag;
+       io->cmd = cmd;
+
+       if (rule->peer_key) {
+               fuzzy_encrypt_cmd (rule, &enccmd->hdr, (guchar *)cmd, sizeof (*cmd));
+               io->io.iov_base = enccmd;
+               io->io.iov_len = sizeof (*enccmd);
+       }
+       else {
+               io->io.iov_base = cmd;
+               io->io.iov_len = sizeof (*cmd);
+       }
+
+       return io;
+}
+
 static void *
 fuzzy_cmd_get_cached (struct fuzzy_rule *rule,
                rspamd_mempool_t *pool,
@@ -1528,7 +1570,29 @@ fuzzy_check_try_read (struct fuzzy_client_session *session)
                while ((rep = fuzzy_process_reply (&p, &r,
                                session->commands, session->rule, &cmd)) != NULL) {
                        if (rep->prob > 0.5) {
-                               fuzzy_insert_result (session, rep, cmd, rep->flag);
+                               if (cmd->cmd == FUZZY_CHECK) {
+                                       fuzzy_insert_result (session, rep, cmd, rep->flag);
+                               }
+                               else if (cmd->cmd == FUZZY_STAT) {
+                                       /* Just set pool variable to extract it in further */
+                                       struct rspamd_fuzzy_stat_entry *pval;
+                                       GList *res;
+
+                                       pval = rspamd_mempool_alloc (task->task_pool, sizeof (*pval));
+                                       pval->fuzzy_cnt = rep->flag;
+                                       pval->name = session->rule->name;
+
+                                       res = rspamd_mempool_get_variable (task->task_pool, "fuzzy_stat");
+
+                                       if (res == NULL) {
+                                               res = g_list_append (NULL, pval);
+                                               rspamd_mempool_set_variable (task->task_pool, "fuzzy_stat",
+                                                               res, (rspamd_mempool_destruct_t)g_list_free);
+                                       }
+                                       else {
+                                               res = g_list_append (res, pval);
+                                       }
+                               }
                        }
                        else if (rep->value == 403) {
                                msg_info_task (
@@ -1936,6 +2000,15 @@ fuzzy_generate_commands (struct rspamd_task *task, struct fuzzy_rule *rule,
 
        res = g_ptr_array_new ();
 
+       if (c == FUZZY_STAT) {
+               io = fuzzy_cmd_stat (rule, c, flag, value, task->task_pool);
+               if (io) {
+                       g_ptr_array_add (res, io);
+               }
+
+               goto end;
+       }
+
        for (i = 0; i < task->text_parts->len; i ++) {
                part = g_ptr_array_index (task->text_parts, i);
 
@@ -2031,7 +2104,7 @@ fuzzy_generate_commands (struct rspamd_task *task, struct fuzzy_rule *rule,
                g_ptr_array_add (res, io);
        }
 #endif
-
+end:
        if (res->len == 0) {
                g_ptr_array_free (res, FALSE);
                return NULL;
@@ -2126,6 +2199,24 @@ fuzzy_symbol_callback (struct rspamd_task *task, void *unused)
        }
 }
 
+void
+fuzzy_stat_command (struct rspamd_task *task)
+{
+       struct fuzzy_rule *rule;
+       GList *cur;
+       GPtrArray *commands;
+
+       cur = fuzzy_module_ctx->fuzzy_rules;
+       while (cur) {
+               rule = cur->data;
+               commands = fuzzy_generate_commands (task, rule, FUZZY_STAT, 0, 0);
+               if (commands != NULL) {
+                       register_fuzzy_client_call (task, rule, commands);
+               }
+               cur = g_list_next (cur);
+       }
+}
+
 static inline gboolean
 register_fuzzy_controller_call (struct rspamd_http_connection_entry *entry,
        struct fuzzy_rule *rule,