Fixed PB-TNC error handling

4 files changed, 32 insertions, 35 deletions

diff --git a/src/libtnccs/plugins/tnccs_20/tnccs_20.c b/src/libtnccs/plugins/tnccs_20/tnccs_20.c
index d8daf9a9e..a1a95733f 100644
--- a/src/libtnccs/plugins/tnccs_20/tnccs_20.c
+++ b/src/libtnccs/plugins/tnccs_20/tnccs_20.c
@@ -172,7 +172,7 @@ METHOD(tls_t, process, status_t,
 	private_tnccs_20_t *this, void *buf, size_t buflen)
 {
 	pb_tnc_batch_t *batch;
-	bool from_server;
+	bool from_server, fatal_header_error = FALSE;
 	status_t status;
 	chunk_t data;
 
@@ -198,7 +198,11 @@ METHOD(tls_t, process, status_t,
 	batch = pb_tnc_batch_create_from_data(data);
 	status = batch->process_header(batch, !this->mutual, this->is_server,
 								   &from_server);
-
+	if (status == FAILED)
+	{
+		fatal_header_error = TRUE;
+		status = VERIFY_ERROR;
+	}
 	this->to_server = this->mutual ? from_server : !this->is_server;
 
 	/* In the mutual case, first batch from TNC server requires a TNC client */
@@ -229,7 +233,8 @@ METHOD(tls_t, process, status_t,
 	}
 	if (status == VERIFY_ERROR)
 	{
-		this->tnccs_handler->handle_errors(this->tnccs_handler, batch);
+		this->tnccs_handler->handle_errors(this->tnccs_handler, batch,
+										   fatal_header_error);
 		status = NEED_MORE;
 	}
 	batch->destroy(batch);
diff --git a/src/libtnccs/plugins/tnccs_20/tnccs_20_client.c b/src/libtnccs/plugins/tnccs_20/tnccs_20_client.c
index 9bfadcb66..80ae4b64e 100644
--- a/src/libtnccs/plugins/tnccs_20/tnccs_20_client.c
+++ b/src/libtnccs/plugins/tnccs_20/tnccs_20_client.c
@@ -456,9 +456,7 @@ static void build_retry_batch(private_tnccs_20_client_t *this)
 METHOD(tnccs_20_handler_t, process, status_t,
 	private_tnccs_20_client_t *this, pb_tnc_batch_t *batch)
 {
-	pb_tnc_msg_t *msg;
 	pb_tnc_batch_type_t batch_type;
-	enumerator_t *enumerator;
 	status_t status;
 
 	batch_type = batch->get_type(batch);
@@ -507,20 +505,9 @@ METHOD(tnccs_20_handler_t, process, status_t,
 	{
 		case FAILED:
 			this->fatal_error = TRUE;
-			this->mutex->lock(this->mutex);
-			change_batch_type(this, PB_BATCH_CLOSE);
-			this->mutex->unlock(this->mutex);
 			status = VERIFY_ERROR;
-			/* fall through to add error messages to outbound batch */
+			break;
 		case VERIFY_ERROR:
-			enumerator = batch->create_error_enumerator(batch);
-			while (enumerator->enumerate(enumerator, &msg))
-			{
-				this->mutex->lock(this->mutex);
-				this->messages->insert_last(this->messages, msg->get_ref(msg));
-				this->mutex->unlock(this->mutex);
-			}
-			enumerator->destroy(enumerator);
 			break;
 		case SUCCESS:
 		default:
@@ -728,11 +715,19 @@ METHOD(tnccs_20_handler_t, add_msg, void,
 }
 
 METHOD(tnccs_20_handler_t, handle_errors, void,
-	private_tnccs_20_client_t *this, pb_tnc_batch_t *batch)
+	private_tnccs_20_client_t *this, pb_tnc_batch_t *batch,
+	bool fatal_header_error)
 {
 	pb_tnc_msg_t *msg;
 	enumerator_t *enumerator;
 
+	if (fatal_header_error || this->fatal_error)
+	{
+		this->mutex->lock(this->mutex);
+		change_batch_type(this, PB_BATCH_CLOSE);
+		this->mutex->unlock(this->mutex);
+	}
+
 	enumerator = batch->create_error_enumerator(batch);
 	while (enumerator->enumerate(enumerator, &msg))
 	{
diff --git a/src/libtnccs/plugins/tnccs_20/tnccs_20_handler.h b/src/libtnccs/plugins/tnccs_20/tnccs_20_handler.h
index d0adb2489..5c4d7a7b4 100644
--- a/src/libtnccs/plugins/tnccs_20/tnccs_20_handler.h
+++ b/src/libtnccs/plugins/tnccs_20/tnccs_20_handler.h
@@ -90,9 +90,11 @@ struct tnccs_20_handler_t {
 	/**
 	 * Handle errors that occurred during PB-TNC batch header processing
 	 *
-	 * @param batch			batch where a fatal error occurred
+	 * @param batch					batch where a fatal error occurred
+	 * @param fatal_header_error	TRUE if fatal error in batch header
 	 */
-	void (*handle_errors)(tnccs_20_handler_t *this, pb_tnc_batch_t *batch);
+	void (*handle_errors)(tnccs_20_handler_t *this, pb_tnc_batch_t *batch,
+						  bool fatal_header_error);
 
 	/**
 	 * Destroys a tnccs_20_handler_t object.
diff --git a/src/libtnccs/plugins/tnccs_20/tnccs_20_server.c b/src/libtnccs/plugins/tnccs_20/tnccs_20_server.c
index c75ecdc39..e6ac21ac1 100644
--- a/src/libtnccs/plugins/tnccs_20/tnccs_20_server.c
+++ b/src/libtnccs/plugins/tnccs_20/tnccs_20_server.c
@@ -285,9 +285,7 @@ static void build_retry_batch(private_tnccs_20_server_t *this)
 METHOD(tnccs_20_handler_t, process, status_t,
 	private_tnccs_20_server_t *this, pb_tnc_batch_t *batch)
 {
-	pb_tnc_msg_t *msg;
 	pb_tnc_batch_type_t batch_type;
-	enumerator_t *enumerator;
 	status_t status;
 
 	batch_type = batch->get_type(batch);
@@ -333,20 +331,9 @@ METHOD(tnccs_20_handler_t, process, status_t,
 	{
 		case FAILED:
 			this->fatal_error = TRUE;
-			this->mutex->lock(this->mutex);
-			change_batch_type(this, PB_BATCH_CLOSE);
-			this->mutex->unlock(this->mutex);
 			status = VERIFY_ERROR;
-			/* fall through to add error messages to outbound batch */
+			break;
 		case VERIFY_ERROR:
-			enumerator = batch->create_error_enumerator(batch);
-			while (enumerator->enumerate(enumerator, &msg))
-			{
-				this->mutex->lock(this->mutex);
-				this->messages->insert_last(this->messages, msg->get_ref(msg));
-				this->mutex->unlock(this->mutex);
-			}
-			enumerator->destroy(enumerator);
 			break;
 		case SUCCESS:
 		default:
@@ -601,11 +588,19 @@ METHOD(tnccs_20_handler_t, add_msg, void,
 }
 
 METHOD(tnccs_20_handler_t, handle_errors, void,
-	private_tnccs_20_server_t *this,  pb_tnc_batch_t *batch)
+	private_tnccs_20_server_t *this,  pb_tnc_batch_t *batch,
+	bool fatal_header_error)
 {
 	pb_tnc_msg_t *msg;
 	enumerator_t *enumerator;
 
+	if (fatal_header_error || this->fatal_error)
+	{
+		this->mutex->lock(this->mutex);
+		change_batch_type(this, PB_BATCH_CLOSE);
+		this->mutex->unlock(this->mutex);
+	}
+
 	enumerator = batch->create_error_enumerator(batch);
 	while (enumerator->enumerate(enumerator, &msg))
 	{