From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Sat, 23 Apr 2016 14:28:40 +0000 (+0100)
Subject: [Feature] Improve getting SMTP data from lua_task
X-Git-Tag: 1.3.0~649
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=428ed3475815065da84711eba30da0b33c05c313;p=rspamd.git

[Feature] Improve getting SMTP data from lua_task
---

diff --git a/src/libserver/task.h b/src/libserver/task.h
index 40ece993f..1f4e8eeea 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -111,6 +111,9 @@ enum rspamd_task_stage {
 #define RSPAMD_TASK_IS_CLASSIFIED(task) (((task)->processed_stages & RSPAMD_TASK_STAGE_CLASSIFIERS))
 #define RSPAMD_TASK_IS_EMPTY(task) (((task)->flags & RSPAMD_TASK_FLAG_EMPTY))
 
+struct rspamd_email_address;
+
+
 /**
  * Worker task structure
  */
@@ -148,10 +151,17 @@ struct rspamd_task {
 	GHashTable *results;							/**< hash table of metric_result indexed by
 													 *    metric's name									*/
 	GPtrArray *tokens;								/**< statistics tokens */
-	InternetAddressList *rcpt_mime;					/**< list of all recipients							*/
-	InternetAddressList *rcpt_envelope;				/**< list of all recipients							*/
+#if 0
+	GPtrArray *rcpt_mime;							/**< list of all recipients (rspamd_email_address)	*/
+	GPtrArray *rcpt_envelope;						/**< list of all recipients	(rspamd_email_address)	*/
+	struct rspamd_email_address *from_mime;
+	struct rspamd_email_address *from_envelope;
+#else
+	InternetAddressList *rcpt_mime;
+	InternetAddressList *rcpt_envelope;
 	InternetAddressList *from_mime;
 	InternetAddressList *from_envelope;
+#endif
 
 	GList *messages;								/**< list of messages that would be reported		*/
 	struct rspamd_re_runtime *re_rt;				/**< regexp runtime									*/
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 6d86d316e..83a039813 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -22,6 +22,7 @@
 #include "images.h"
 #include "cfg_file.h"
 #include "utlist.h"
+#include "xxhash.h"
 
 /***
  * @module rspamd_task
@@ -1453,37 +1454,71 @@ lua_task_get_dns_req (lua_State *L)
 	return 1;
 }
 
+enum rspamd_address_type {
+	RSPAMD_ADDRESS_ANY = 0,
+	RSPAMD_ADDRESS_SMTP = 1,
+	RSPAMD_ADDRESS_MIME = 2,
+	RSPAMD_ADDRESS_RAW_ANY = 3,
+	RSPAMD_ADDRESS_RAW_SMTP = 4,
+	RSPAMD_ADDRESS_RAW_MIME = 5,
+	RSPAMD_ADDRESS_MAX
+};
+
 /*
  * Convert element at the specified position to the type
  * for get_from/get_recipients
  */
-static gint
+static enum rspamd_address_type
 lua_task_str_to_get_type (lua_State *L, gint pos)
 {
 	const gchar *type = NULL;
+	gint ret = RSPAMD_ADDRESS_ANY;
+	guint64 h;
+	gsize sz;
 
 	/* Get what value */
 
 	if (lua_type (L, pos) == LUA_TNUMBER) {
-		return lua_tonumber (L, pos);
+		ret = lua_tonumber (L, pos);
+
+		if (ret >= RSPAMD_ADDRESS_ANY && ret < RSPAMD_ADDRESS_MAX) {
+			return ret;
+		}
+
+		return RSPAMD_ADDRESS_ANY;
 	}
 	else if (lua_type (L, pos) == LUA_TSTRING) {
-		type = lua_tostring (L, pos);
-
-		if (type) {
-			if (strcmp (type, "mime") == 0) {
-				return 2;
-			}
-			else if (strcmp (type, "any") == 0) {
-				return 0;
-			}
-			else if (strcmp (type, "smtp") == 0 || strcmp (type, "envelope") == 0) {
-				return 1;
+		type = lua_tolstring (L, pos, &sz);
+
+		if (type && sz > 0) {
+			h = XXH64 (type, sz, 0xdeadbabe);
+
+			switch (h) {
+			case 0xDA081341FB600389ULL: /* mime */
+				ret = RSPAMD_ADDRESS_MIME;
+				break;
+			case 0xEEC8A7832F8C43ACULL: /* any */
+				ret = RSPAMD_ADDRESS_ANY;
+				break;
+			case 0x472274D5193B2A80ULL: /* smtp */
+			case 0xEFE0F586CC9F14A9ULL: /* envelope */
+				ret = RSPAMD_ADDRESS_SMTP;
+				break;
+			case 0x9DA887501690DE20ULL: /* raw_mime */
+				ret = RSPAMD_ADDRESS_RAW_MIME;
+				break;
+			case 0x6B54FE02DEB595A4ULL: /* raw_smtp */
+			case 0xE0E596C861777B02ULL: /* raw_envelope */
+				ret = RSPAMD_ADDRESS_RAW_SMTP;
+				break;
+			case 0x2C49DBE3A10A0197ULL: /* raw_any */
+				ret = RSPAMD_ADDRESS_RAW_ANY;
+				break;
 			}
 		}
 	}
 
-	return 0;
+	return ret;
 }
 
 static gint
@@ -1500,15 +1535,15 @@ lua_task_get_recipients (lua_State *L)
 		}
 
 		switch (what) {
-		case 1:
+		case RSPAMD_ADDRESS_SMTP:
 			/* Here we check merely envelope rcpt */
 			addrs = task->rcpt_envelope;
 			break;
-		case 2:
+		case RSPAMD_ADDRESS_MIME:
 			/* Here we check merely mime rcpt */
 			addrs = task->rcpt_mime;
 			break;
-		case 0:
+		case RSPAMD_ADDRESS_ANY:
 		default:
 			if (task->rcpt_envelope) {
 				addrs = task->rcpt_envelope;