[Fix] Fix multiple request headers structure

diff --git a/src/libmime/message.c b/src/libmime/message.c
index 87a0123fcb981d075b3ca1e326a48efabacd49a5..a3b673f28452ea2dacb96dcc9e8c276ba50ecfa8 100644 (file)
--- a/src/libmime/message.c
+++ b/src/libmime/message.c
@@ -1271,9 +1271,7 @@ rspamd_message_from_data (struct rspamd_task *task, GByteArray *data,
                }
        }
 
-       srch.begin = "Content-Type";
-       srch.len = sizeof ("Content-Type") - 1;
-       tok = g_hash_table_lookup (task->request_headers, &srch);
+       tok = rspamd_task_get_request_header (task, "Content-Type");
 
        if (tok) {
                /* We have Content-Type defined */

diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index ffefc125fdf6aa0db189e8172adf9a3aa63e1a63..f508a2f3ae0cac5ea55960a560fc8c1319d5f364 100644 (file)
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -251,7 +251,7 @@ rspamd_protocol_handle_url (struct rspamd_task *task,
                        value = v;
                        /* Steal strings */
                        g_hash_table_iter_steal (&it);
-                       g_hash_table_replace (task->request_headers, key, value);
+                       rspamd_task_add_request_header (task, key, value);
                        msg_debug_task ("added header \"%T\" -> \"%T\" from HTTP query",
                                        key, value);
                }
@@ -290,8 +290,6 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
                        hn_tok = rspamd_ftok_map (hn);
                        hv_tok = rspamd_ftok_map (hv);
 
-                       g_hash_table_replace (task->request_headers, hn_tok, hv_tok);
-
                        switch (*hn_tok->begin) {
                        case 'd':
                        case 'D':
@@ -469,6 +467,8 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
                                debug_task ("unknown header: %V", hn);
                                break;
                        }
+
+                       rspamd_task_add_request_header (task, hn_tok, hv_tok);
                }
        }
 

diff --git a/src/libserver/task.c b/src/libserver/task.c
index 086a25e7d449b259f2c32768bfc05c2a4cd5a3bf..95e435b6ed70be4c2260010c5c56ce27091e5e25 100644 (file)
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -37,6 +37,23 @@ rspamd_task_quark (void)
        return g_quark_from_static_string ("task-error");
 }
 
+static void
+rspamd_request_header_dtor (gpointer p)
+{
+       GPtrArray *ar = p;
+       guint i;
+       rspamd_ftok_t *tok;
+
+       if (ar) {
+               for (i = 0; i < ar->len; i ++) {
+                       tok = g_ptr_array_index (ar, i);
+                       rspamd_fstring_mapped_ftok_free (tok);
+               }
+
+               g_ptr_array_free (ar, TRUE);
+       }
+}
+
 /*
  * Create new task
  */
@@ -74,7 +91,7 @@ rspamd_task_new (struct rspamd_worker *worker, struct rspamd_config *cfg)
                        rspamd_strcase_equal);
        new_task->request_headers = g_hash_table_new_full (rspamd_ftok_icase_hash,
                        rspamd_ftok_icase_equal, rspamd_fstring_mapped_ftok_free,
-                       rspamd_fstring_mapped_ftok_free);
+                       rspamd_request_header_dtor);
        rspamd_mempool_add_destructor (new_task->task_pool,
                (rspamd_mempool_destruct_t) g_hash_table_unref,
                new_task->request_headers);
