[Feature] Allow to specify maximum number of shots for symbols

4 files changed, 42 insertions, 10 deletions

diff --git a/src/libserver/cfg_file.h b/src/libserver/cfg_file.h
index f6798c9e3..67a32c608 100644
--- a/src/libserver/cfg_file.h
+++ b/src/libserver/cfg_file.h
@@ -95,7 +95,6 @@ struct rspamd_symbols_group {
 	gboolean one_shot;
 };
 
-#define RSPAMD_SYMBOL_FLAG_ONESHOT (1 << 0)
 #define RSPAMD_SYMBOL_FLAG_IGNORE (1 << 1)
 #define RSPAMD_SYMBOL_FLAG_ONEPARAM (1 << 2)
 
@@ -111,6 +110,7 @@ struct rspamd_symbol {
 	struct rspamd_symbols_group *gr;
 	GList *groups;
 	guint flags;
+	gint nshots;
 };
 
 
@@ -313,6 +313,7 @@ struct rspamd_config {
 	gsize max_message;                              /**< maximum size for messages							*/
 	gsize max_pic_size;                             /**< maximum size for a picture to process				*/
 	gsize images_cache_size;                        /**< size of LRU cache for DCT data from images			*/
+	gint default_max_shots;                         /**< default maximum count of symbols hits permitted (-1 for unlimited) */
 
 	enum rspamd_log_type log_type;                  /**< log type											*/
 	gint log_facility;                              /**< log facility in case of syslog						*/
@@ -590,13 +591,15 @@ gboolean rspamd_init_filters (struct rspamd_config *cfg, bool reconfig);
  * @param one_shot TRUE if symbol can add its score once
  * @param rewrite_existing TRUE if we need to rewrite the existing symbol
  * @param priority use the following priority for a symbol
+ * @param nshots means maximum number of hits for a symbol in metric (-1 for unlimited)
  * @return TRUE if symbol has been inserted or FALSE if symbol already exists with higher priority
  */
 gboolean rspamd_config_add_metric_symbol (struct rspamd_config *cfg,
 		const gchar *metric,
 		const gchar *symbol, gdouble score, const gchar *description,
 		const gchar *group, guint flags,
-		guint priority);
+		guint priority,
+		gint nshots);
 
 /**
  * Sets action score for a specified metric with the specified priority
diff --git a/src/libserver/cfg_rcl.c b/src/libserver/cfg_rcl.c
index f002587fa..8ed935020 100644
--- a/src/libserver/cfg_rcl.c
+++ b/src/libserver/cfg_rcl.c
@@ -352,15 +352,23 @@ rspamd_rcl_symbol_handler (rspamd_mempool_t *pool, const ucl_object_t *obj,
 	const gchar *description = NULL;
 	gdouble score = 0.0;
 	guint priority = 1, flags = 0;
+	gint nshots;
 
 	g_assert (key != NULL);
 	metric = sd->metric;
 	g_assert (metric != NULL);
 	cfg = sd->cfg;
+	nshots = cfg->default_max_shots;
 
 	if ((elt = ucl_object_lookup (obj, "one_shot")) != NULL) {
 		if (ucl_object_toboolean (elt)) {
-			flags |= RSPAMD_SYMBOL_FLAG_ONESHOT;
+			nshots = 1;
+		}
+	}
+
+	if ((elt = ucl_object_lookup (obj, "any_shot")) != NULL) {
+		if (ucl_object_toboolean (elt)) {
+			nshots = -1;
 		}
 	}
 
@@ -376,6 +384,10 @@ rspamd_rcl_symbol_handler (rspamd_mempool_t *pool, const ucl_object_t *obj,
 		}
 	}
 
+	if ((elt = ucl_object_lookup (obj, "nshots")) != NULL) {
+		nshots = ucl_object_toint (elt);
+	}
+
 	elt = ucl_object_lookup_any (obj, "score", "weight", NULL);
 	if (elt) {
 		score = ucl_object_todouble (elt);
@@ -396,11 +408,11 @@ rspamd_rcl_symbol_handler (rspamd_mempool_t *pool, const ucl_object_t *obj,
 
 	if (sd->gr) {
 		rspamd_config_add_metric_symbol (cfg, metric->name, key, score,
-				description, sd->gr->name, flags, priority);
+				description, sd->gr->name, flags, priority, nshots);
 	}
 	else {
 		rspamd_config_add_metric_symbol (cfg, metric->name, key, score,
-				description, NULL, flags, priority);
+				description, NULL, flags, priority, nshots);
 	}
 
 	return TRUE;
@@ -1506,7 +1518,8 @@ rspamd_rcl_composite_handler (rspamd_mempool_t *pool,
 		}
 
 		rspamd_config_add_metric_symbol (cfg, metric, composite_name, score,
-				description, group, FALSE, FALSE);
+				description, group, FALSE, FALSE,
+				1);
 	}
 
 	val = ucl_object_lookup (obj, "policy");
@@ -2178,6 +2191,12 @@ rspamd_rcl_config_init (struct rspamd_config *cfg)
 			G_STRUCT_OFFSET (struct rspamd_config, compat_messages),
 			0,
 			"Use pre 1.4 style of messages in the protocol");
+	rspamd_rcl_add_default_handler (sub,
+			"max_shots",
+			rspamd_rcl_parse_struct_integer,
+			G_STRUCT_OFFSET (struct rspamd_config, default_max_shots),
+			0,
+			"Maximum number of hits per a single symbol (default: 100)");
 
 	/* Neighbours configuration */
 	rspamd_rcl_add_section_doc (&sub->subsections, "neighbours", "name",
diff --git a/src/libserver/cfg_utils.c b/src/libserver/cfg_utils.c
index 791d24fa0..2541c1aac 100644
--- a/src/libserver/cfg_utils.c
+++ b/src/libserver/cfg_utils.c
@@ -41,6 +41,7 @@
 #define DEFAULT_WORDS_DECAY 200
 #define DEFAULT_MAX_MESSAGE (50 * 1024 * 1024)
 #define DEFAULT_MAX_PIC (1 * 1024 * 1024)
+#define DEFAULT_MAX_SHOTS 100
 
 struct rspamd_ucl_map_cbdata {
 	struct rspamd_config *cfg;
@@ -179,6 +180,7 @@ rspamd_config_new (void)
 #ifdef WITH_HIREDIS
 	cfg->redis_pool = rspamd_redis_pool_init ();
 #endif
+	cfg->default_max_shots = DEFAULT_MAX_SHOTS;
 
 	REF_INIT_RETAIN (cfg, rspamd_config_free);
 
@@ -660,6 +662,12 @@ rspamd_config_post_load (struct rspamd_config *cfg,
 	cfg->clock_res = 1;
 #endif
 
+	if (cfg->one_shot_mode) {
+		msg_info_config ("enabling one shot mode (was %d max shots)",
+				cfg->default_max_shots);
+		cfg->default_max_shots = 1;
+	}
+
 	rspamd_regexp_library_init ();
 	rspamd_multipattern_library_init (cfg->hs_cache_dir,
 			cfg->libs_ctx->crypto_ctx);
@@ -1374,7 +1382,7 @@ static void
 rspamd_config_new_metric_symbol (struct rspamd_config *cfg,
 		struct rspamd_metric *metric, const gchar *symbol,
 		gdouble score, const gchar *description, const gchar *group,
-		guint flags, guint priority)
+		guint flags, guint priority, gint nshots)
 {
 	struct rspamd_symbols_group *sym_group;
 	struct rspamd_symbol *sym_def;
@@ -1391,6 +1399,7 @@ rspamd_config_new_metric_symbol (struct rspamd_config *cfg,
 	sym_def->name = rspamd_mempool_strdup (cfg->cfg_pool, symbol);
 	sym_def->priority = priority;
 	sym_def->flags = flags;
+	sym_def->nshots = nshots;
 
 	if (description) {
 		sym_def->description = rspamd_mempool_strdup (cfg->cfg_pool, description);
@@ -1433,7 +1442,7 @@ gboolean
 rspamd_config_add_metric_symbol (struct rspamd_config *cfg,
 		const gchar *metric_name, const gchar *symbol,
 		gdouble score, const gchar *description, const gchar *group,
-		guint flags, guint priority)
+		guint flags, guint priority, gint nshots)
 {
 	struct rspamd_symbol *sym_def;
 	struct rspamd_metric *metric;
@@ -1482,6 +1491,7 @@ rspamd_config_add_metric_symbol (struct rspamd_config *cfg,
 			*sym_def->weight_ptr = score;
 			sym_def->score = score;
 			sym_def->flags = flags;
+			sym_def->nshots = nshots;
 
 			if (description) {
 				sym_def->description = rspamd_mempool_strdup (cfg->cfg_pool,
@@ -1495,7 +1505,7 @@ rspamd_config_add_metric_symbol (struct rspamd_config *cfg,
 	}
 
 	rspamd_config_new_metric_symbol (cfg, metric, symbol, score, description,
-			group, flags, priority);
+			group, flags, priority, nshots);
 
 	return TRUE;
 }
diff --git a/src/libserver/dynamic_cfg.c b/src/libserver/dynamic_cfg.c
index 6e319ed36..950b9719c 100644
--- a/src/libserver/dynamic_cfg.c
+++ b/src/libserver/dynamic_cfg.c
@@ -81,7 +81,7 @@ apply_dynamic_conf (const ucl_object_t *top, struct rspamd_config *cfg)
 					 */
 					rspamd_config_add_metric_symbol (cfg, real_metric->name,
 							ucl_object_tostring (n), nscore, NULL, NULL,
-							0, priority);
+							0, priority, cfg->default_max_shots);
 				}
 				else {
 					msg_info (