From 9bcf7906eeb645e18f6ac95b5544f0da2c169eb2 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Thu, 29 Jan 2015 23:06:45 +0000
Subject: [PATCH] Fix keys initialization.

---
 src/client/rspamdclient.c |  2 +-
 src/libserver/protocol.c  |  5 +++++
 src/libserver/task.c      |  3 +++
 src/libserver/task.h      |  1 +
 src/libutil/http.c        | 14 +++++++++++++-
 src/libutil/http.h        |  8 +++++++-
 src/worker.c              |  6 +++++-
 7 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/client/rspamdclient.c b/src/client/rspamdclient.c
index c95f982cf..1d119eb18 100644
--- a/src/client/rspamdclient.c
+++ b/src/client/rspamdclient.c
@@ -251,7 +251,7 @@ rspamd_client_destroy (struct rspamd_client_connection *conn)
 			g_string_free (conn->key, TRUE);
 		}
 		if (conn->keypair) {
-			rspamd_http_connection_key_destroy (conn->keypair);
+			rspamd_http_connection_key_unref (conn->keypair);
 		}
 		g_string_free (conn->server_name, TRUE);
 		g_slice_free1 (sizeof (struct rspamd_client_connection), conn);
diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index 44bfe5a4e..bc64b7c31 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -905,6 +905,11 @@ rspamd_protocol_write_reply (struct rspamd_task *task)
 	const gchar *ctype = "application/json";
 
 	msg = rspamd_http_new_message (HTTP_RESPONSE);
+
+	if (task->peer_key) {
+		msg->peer_key = g_string_new (task->peer_key->str);
+		msg_info ("<%s> writing encrypted reply", task->message_id);
+	}
 	if (!task->is_json) {
 		/* Turn compatibility on */
 		msg->method = HTTP_SYMBOLS;
diff --git a/src/libserver/task.c b/src/libserver/task.c
index bfa9184fb..18ba08dad 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -266,6 +266,9 @@ rspamd_task_free (struct rspamd_task *task, gboolean is_soft)
 		if (task->settings != NULL) {
 			ucl_object_unref (task->settings);
 		}
+		if (task->peer_key != NULL) {
+			g_string_free (task->peer_key, TRUE);
+		}
 		rspamd_mempool_delete (task->task_pool);
 		g_slice_free1 (sizeof (struct rspamd_task), task);
 	}
diff --git a/src/libserver/task.h b/src/libserver/task.h
index 1261dd134..9c05def82 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -147,6 +147,7 @@ struct rspamd_task {
 	} pre_result;                                               /**< Result of pre-filters							*/
 
 	ucl_object_t *settings;                                     /**< Settings applied to task						*/
+	GString *peer_key;											/**< Peer's pubkey									*/
 };
 
 /**
diff --git a/src/libutil/http.c b/src/libutil/http.c
index 8fd0e7448..1b2338509 100644
--- a/src/libutil/http.c
+++ b/src/libutil/http.c
@@ -1197,6 +1197,7 @@ rspamd_http_new_message (enum http_parser_type type)
 	new->port = 80;
 	new->type = type;
 	new->method = HTTP_GET;
+	new->peer_key = NULL;
 
 	return new;
 }
@@ -1761,7 +1762,7 @@ rspamd_http_connection_set_key (struct rspamd_http_connection *conn,
 }
 
 void
-rspamd_http_connection_key_destroy (gpointer key)
+rspamd_http_connection_key_unref (gpointer key)
 {
 	struct rspamd_http_keypair *kp = (struct rspamd_http_keypair *)key;
 
@@ -1769,6 +1770,17 @@ rspamd_http_connection_key_destroy (gpointer key)
 	REF_RELEASE (kp);
 }
 
+gpointer
+rspamd_http_connection_key_ref (gpointer key)
+{
+	struct rspamd_http_keypair *kp = (struct rspamd_http_keypair *)key;
+
+	g_assert (key != NULL);
+	REF_RETAIN (kp);
+
+	return kp;
+}
+
 GString *
 rspamd_http_connection_make_peer_key (const gchar *key)
 {
diff --git a/src/libutil/http.h b/src/libutil/http.h
index b52671e10..83a284a52 100644
--- a/src/libutil/http.h
+++ b/src/libutil/http.h
@@ -193,7 +193,13 @@ GString *rspamd_http_connection_print_key (gpointer key, guint how);
  * Release key pointed by an opaque pointer
  * @param key opaque key structure
  */
-void rspamd_http_connection_key_destroy (gpointer key);
+void rspamd_http_connection_key_unref (gpointer key);
+
+/**
+ * Increase refcount for a key pointed by an opaque pointer
+ * @param key opaque key structure
+ */
+gpointer rspamd_http_connection_key_ref (gpointer key);
 
 GString *rspamd_http_connection_make_peer_key (const gchar *key);
 
diff --git a/src/worker.c b/src/worker.c
index 0a6ce4a5e..91797a6dc 100644
--- a/src/worker.c
+++ b/src/worker.c
@@ -129,6 +129,10 @@ rspamd_worker_body_handler (struct rspamd_http_connection *conn,
 		return 0;
 	}
 
+	if (msg->peer_key) {
+		task->peer_key = g_string_new (msg->peer_key->str);
+	}
+
 	if (!rspamd_task_process (task, msg, ctx->classify_pool, TRUE)) {
 		task->state = WRITE_REPLY;
 	}
@@ -354,7 +358,7 @@ start_worker (struct rspamd_worker *worker)
 	rspamd_log_close (rspamd_main->logger);
 
 	if (ctx->key) {
-		rspamd_http_connection_key_destroy (ctx->key);
+		rspamd_http_connection_key_unref (ctx->key);
 	}
 
 	exit (EXIT_SUCCESS);
-- 
2.39.5