@@ -283,7 +300,7 @@ rspamd_task_load_message (struct rspamd_task *task,
        gchar filepath[PATH_MAX], *fp;
        gint fd, flen;
        gulong offset = 0, shmem_size = 0;
-       rspamd_ftok_t srch, *tok;
+       rspamd_ftok_t *tok;
        gpointer map;
        struct stat st;
        struct rspamd_task_map *m;
@@ -299,9 +316,7 @@ rspamd_task_load_message (struct rspamd_task *task,
                rspamd_protocol_handle_headers (task, msg);
        }
 
-       srch.begin = "shm";
-       srch.len = 3;
-       tok = g_hash_table_lookup (task->request_headers, &srch);
+       tok = rspamd_task_get_request_header (task, "shm");
 
        if (tok) {
                /* Shared memory part */
@@ -349,9 +364,7 @@ rspamd_task_load_message (struct rspamd_task *task,
 
                close (fd);
 
-               srch.begin = "shm-offset";
-               srch.len = 10;
-               tok = g_hash_table_lookup (task->request_headers, &srch);
+               tok = rspamd_task_get_request_header (task, "shm-offset");
 
                if (tok) {
                        rspamd_strtoul (tok->begin, tok->len, &offset);
@@ -365,11 +378,10 @@ rspamd_task_load_message (struct rspamd_task *task,
                        }
                }
 
-               srch.begin = "shm-length";
-               srch.len = 10;
-               tok = g_hash_table_lookup (task->request_headers, &srch);
+               tok = rspamd_task_get_request_header (task, "shm-length");
                shmem_size = st.st_size;
 
+
                if (tok) {
                        rspamd_strtoul (tok->begin, tok->len, &shmem_size);
 
@@ -397,14 +409,10 @@ rspamd_task_load_message (struct rspamd_task *task,
                return TRUE;
        }
 
-       srch.begin = "file";
-       srch.len = 4;
-       tok = g_hash_table_lookup (task->request_headers, &srch);
+       tok = rspamd_task_get_request_header (task, "file");
 
        if (tok == NULL) {
-               srch.begin = "path";
-               srch.len = 4;
-               tok = g_hash_table_lookup (task->request_headers, &srch);
+               tok = rspamd_task_get_request_header (task, "path");
        }
 
        if (tok) {
@@ -1313,3 +1321,59 @@ rspamd_task_get_required_score (struct rspamd_task *task, struct metric_result *
 
        return NAN;
 }
+
+rspamd_ftok_t *
+rspamd_task_get_request_header (struct rspamd_task *task,
+               const gchar *name)
+{
+       GPtrArray *ret;
+       rspamd_ftok_t srch;
+
+       srch.begin = (gchar *)name;
+       srch.len = strlen (name);
+
+       ret = g_hash_table_lookup (task->request_headers, &srch);
+
+       if (ret) {
+               return (rspamd_ftok_t *)g_ptr_array_index (ret, 0);
+       }
+
+       return NULL;
+}
+
+GPtrArray*
+rspamd_task_get_request_header_multiple (struct rspamd_task *task,
+               const gchar *name)
+{
+       GPtrArray *ret;
+       rspamd_ftok_t srch;
+
+       srch.begin = (gchar *)name;
+       srch.len = strlen (name);
+
+       ret = g_hash_table_lookup (task->request_headers, &srch);
+
+       return ret;
+}
+
+
+void
+rspamd_task_add_request_header (struct rspamd_task *task,
+               rspamd_ftok_t *name, rspamd_ftok_t *value)
+{
+       GPtrArray *ret;
+
+       ret = g_hash_table_lookup (task->request_headers, name);
+
+       if (ret) {
+               g_ptr_array_add (ret, value);
+
+               /* We need to free name token */
+               rspamd_fstring_mapped_ftok_free (name);
+       }
+       else {
+               ret = g_ptr_array_sized_new (2);
+               g_ptr_array_add (ret, value);
+               g_hash_table_replace (task->request_headers, name, ret);
+       }
+}

diff --git a/src/libserver/task.h b/src/libserver/task.h
index 80fdb82a368a3e44dd9e124b37e3993a29f22799..c76dd614e4bc9baac111659b6f1146a2fb5877e9 100644 (file)
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -276,6 +276,33 @@ struct metric_result;
 gdouble rspamd_task_get_required_score (struct rspamd_task *task,
                struct metric_result *m);
 
+/**
+ * Returns the first header as value for a header
+ * @param task
+ * @param name
+ * @return
+ */
+rspamd_ftok_t * rspamd_task_get_request_header (struct rspamd_task *task,
+               const gchar *name);
+
+/**
+ * Returns all headers with the specific name
+ * @param task
+ * @param name
+ * @return
+ */
+GPtrArray* rspamd_task_get_request_header_multiple (struct rspamd_task *task,
+               const gchar *name);
+
+/**
+ * Adds a new request header to task (name and value should be mapped to fstring)
+ * @param task
+ * @param name
+ * @param value
+ */
+void rspamd_task_add_request_header (struct rspamd_task *task,
+               rspamd_ftok_t *name, rspamd_ftok_t *value);
+
 /**
  * Write log line about the specified task if needed
  */

diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 1638f8f4990115b3f477a9ecf399d442800bd14b..511fa4a6f9e828cb07ce161a6956fb3c98571f69 100644 (file)
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -1268,19 +1268,15 @@ lua_task_get_parts (lua_State * L)
 static gint
 lua_task_get_request_header (lua_State *L)
 {
-       rspamd_ftok_t *hdr, srch;
+       rspamd_ftok_t *hdr;
        struct rspamd_task *task = lua_check_task (L, 1);
        const gchar *s;
        struct rspamd_lua_text *t;
-       gsize len;
 
-       s = luaL_checklstring (L, 2, &len);
+       s = luaL_checkstring (L, 2);
 
        if (s && task) {
-               srch.begin = (gchar *)s;
-               srch.len = len;
-
-               hdr = g_hash_table_lookup (task->request_headers, &srch);
+               hdr = rspamd_task_get_request_header (task, s);
 
                if (hdr) {
                        t = lua_newuserdata (L, sizeof (*t));
@@ -1309,7 +1305,7 @@ lua_task_set_request_header (lua_State *L)
        const gchar *s, *v = NULL;
        rspamd_fstring_t *buf;
        struct rspamd_lua_text *t;
-       rspamd_ftok_t *hdr, srch, *new_name;
+       rspamd_ftok_t *hdr, *new_name;
        gsize len, vlen;
 
        s = luaL_checklstring (L, 2, &len);
@@ -1328,25 +1324,12 @@ lua_task_set_request_header (lua_State *L)
                }
 
                if (v != NULL) {
-                       srch.begin = (gchar *)s;
-                       srch.len = len;
-
-                       hdr = g_hash_table_lookup (task->request_headers, &srch);
-
-                       if (!hdr) {
-                               new_name = &srch;
-                       }
-                       else {
-                               /* Not found, need to allocate */
-                               buf = rspamd_fstring_new_init (srch.begin, srch.len);
-                               new_name = rspamd_ftok_map (buf);
-                       }
-
                        buf = rspamd_fstring_new_init (v, vlen);
                        hdr = rspamd_ftok_map (buf);
+                       buf = rspamd_fstring_new_init (s, len);
+                       new_name = rspamd_ftok_map (buf);
 
-                       /* This does not destroy key if it exists */
-                       g_hash_table_insert (task->request_headers, new_name, hdr);
+                       rspamd_task_add_request_header (task, new_name, hdr);
                }
 
        }