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

9 files changed, 99 insertions, 33 deletions

diff --git a/src/libmime/filter.c b/src/libmime/filter.c
index 8b4b30534..2a6bbae6d 100644
--- a/src/libmime/filter.c
+++ b/src/libmime/filter.c
@@ -106,6 +106,7 @@ insert_metric_result (struct rspamd_task *task,
 	struct rspamd_symbol *sdef;
 	struct rspamd_symbols_group *gr = NULL;
 	const ucl_object_t *mobj, *sobj;
+	gint max_shots;
 
 	metric_res = rspamd_create_metric_result (task, metric->name);
 
@@ -142,14 +143,24 @@ insert_metric_result (struct rspamd_task *task,
 
 	/* Add metric score */
 	if ((s = g_hash_table_lookup (metric_res->symbols, symbol)) != NULL) {
-		if (sdef && (sdef->flags & RSPAMD_SYMBOL_FLAG_ONESHOT)) {
-			/*
-			 * For one shot symbols we do not need to add them again, so
-			 * we just force single behaviour here
-			 */
+		if (single) {
+			max_shots = 1;
+		}
+		else {
+			if (sdef) {
+				max_shots = sdef->nshots;
+			}
+			else {
+				max_shots = task->cfg->default_max_shots;
+			}
+		}
+
+		if (!single && (max_shots > 0 && (s->nshots >= max_shots))) {
 			single = TRUE;
 		}
 
+		s->nshots ++;
+
 		if (rspamd_task_add_result_option (task, s, opt)) {
 			if (!single) {
 				diff = w;
@@ -208,6 +219,7 @@ insert_metric_result (struct rspamd_task *task,
 
 		s->name = symbol;
 		s->sym = sdef;
+		s->nshots = 1;
 
 		w = rspamd_check_group_score (task, symbol, gr, gr_score, w);
 
@@ -285,7 +297,7 @@ rspamd_task_insert_result (struct rspamd_task *task,
 	const gchar *opt)
 {
 	return insert_result_common (task, symbol, flag, opt,
-			task->cfg->one_shot_mode);
+			FALSE);
 }
 
 /* Insert result as a single option */
@@ -307,7 +319,8 @@ rspamd_task_add_result_option (struct rspamd_task *task,
 
 	if (s && opt) {
 		if (s->options && !(s->sym &&
-				(s->sym->flags & RSPAMD_SYMBOL_FLAG_ONEPARAM))) {
+				(s->sym->flags & RSPAMD_SYMBOL_FLAG_ONEPARAM)) &&
+				g_hash_table_size (s->options) < task->cfg->default_max_shots) {
 			/* Append new options */
 			if (!g_hash_table_lookup (s->options, opt)) {
 				opt_cpy = rspamd_mempool_strdup (task->task_pool, opt);
diff --git a/src/libmime/filter.h b/src/libmime/filter.h
index a221960ac..920009d9c 100644
--- a/src/libmime/filter.h
+++ b/src/libmime/filter.h
@@ -22,6 +22,7 @@ struct rspamd_symbol_result {
 	GHashTable *options;                            /**< list of symbol's options				*/
 	const gchar *name;
 	struct rspamd_symbol *sym;						/**< symbol configuration					*/
+	guint nshots;
 };
 
 /**
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 (
diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c
index 8b3d995b9..e83e9a699 100644
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -1204,23 +1204,27 @@ lua_config_register_symbol (lua_State * L)
 			*description = NULL, *group = NULL;
 	double weight = 0, score = NAN;
 	gboolean one_shot = FALSE;
-	gint ret = -1, cbref = -1, type, flags = 0;
+	gint ret = -1, cbref = -1, type, flags = 0, nshots = 0;
 	gint64 parent = 0, priority = 0;
 	GError *err = NULL;
 
 	if (cfg) {
 		if (!rspamd_lua_parse_table_arguments (L, 2, &err,
 				"name=S;weigth=N;callback=F;flags=S;type=S;priority=I;parent=I;"
-				"score=D;description=S;group=S;one_shot=B",
+				"score=D;description=S;group=S;one_shot=B;nshots=I",
 				&name, &weight, &cbref, &flags_str, &type_str,
 				&priority, &parent,
-				&score, &description, &group, &one_shot)) {
+				&score, &description, &group, &one_shot, &nshots)) {
 			msg_err_config ("bad arguments: %e", err);
 			g_error_free (err);
 
 			return luaL_error (L, "invalid arguments");
 		}
 
+		if (nshots == 0) {
+			nshots = cfg->default_max_shots;
+		}
+
 		type = lua_parse_symbol_type (type_str);
 
 		if (!name && !(type & SYMBOL_TYPE_CALLBACK)) {
@@ -1244,11 +1248,11 @@ lua_config_register_symbol (lua_State * L)
 
 		if (!isnan (score)) {
 			if (one_shot) {
-				flags |= RSPAMD_SYMBOL_FLAG_ONESHOT;
+				nshots = 1;
 			}
 
 			rspamd_config_add_metric_symbol (cfg, DEFAULT_METRIC, name,
-					score, description, group, flags, (guint)priority);
+					score, description, group, flags, (guint)priority, nshots);
 		}
 	}
 	else {
@@ -1484,13 +1488,15 @@ lua_config_set_metric_symbol (lua_State * L)
 	GError *err = NULL;
 	gdouble priority = 0.0;
 	guint flags = 0;
+	gint nshots = 0;
 
 	if (cfg) {
 
 		if (lua_type (L, 2) == LUA_TTABLE) {
 			if (!rspamd_lua_parse_table_arguments (L, 2, &err,
 					"*name=S;score=N;description=S;"
-					"group=S;one_shot=B;one_param=B;metric=S;priority=N;flags=S",
+					"group=S;one_shot=B;one_param=B;metric=S;priority=N;flags=S;"
+					"nshots=I",
 					&name, &weight, &description,
 					&group, &one_shot, &one_param,
 					&metric_name, &priority, &flags_str)) {
@@ -1522,9 +1528,13 @@ lua_config_set_metric_symbol (lua_State * L)
 			metric_name = DEFAULT_METRIC;
 		}
 
+		if (nshots == 0) {
+			nshots = cfg->default_max_shots;
+		}
+
 		metric = g_hash_table_lookup (cfg->metrics, metric_name);
 		if (one_shot) {
-			flags |= RSPAMD_SYMBOL_FLAG_ONESHOT;
+			nshots = 1;
 		}
 		if (one_param) {
 			flags |= RSPAMD_SYMBOL_FLAG_ONEPARAM;
@@ -1532,7 +1542,7 @@ lua_config_set_metric_symbol (lua_State * L)
 
 		if (flags_str) {
 			if (strstr (flags_str, "one_shot") != NULL) {
-				flags |= RSPAMD_SYMBOL_FLAG_ONESHOT;
+				nshots = 1;
 			}
 			if (strstr (flags_str, "ignore") != NULL) {
 				flags |= RSPAMD_SYMBOL_FLAG_IGNORE;
@@ -1547,7 +1557,7 @@ lua_config_set_metric_symbol (lua_State * L)
 		}
 		else if (name != NULL && weight != 0) {
 			rspamd_config_add_metric_symbol (cfg, metric_name, name,
-					weight, description, group, flags, (guint)priority);
+					weight, description, group, flags, (guint)priority, nshots);
 		}
 	}
 	else {
@@ -1748,7 +1758,7 @@ lua_config_newindex (lua_State *L)
 {
 	struct rspamd_config *cfg = lua_check_config (L, 1);
 	const gchar *name;
-	gint id;
+	gint id, nshots = 0;
 	gboolean optional = FALSE;
 
 	name = luaL_checkstring (L, 2);
@@ -1893,7 +1903,7 @@ lua_config_newindex (lua_State *L)
 
 					if (lua_type (L, -1) == LUA_TBOOLEAN) {
 						if (lua_toboolean (L, -1)) {
-							flags |= RSPAMD_SYMBOL_FLAG_ONESHOT;
+							nshots = 1;
 						}
 					}
 					lua_pop (L, 1);
@@ -1913,7 +1923,7 @@ lua_config_newindex (lua_State *L)
 					 * since we are defining default values here
 					 */
 					rspamd_config_add_metric_symbol (cfg, NULL, name, score,
-							description, group, flags, 0);
+							description, group, flags, 0, nshots);
 				}
 				else {
 					lua_pop (L, 1);
diff --git a/src/plugins/dkim_check.c b/src/plugins/dkim_check.c
index aa51e7592..3abab69a4 100644
--- a/src/plugins/dkim_check.c
+++ b/src/plugins/dkim_check.c
@@ -530,7 +530,7 @@ dkim_module_config (struct rspamd_config *cfg)
 
 					rspamd_config_add_metric_symbol (cfg, DEFAULT_METRIC,
 							"DKIM_SIGN", 0.0, "DKIM signature fake symbol",
-							"dkim", RSPAMD_SYMBOL_FLAG_IGNORE, 1);
+							"dkim", RSPAMD_SYMBOL_FLAG_IGNORE, 1, 1);
 
 				}
 				else {
diff --git a/src/plugins/regexp.c b/src/plugins/regexp.c
index f2d3ab2c3..f73a31bd4 100644
--- a/src/plugins/regexp.c
+++ b/src/plugins/regexp.c
@@ -127,7 +127,7 @@ regexp_module_config (struct rspamd_config *cfg)
 	struct regexp_module_item *cur_item = NULL;
 	const ucl_object_t *sec, *value, *elt;
 	ucl_object_iter_t it = NULL;
-	gint res = TRUE, id, nre = 0, nlua = 0;
+	gint res = TRUE, id, nre = 0, nlua = 0, nshots = cfg->default_max_shots;
 
 	if (!rspamd_config_is_module_enabled (cfg, "regexp")) {
 		return TRUE;
@@ -282,10 +282,20 @@ regexp_module_config (struct rspamd_config *cfg)
 
 				if (elt) {
 					if (ucl_object_toboolean (elt)) {
-						flags |= RSPAMD_SYMBOL_FLAG_ONESHOT;
+						nshots = 1;
 					}
 				}
 
+				if ((elt = ucl_object_lookup (value, "any_shot")) != NULL) {
+					if (ucl_object_toboolean (elt)) {
+						nshots = -1;
+					}
+				}
+
+				if ((elt = ucl_object_lookup (value, "nshots")) != NULL) {
+					nshots = ucl_object_toint (elt);
+				}
+
 				elt = ucl_object_lookup (value, "one_param");
 
 				if (elt) {
@@ -301,7 +311,7 @@ regexp_module_config (struct rspamd_config *cfg)
 				}
 
 				rspamd_config_add_metric_symbol (cfg, metric, cur_item->symbol,
-						score, description, group, flags, priority);
+						score, description, group, flags, priority, nshots);
 			}
 		}
 		else {