[Minor] Rework headers extracting API

3 files changed, 103 insertions, 54 deletions

diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index 6c64c76d4..838e0fe7a 100644
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -208,18 +208,22 @@ struct rspamd_lua_ip * lua_check_ip (lua_State * L, gint pos);
 
 struct rspamd_lua_text * lua_check_text (lua_State * L, gint pos);
 
+enum rspamd_lua_task_header_type {
+	RSPAMD_TASK_HEADER_PUSH_SIMPLE = 0,
+	RSPAMD_TASK_HEADER_PUSH_RAW,
+	RSPAMD_TASK_HEADER_PUSH_FULL,
+	RSPAMD_TASK_HEADER_PUSH_COUNT,
+};
 
 gint rspamd_lua_push_header (lua_State *L,
-		struct rspamd_mime_header *h,
-		gboolean full,
-		gboolean raw);
+							 struct rspamd_mime_header *h,
+							 enum rspamd_lua_task_header_type how);
 /**
  * Push specific header to lua
  */
 gint rspamd_lua_push_header_array (lua_State *L,
-		GPtrArray *hdrs,
-		gboolean full,
-		gboolean raw);
+								   GPtrArray *hdrs,
+								   enum rspamd_lua_task_header_type how);
 
 /**
  * Check for task at the specified position
diff --git a/src/lua/lua_mimepart.c b/src/lua/lua_mimepart.c
index dd4e998a1..6ed06f22d 100644
--- a/src/lua/lua_mimepart.c
+++ b/src/lua/lua_mimepart.c
@@ -277,6 +277,15 @@ end
  */
 LUA_FUNCTION_DEF (mimepart, get_header_full);
 /***
+ * @method mimepart:get_header_count(name[, case_sensitive])
+ * Lightweight version if you need just a header's count
+ *  * By default headers are searched in caseless matter.
+ * @param {string} name name of header to get
+ * @param {boolean} case_sensitive case sensitiveness flag to search for a header
+ * @return {number} number of header's occurrencies or 0 if not found
+ */
+LUA_FUNCTION_DEF (mimepart, get_header_count);
+/***
  * @method mime_part:get_content()
  * Get the parsed content of part
  * @return {text} opaque text object (zero-copy if not casted to lua string)
@@ -424,6 +433,7 @@ static const struct luaL_reg mimepartlib_m[] = {
 	LUA_INTERFACE_DEF (mimepart, get_header),
 	LUA_INTERFACE_DEF (mimepart, get_header_raw),
 	LUA_INTERFACE_DEF (mimepart, get_header_full),
+	LUA_INTERFACE_DEF (mimepart, get_header_count),
 	LUA_INTERFACE_DEF (mimepart, is_image),
 	LUA_INTERFACE_DEF (mimepart, get_image),
 	LUA_INTERFACE_DEF (mimepart, is_archive),
@@ -1186,7 +1196,7 @@ lua_mimepart_get_filename (lua_State * L)
 }
 
 static gint
-lua_mimepart_get_header_common (lua_State *L, gboolean full, gboolean raw)
+lua_mimepart_get_header_common (lua_State *L, enum rspamd_lua_task_header_type how)
 {
 	struct rspamd_mime_part *part = lua_check_mimepart (L);
 	const gchar *name;
@@ -1199,7 +1209,7 @@ lua_mimepart_get_header_common (lua_State *L, gboolean full, gboolean raw)
 		ar = rspamd_message_get_header_from_hash (part->raw_headers, NULL,
 				name, FALSE);
 
-		return rspamd_lua_push_header_array (L, ar, full, raw);
+		return rspamd_lua_push_header_array (L, ar, how);
 	}
 
 	lua_pushnil (L);
@@ -1210,19 +1220,25 @@ lua_mimepart_get_header_common (lua_State *L, gboolean full, gboolean raw)
 static gint
 lua_mimepart_get_header_full (lua_State * L)
 {
-	return lua_mimepart_get_header_common (L, TRUE, TRUE);
+	return lua_mimepart_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_FULL);
 }
 
 static gint
 lua_mimepart_get_header (lua_State * L)
 {
-	return lua_mimepart_get_header_common (L, FALSE, FALSE);
+	return lua_mimepart_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_SIMPLE);
 }
 
 static gint
 lua_mimepart_get_header_raw (lua_State * L)
 {
-	return lua_mimepart_get_header_common (L, FALSE, TRUE);
+	return lua_mimepart_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_RAW);
+}
+
+static gint
+lua_mimepart_get_header_count (lua_State * L)
+{
+	return lua_mimepart_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_COUNT);
 }
 
 static gint
@@ -1418,7 +1434,7 @@ static gint
 lua_mimepart_headers_foreach (lua_State *L)
 {
 	struct rspamd_mime_part *part = lua_check_mimepart (L);
-	gboolean full = FALSE, raw = FALSE;
+	enum rspamd_lua_task_header_type how = RSPAMD_TASK_HEADER_PUSH_SIMPLE;
 	struct rspamd_lua_regexp *re = NULL;
 	GList *cur;
 	struct rspamd_mime_header *hdr;
@@ -1429,8 +1445,8 @@ lua_mimepart_headers_foreach (lua_State *L)
 			lua_pushstring (L, "full");
 			lua_gettable (L, 3);
 
-			if (lua_isboolean (L, -1)) {
-				full = lua_toboolean (L, -1);
+			if (lua_isboolean (L, -1) && lua_toboolean (L, -1)) {
+				how = RSPAMD_TASK_HEADER_PUSH_FULL;
 			}
 
 			lua_pop (L, 1);
@@ -1438,8 +1454,8 @@ lua_mimepart_headers_foreach (lua_State *L)
 			lua_pushstring (L, "raw");
 			lua_gettable (L, 3);
 
-			if (lua_isboolean (L, -1)) {
-				raw = lua_toboolean (L, -1);
+			if (lua_isboolean (L, -1) && lua_toboolean (L, -1)) {
+				how = RSPAMD_TASK_HEADER_PUSH_RAW;
 			}
 
 			lua_pop (L, 1);
@@ -1472,7 +1488,7 @@ lua_mimepart_headers_foreach (lua_State *L)
 				old_top = lua_gettop (L);
 				lua_pushvalue (L, 2);
 				lua_pushstring (L, hdr->name);
-				rspamd_lua_push_header (L, hdr, full, raw);
+				rspamd_lua_push_header (L, hdr, how);
 
 				if (lua_pcall (L, 2, LUA_MULTRET, 0) != 0) {
 					msg_err ("call to header_foreach failed: %s",
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index cffaa269c..94239e5ac 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -291,7 +291,15 @@ function check_header_delimiter_tab(task, header_name)
 end
  */
 LUA_FUNCTION_DEF (task, get_header_full);
