commit 8c9a5fc: [Minor] Allow to specify flags mode when getting urls

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2020-05-18 16:56:25 +0100
URL: https://github.com/rspamd/rspamd/commit/8c9a5fc40ae7275e9101ecfe65bef6d150260af1 (HEAD -> master)

[Minor] Allow to specify flags mode when getting urls

---
 src/lua/lua_url.c | 35 +++++++++++++++++++++++++++++++----
 src/lua/lua_url.h |  4 ++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/lua/lua_url.c b/src/lua/lua_url.c
index b49146292..29f75995b 100644
--- a/src/lua/lua_url.c
+++ b/src/lua/lua_url.c
@@ -933,8 +933,18 @@ lua_tree_url_callback (gpointer key, gpointer value, gpointer ud)
 	struct rspamd_url *url = (struct rspamd_url *)value;
 	struct lua_tree_cb_data *cb = ud;
 
-	if ((url->protocol & cb->protocols_mask) == url->protocol &&
-		(url->flags == (url->flags & cb->flags_mask))) {
+	if ((url->protocol & cb->protocols_mask) == url->protocol) {
+
+		if (cb->flags_mode == url_flags_mode_include_any) {
+			if (url->flags != (url->flags & cb->flags_mask)) {
+				return;
+			}
+		}
+		else {
+			if ((url->flags & cb->flags_mask) != cb->flags_mask) {
+				return;
+			}
+		}
 
 		if (cb->skip_prob > 0) {
 			gdouble coin = rspamd_random_double_fast_seed (cb->xoroshiro_state);
@@ -966,6 +976,9 @@ lua_url_cbdata_fill (lua_State *L,
 	guint flags_mask = default_flags;
 	gboolean seen_flags = FALSE, seen_protocols = FALSE;
 
+	memset (cbd, 0, sizeof (*cbd));
+	cbd->flags_mode = url_flags_mode_include_any;
+
 	if (pos_arg_type == LUA_TBOOLEAN) {
 		protocols_mask = default_protocols;
 		if (lua_toboolean (L, 2)) {
@@ -980,6 +993,21 @@ lua_url_cbdata_fill (lua_State *L,
 			if (lua_istable (L, -1)) {
 				gint top = lua_gettop (L);
 
+				lua_getfield (L, pos, "flags_mode");
+				if (lua_isstring (L, -1)) {
+					const gchar *mode_str = lua_tostring (L, -1);
+
+					if (strcmp (mode_str, "explicit") == 0) {
+						cbd->flags_mode = url_flags_mode_include_explicit;
+						/*
+						 * Ignore default flags in this mode and include
+						 * merely flags specified by a caller
+						 */
+						flags_mask = 0;
+					}
+				}
+				lua_pop (L, 1);
+
 				for (lua_pushnil (L); lua_next (L, top); lua_pop (L, 1)) {
 					int nmask = 0;
 					const gchar *fname = lua_tostring (L, -1);
@@ -993,6 +1021,7 @@ lua_url_cbdata_fill (lua_State *L,
 						return FALSE;
 					}
 				}
+
 				seen_flags = TRUE;
 			}
 			else {
@@ -1136,8 +1165,6 @@ lua_url_cbdata_fill (lua_State *L,
 		}
 	}
 
-	memset (cbd, 0, sizeof (*cbd));
-
 	cbd->i = 1;
 	cbd->L = L;
 	cbd->max_urls = max_urls;
diff --git a/src/lua/lua_url.h b/src/lua/lua_url.h
index 0ea2186d8..c074a579a 100644
--- a/src/lua/lua_url.h
+++ b/src/lua/lua_url.h
@@ -28,6 +28,10 @@ struct lua_tree_cb_data {
 	int metatable_pos;
 	guint flags_mask;
 	guint protocols_mask;
+	enum {
+		url_flags_mode_include_any,
+		url_flags_mode_include_explicit,
+	} flags_mode;
 	gsize max_urls;
 	gdouble skip_prob;
 	guint64 xoroshiro_state[4];