Store max hits inside rspamd_regexp_t

3 files changed, 57 insertions, 0 deletions

diff --git a/src/libutil/regexp.c b/src/libutil/regexp.c
index 9db031109..39ff5ebf2 100644
--- a/src/libutil/regexp.c
+++ b/src/libutil/regexp.c
@@ -53,6 +53,7 @@ struct rspamd_regexp_s {
 	gpointer ud;
 	gpointer re_class;
 	guint64 cache_id;
+	guint max_hits;
 	gint flags;
 	gint pcre_flags;
 	gint ncaptures;
@@ -250,6 +251,7 @@ fin:
 	res->pattern = real_pattern;
 	res->cache_id = RSPAMD_INVALID_ID;
 	res->pcre_flags = regexp_flags;
+	res->max_hits = 0;
 
 	if (rspamd_flags & RSPAMD_REGEXP_FLAG_RAW) {
 		res->raw_re = r;
@@ -523,6 +525,26 @@ rspamd_regexp_get_ncaptures (rspamd_regexp_t *re)
 	return re->ncaptures;
 }
 
+guint
+rspamd_regexp_get_maxhits (rspamd_regexp_t *re)
+{
+	g_assert (re != NULL);
+
+	return re->max_hits;
+}
+
+guint
+rspamd_regexp_set_maxhits (rspamd_regexp_t *re, guint new_maxhits)
+{
+	guint old_hits;
+
+	g_assert (re != NULL);
+	old_hits = re->max_hits;
+	re->max_hits = new_maxhits;
+
+	return old_hits;
+}
+
 guint64
 rspamd_regexp_get_cache_id (rspamd_regexp_t *re)
 {
diff --git a/src/libutil/regexp.h b/src/libutil/regexp.h
index 8d3720f3e..0b585ceec 100644
--- a/src/libutil/regexp.h
+++ b/src/libutil/regexp.h
@@ -122,6 +122,16 @@ guint rspamd_regexp_get_pcre_flags (rspamd_regexp_t *re);
 guint rspamd_regexp_get_flags (rspamd_regexp_t *re);
 
 /**
+ * Set regexp maximum hits
+ */
+guint rspamd_regexp_get_maxhits (rspamd_regexp_t *re);
+
+/**
+ * Get regexp maximum hits
+ */
+guint rspamd_regexp_set_maxhits (rspamd_regexp_t *re, guint new_maxhits);
+
+/**
  * Returns number of backreferences in a regexp
  */
 gint rspamd_regexp_get_nbackrefs (rspamd_regexp_t *re);
diff --git a/src/lua/lua_regexp.c b/src/lua/lua_regexp.c
index b878d2a8d..d19ed83e6 100644
--- a/src/lua/lua_regexp.c
+++ b/src/lua/lua_regexp.c
@@ -42,6 +42,7 @@ LUA_FUNCTION_DEF (regexp, create_cached);
 LUA_FUNCTION_DEF (regexp, get_cached);
 LUA_FUNCTION_DEF (regexp, get_pattern);
 LUA_FUNCTION_DEF (regexp, set_limit);
+LUA_FUNCTION_DEF (regexp, set_max_hits);
 LUA_FUNCTION_DEF (regexp, search);
 LUA_FUNCTION_DEF (regexp, match);
 LUA_FUNCTION_DEF (regexp, matchn);
@@ -52,6 +53,7 @@ LUA_FUNCTION_DEF (regexp, gc);
 static const struct luaL_reg regexplib_m[] = {
 	LUA_INTERFACE_DEF (regexp, get_pattern),
 	LUA_INTERFACE_DEF (regexp, set_limit),
+	LUA_INTERFACE_DEF (regexp, set_max_hits),
 	LUA_INTERFACE_DEF (regexp, match),
 	LUA_INTERFACE_DEF (regexp, matchn),
 	LUA_INTERFACE_DEF (regexp, search),
@@ -267,6 +269,29 @@ lua_regexp_set_limit (lua_State *L)
 }
 
 /***
+ * @method re:set_max_hits(lim)
+ * Set maximum number of hits returned by a regexp
+ * @param {number} lim limit in bytes
+ */
+static int
+lua_regexp_set_max_hits (lua_State *L)
+{
+	struct rspamd_lua_regexp *re = lua_check_regexp (L);
+	guint lim;
+
+	lim = luaL_checknumber (L, 2);
+
+	if (re && re->re && !IS_DESTROYED (re)) {
+		lua_pushnumber (L, rspamd_regexp_set_maxhits (re->re, lim));
+	}
+	else {
+		lua_pushnil (L);
+	}
+
+	return 1;
+}
+
+/***
  * @method re:search(line[, raw[, capture]])
  * Search line in regular expression object. If line matches then this
  * function returns the table of captured strings. Otherwise, nil is returned.