From d8555003485d0a47d88ec7a80a717c730e85ffbd Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Thu, 4 Aug 2016 18:31:28 +0100 Subject: [PATCH] [Fix] Fix multiple request headers structure --- src/libmime/message.c | 4 +- src/libserver/protocol.c | 6 +-- src/libserver/task.c | 98 +++++++++++++++++++++++++++++++++------- src/libserver/task.h | 27 +++++++++++ src/lua/lua_task.c | 31 +++---------- 5 files changed, 119 insertions(+), 47 deletions(-) diff --git a/src/libmime/message.c b/src/libmime/message.c index 87a0123fc..a3b673f28 100644 --- 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 ffefc125f..f508a2f3a 100644 --- 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 086a25e7d..95e435b6e 100644 --- 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 80fdb82a3..c76dd614e 100644 --- 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 1638f8f49..511fa4a6f 100644 --- 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); } } -- 2.39.5