-
+/***
+ * @method task:get_header_count(name[, case_sensitive])
+ * Lightweight version if you need just a header's count
+ *  * By default headers are searched in caseless matter.
+ * @param {string} name name of header to get
+ * @param {boolean} case_sensitive case sensitiveness flag to search for a header
+ * @return {number} number of header's occurrencies or 0 if not found
+ */
+LUA_FUNCTION_DEF (task, get_header_count);
 /***
  * @method task:get_raw_headers()
  * Get all undecoded headers of a message as a string
@@ -931,6 +939,7 @@ static const struct luaL_reg tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, get_header),
 	LUA_INTERFACE_DEF (task, get_header_raw),
 	LUA_INTERFACE_DEF (task, get_header_full),
+	LUA_INTERFACE_DEF (task, get_header_count),
 	LUA_INTERFACE_DEF (task, get_raw_headers),
 	LUA_INTERFACE_DEF (task, get_received_headers),
 	LUA_INTERFACE_DEF (task, get_queue_id),
@@ -1896,13 +1905,15 @@ lua_task_set_request_header (lua_State *L)
 	return 0;
 }
 
+
+
 gint
 rspamd_lua_push_header (lua_State *L, struct rspamd_mime_header *rh,
-		gboolean full, gboolean raw)
+						enum rspamd_lua_task_header_type how)
 {
-	const gchar *val;
 
-	if (full) {
+	switch (how) {
+	case RSPAMD_TASK_HEADER_PUSH_FULL:
 		/* Create new associated table for a header */
 		lua_createtable (L, 0, 7);
 		rspamd_lua_table_set (L, "name",	 rh->name);
