Fix writing reply to a client.

Issue:	#94
Reported by:	@jpastuszek

4 files changed, 11 insertions, 3 deletions

diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index ac41dca0c..26efc9b7c 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -855,7 +855,7 @@ rspamd_protocol_write_reply (struct rspamd_task *task)
 	}
 	msg->date = time (NULL);
 
-	task->state = CLOSING_CONNECTION;
+	task->state = WRITING_REPLY;
 
 	debug_task ("writing reply to client");
 	if (task->error_code != 0) {
diff --git a/src/libserver/task.c b/src/libserver/task.c
index 860bb2b03..8eb0c7932 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -335,6 +335,7 @@ rspamd_task_process (struct rspamd_task *task,
 			/* Call write_socket to write reply and exit */
 			task->state = WRITE_REPLY;
 		}
+		task->s->wanna_die = TRUE;
 	}
 	else {
 		rspamd_lua_call_pre_filters (task);
diff --git a/src/libserver/task.h b/src/libserver/task.h
index bf14d2174..2fd634e54 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -69,6 +69,7 @@ struct rspamd_task {
 		WAIT_FILTER,
 		WAIT_POST_FILTER,
 		WRITE_REPLY,
+		WRITING_REPLY,
 		CLOSING_CONNECTION
 	} state;                                                    /**< current session state							*/
 	enum rspamd_command cmd;                                    /**< command										*/
diff --git a/src/worker.c b/src/worker.c
index c0e8cacd2..26e59ace2 100644
--- a/src/worker.c
+++ b/src/worker.c
@@ -141,7 +141,7 @@ rspamd_worker_error_handler (struct rspamd_http_connection *conn, GError *err)
 
 	msg_info ("abnormally closing connection from: %s, error: %s",
 		rspamd_inet_address_to_string (&task->client_addr), err->message);
-	if (task->state != CLOSING_CONNECTION) {
+	if (task->state != CLOSING_CONNECTION && task->state != WRITING_REPLY) {
 		/* We still need to write a reply */
 		task->error_code = err->code;
 		task->last_error =
@@ -166,7 +166,7 @@ rspamd_worker_finish_handler (struct rspamd_http_connection *conn,
 			rspamd_inet_address_to_string (&task->client_addr));
 		destroy_session (task->s);
 	}
-	else {
+	else if (task->state != WRITING_REPLY) {
 		/*
 		 * If all filters have finished their tasks, this function will trigger
 		 * writing a reply.
@@ -174,6 +174,12 @@ rspamd_worker_finish_handler (struct rspamd_http_connection *conn,
 		task->s->wanna_die = TRUE;
 		check_session_pending (task->s);
 	}
+	else {
+		/*
+		 * We are going to write a reply to a client
+		 */
+		task->state = CLOSING_CONNECTION;
+	}
 
 	return 0;
 }