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

5 files changed, 146 insertions, 7 deletions

diff --git a/src/client/rspamc.c b/src/client/rspamc.c
index 68c33472d..6c99f9984 100644
--- 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 1021939ae..085eda5d4 100644
--- 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 d82ef342c..247653c1d 100644
--- 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 aca6ab952..c23672ed8 100644
--- 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 d2b0137f2..f775e9e09 100644
--- 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,