[Feature] Improve getting SMTP data from lua_task

diff --git a/src/libserver/task.h b/src/libserver/task.h
index 40ece993f10f4f0e8703db886658302802e5503f..1f4e8eeeaabc735fb5774acbad8fc331a0515902 100644 (file)
--- 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 6d86d316e44cebb274352d2613d97a656f73e183..83a039813ac13f42d4420719e861710a06a6f298 100644 (file)
--- 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;