Add support of custom request and reply headers.

3 files changed, 46 insertions, 17 deletions

diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index 9d3e842cb..1b174b122 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -230,12 +230,13 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 	struct rspamd_http_message *msg)
 {
 	gchar *headern, *err, *tmp;
-	gboolean res = TRUE;
+	gboolean res = TRUE, validh;
 	struct rspamd_http_header *h;
 
 	LL_FOREACH (msg->headers, h)
 	{
 		headern = h->name->str;
+		validh = TRUE;
 
 		switch (headern[0]) {
 		case 'd':
@@ -247,7 +248,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 			}
 			else {
 				debug_task ("wrong header: %s", headern);
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'h':
@@ -262,7 +263,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 			}
 			else {
 				debug_task ("wrong header: %s", headern);
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'f':
@@ -273,7 +274,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 			}
 			else {
 				debug_task ("wrong header: %s", headern);
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'j':
@@ -283,7 +284,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 			}
 			else {
 				debug_task ("wrong header: %s", headern);
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'q':
@@ -294,7 +295,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 			}
 			else {
 				debug_task ("wrong header: %s", headern);
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'r':
@@ -310,7 +311,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 			}
 			else {
 				msg_info ("wrong header: %s", headern);
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'i':
@@ -325,7 +326,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 			}
 			else {
 				debug_task ("wrong header: %s", headern);
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'p':
@@ -338,7 +339,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 				}
 			}
 			else {
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 's':
@@ -347,7 +348,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 				task->subject = h->value->str;
 			}
 			else {
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'u':
@@ -356,7 +357,7 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 				task->user = h->value->str;
 			}
 			else {
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		case 'l':
@@ -367,14 +368,21 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 				}
 			}
 			else {
-				res = FALSE;
+				validh = FALSE;
 			}
 			break;
 		default:
-			debug_task ("wrong header: %s", headern);
-			res = FALSE;
+			debug_task ("unknown header: %s", headern);
+			validh = FALSE;
 			break;
 		}
+
+		if (!validh) {
+			res = FALSE;
+			g_hash_table_replace (task->request_headers,
+				g_string_new_len(h->name->str, h->name->len),
+				g_string_new_len(h->value->str, h->value->len));
+		}
 	}
 
 	if (!res && task->cfg->strict_protocol_headers) {
diff --git a/src/libserver/task.c b/src/libserver/task.c
index d71e0412e..59f015f1b 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -41,6 +41,15 @@ rcpt_destruct (void *pointer)
 	}
 }
 
+
+static void
+gstring_destruct (gpointer ptr)
+{
+	GString *s = (GString *)ptr;
+
+	g_string_free (s, TRUE);
+}
+
 /*
  * Create new task
  */
@@ -77,16 +86,26 @@ rspamd_task_new (struct rspamd_worker *worker)
 		(rspamd_mempool_destruct_t) rcpt_destruct, new_task);
 	new_task->results = g_hash_table_new (rspamd_str_hash, rspamd_str_equal);
 	rspamd_mempool_add_destructor (new_task->task_pool,
-		(rspamd_mempool_destruct_t) g_hash_table_destroy,
+		(rspamd_mempool_destruct_t) g_hash_table_unref,
 		new_task->results);
 	new_task->re_cache = g_hash_table_new (rspamd_str_hash, rspamd_str_equal);
 	rspamd_mempool_add_destructor (new_task->task_pool,
-		(rspamd_mempool_destruct_t) g_hash_table_destroy,
+		(rspamd_mempool_destruct_t) g_hash_table_unref,
 		new_task->re_cache);
 	new_task->raw_headers = g_hash_table_new (rspamd_strcase_hash,
 			rspamd_strcase_equal);
+	new_task->request_headers = g_hash_table_new_full ((GHashFunc)g_string_hash,
+		(GEqualFunc)g_string_equal, gstring_destruct, gstring_destruct);
+	rspamd_mempool_add_destructor (new_task->task_pool,
+		(rspamd_mempool_destruct_t) g_hash_table_unref,
+		new_task->request_headers);
+	new_task->reply_headers = g_hash_table_new_full ((GHashFunc)g_string_hash,
+		(GEqualFunc)g_string_equal, gstring_destruct, gstring_destruct);
+	rspamd_mempool_add_destructor (new_task->task_pool,
+		(rspamd_mempool_destruct_t) g_hash_table_unref,
+		new_task->reply_headers);
 	rspamd_mempool_add_destructor (new_task->task_pool,
-		(rspamd_mempool_destruct_t) g_hash_table_destroy,
+		(rspamd_mempool_destruct_t) g_hash_table_unref,
 		new_task->raw_headers);
 	new_task->emails = g_tree_new (compare_email_func);
 	rspamd_mempool_add_destructor (new_task->task_pool,
diff --git a/src/libserver/task.h b/src/libserver/task.h
index d9ec9d8d2..84262af6f 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -91,6 +91,8 @@ struct rspamd_task {
 	gchar *user;                                                    /**< user to deliver								*/
 	gchar *subject;                                             /**< subject (for non-mime)							*/
 	gchar *hostname;                                            /**< hostname reported by MTA						*/
+	GHashTable *request_headers;                                /**< HTTP headers in a request						*/
+	GHashTable *reply_headers;                                  /**< Custom reply headers							*/
 	GString *msg;                                               /**< message buffer									*/
 	struct rspamd_http_connection *http_conn;                   /**< HTTP server connection							*/
 	struct rspamd_async_session * s;                             /**< async session object							*/