@@ -1925,21 +1936,27 @@ rspamd_lua_push_header (lua_State *L, struct rspamd_mime_header *rh,
 		lua_pushstring (L, "order");
 		lua_pushnumber (L, rh->order);
 		lua_settable (L, -3);
-	}
-	else {
-		if (!raw) {
-			val = rh->decoded;
+		break;
+	case RSPAMD_TASK_HEADER_PUSH_RAW:
+		if (rh->value) {
+			lua_pushstring (L, rh->value);
 		}
 		else {
-			val = rh->value;
+			lua_pushnil (L);
 		}
-
-		if (val) {
-			lua_pushstring (L, val);
+		break;
+	case RSPAMD_TASK_HEADER_PUSH_SIMPLE:
+		if (rh->decoded) {
+			lua_pushstring (L, rh->decoded);
 		}
 		else {
 			lua_pushnil (L);
 		}
+		break;
+	case RSPAMD_TASK_HEADER_PUSH_COUNT:
+	default:
+		g_assert_not_reached ();
+		break;
 	}
 
 	return 1;
@@ -1947,39 +1964,45 @@ rspamd_lua_push_header (lua_State *L, struct rspamd_mime_header *rh,
 
 gint
 rspamd_lua_push_header_array (lua_State * L,
-		GPtrArray *ar,
-		gboolean full,
-		gboolean raw)
+							  GPtrArray *ar,
+							  enum rspamd_lua_task_header_type how)
 {
 
 	struct rspamd_mime_header *rh;
 	guint i;
 
-
 	if (ar == NULL || ar->len == 0) {
-		lua_pushnil (L);
+		if (how == RSPAMD_TASK_HEADER_PUSH_COUNT) {
+			lua_pushnumber (L, 0);
+		}
+		else {
+			lua_pushnil (L);
+		}
+
 		return 1;
 	}
 
-	if (full) {
+	if (how == RSPAMD_TASK_HEADER_PUSH_FULL) {
 		lua_createtable (L, ar->len, 0);
-	}
-
-	PTR_ARRAY_FOREACH (ar, i, rh) {
-		if (full) {
-			rspamd_lua_push_header (L, rh, full, raw);
+		PTR_ARRAY_FOREACH (ar, i, rh) {
+			rspamd_lua_push_header (L, rh, how);
 			lua_rawseti (L, -2, i + 1);
 		}
-		else {
-			return rspamd_lua_push_header (L, rh, full, raw);
-		}
+	}
+	else if (how == RSPAMD_TASK_HEADER_PUSH_COUNT) {
+		lua_pushnumber (L, ar->len);
+	}
+	else {
+		rh = g_ptr_array_index (ar, 0);
+
+		return rspamd_lua_push_header (L, rh, how);
 	}
 
 	return 1;
 }
 
 static gint
-lua_task_get_header_common (lua_State *L, gboolean full, gboolean raw)
+lua_task_get_header_common (lua_State *L, enum rspamd_lua_task_header_type how)
 {
 	gboolean strong = FALSE;
 	struct rspamd_task *task = lua_check_task (L, 1);
@@ -1995,7 +2018,7 @@ lua_task_get_header_common (lua_State *L, gboolean full, gboolean raw)
 
 		ar = rspamd_message_get_header_array (task, name, strong);
 
-		return rspamd_lua_push_header_array (L, ar, full, raw);
+		return rspamd_lua_push_header_array (L, ar, how);
 	}
 	else {
 		return luaL_error (L, "invalid arguments");
@@ -2005,19 +2028,25 @@ lua_task_get_header_common (lua_State *L, gboolean full, gboolean raw)
 static gint
 lua_task_get_header_full (lua_State * L)
 {
-	return lua_task_get_header_common (L, TRUE, TRUE);
+	return lua_task_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_FULL);
 }
 
 static gint
 lua_task_get_header (lua_State * L)
 {
-	return lua_task_get_header_common (L, FALSE, FALSE);
+	return lua_task_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_SIMPLE);
 }
 
 static gint
 lua_task_get_header_raw (lua_State * L)
 {
-	return lua_task_get_header_common (L, FALSE, TRUE);
+	return lua_task_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_RAW);
+}
+
+static gint
+lua_task_get_header_count (lua_State * L)
+{
+	return lua_task_get_header_common (L, RSPAMD_TASK_HEADER_PUSH_COUNT);
 }
 
 static gint
@@ -4668,7 +4697,7 @@ static gint
 lua_task_headers_foreach (lua_State *L)
 {
 	struct rspamd_task *task = lua_check_task (L, 1);
-	gboolean full = FALSE, raw = FALSE;
+	enum rspamd_lua_task_header_type how = RSPAMD_TASK_HEADER_PUSH_SIMPLE;
 	struct rspamd_lua_regexp *re = NULL;
 	GList *cur;
 	struct rspamd_mime_header *hdr;
@@ -4679,8 +4708,8 @@ lua_task_headers_foreach (lua_State *L)
 			lua_pushstring (L, "full");
 			lua_gettable (L, 3);
 
-			if (lua_isboolean (L, -1)) {
-				full = lua_toboolean (L, -1);
+			if (lua_isboolean (L, -1) && lua_toboolean (L, -1)) {
+				how = RSPAMD_TASK_HEADER_PUSH_FULL;
 			}
 
 			lua_pop (L, 1);
@@ -4688,8 +4717,8 @@ lua_task_headers_foreach (lua_State *L)
 			lua_pushstring (L, "raw");
 			lua_gettable (L, 3);
 
-			if (lua_isboolean (L, -1)) {
-				raw = lua_toboolean (L, -1);
+			if (lua_isboolean (L, -1) && lua_toboolean (L, -1)) {
+				how = RSPAMD_TASK_HEADER_PUSH_RAW;
 			}
 
 			lua_pop (L, 1);
@@ -4722,7 +4751,7 @@ lua_task_headers_foreach (lua_State *L)
 				old_top = lua_gettop (L);
 				lua_pushvalue (L, 2);
 				lua_pushstring (L, hdr->name);
-				rspamd_lua_push_header (L, hdr, full, raw);
+				rspamd_lua_push_header (L, hdr, how);
 
 				if (lua_pcall (L, 2, LUA_MULTRET, 0) != 0) {
 					msg_err ("call to header_foreach failed: %s",