From ab1d77168c98918b4fa2d93c0e7017212c6ad5f6 Mon Sep 17 00:00:00 2001
From: Tobias Brunner <tobias@strongswan.org>
Date: Thu, 28 Jun 2018 10:44:40 +0200
Subject: [PATCH] ikev2: Allow tasks to do work after generating
 requests/responses

---
 src/libcharon/sa/ikev2/task_manager_v2.c | 72 +++++++++++++++++++++---
 1 file changed, 65 insertions(+), 7 deletions(-)

diff --git a/src/libcharon/sa/ikev2/task_manager_v2.c b/src/libcharon/sa/ikev2/task_manager_v2.c
index f45d074e12..60cb8764f0 100644
--- a/src/libcharon/sa/ikev2/task_manager_v2.c
+++ b/src/libcharon/sa/ikev2/task_manager_v2.c
@@ -471,6 +471,7 @@ METHOD(task_manager_t, initiate, status_t,
 	message_t *message;
 	host_t *me, *other;
 	exchange_type_t exchange = 0;
+	bool result;
 
 	if (this->initiating.type != EXCHANGE_TYPE_UNDEFINED)
 	{
@@ -685,17 +686,47 @@ METHOD(task_manager_t, initiate, status_t,
 		return initiate(this);
 	}
 
-	if (!generate_message(this, message, &this->initiating.packets))
+	result = generate_message(this, message, &this->initiating.packets);
+
+	if (result)
 	{
-		/* message generation failed. There is nothing more to do than to
-		 * close the SA */
-		message->destroy(message);
-		flush(this);
-		charon->bus->ike_updown(charon->bus, this->ike_sa, FALSE);
-		return DESTROY_ME;
+		enumerator = array_create_enumerator(this->active_tasks);
+		while (enumerator->enumerate(enumerator, &task))
+		{
+			if (!task->post_build)
+			{
+				continue;
+			}
+			switch (task->post_build(task, message))
+			{
+				case SUCCESS:
+					array_remove_at(this->active_tasks, enumerator);
+					task->destroy(task);
+					break;
+				case NEED_MORE:
+					break;
+				default:
+					/* critical failure, destroy IKE_SA */
+					result = FALSE;
+					break;
+			}
+		}
+		enumerator->destroy(enumerator);
 	}
 	message->destroy(message);
 
+	if (!result)
+	{	/* message generation failed. There is nothing more to do than to
+		 * close the SA */
+		flush(this);
+		if (this->ike_sa->get_state(this->ike_sa) != IKE_CONNECTING &&
+			this->ike_sa->get_state(this->ike_sa) != IKE_REKEYED)
+		{
+			charon->bus->ike_updown(charon->bus, this->ike_sa, FALSE);
+		}
+		return DESTROY_ME;
+	}
+
 	array_compress(this->active_tasks);
 	array_compress(this->queued_tasks);
 
@@ -946,7 +977,34 @@ static status_t build_response(private_task_manager_t *this, message_t *request)
 	/* message complete, send it */
 	clear_packets(this->responding.packets);
 	result = generate_message(this, message, &this->responding.packets);
+
+	if (result && !delete)
+	{
+		enumerator = array_create_enumerator(this->passive_tasks);
+		while (enumerator->enumerate(enumerator, &task))
+		{
+			if (!task->post_build)
+			{
+				continue;
+			}
+			switch (task->post_build(task, message))
+			{
+				case SUCCESS:
+					array_remove_at(this->passive_tasks, enumerator);
+					task->destroy(task);
+					break;
+				case NEED_MORE:
+					break;
+				default:
+					/* critical failure, destroy IKE_SA */
+					result = FALSE;
+					break;
+			}
+		}
+		enumerator->destroy(enumerator);
+	}
 	message->destroy(message);
+
 	if (id)
 	{
 		id->set_responder_spi(id, responder_spi);