set configuration of imv_test with each TNC handshake

3 files changed, 29 insertions, 8 deletions

diff --git a/src/libimcv/plugins/imv_test/imv_test.c b/src/libimcv/plugins/imv_test/imv_test.c
index 37178867d..7e1f05901 100644
--- a/src/libimcv/plugins/imv_test/imv_test.c
+++ b/src/libimcv/plugins/imv_test/imv_test.c
@@ -66,6 +66,7 @@ TNC_Result TNC_IMV_NotifyConnectionChange(TNC_IMVID imv_id,
 										  TNC_ConnectionState new_state)
 {
 	imv_state_t *state;
+	imv_test_state_t *test_state;
 	int rounds;
 
 	if (!imv_test)
@@ -76,12 +77,21 @@ TNC_Result TNC_IMV_NotifyConnectionChange(TNC_IMVID imv_id,
 	switch (new_state)
 	{
 		case TNC_CONNECTION_STATE_CREATE:
-			rounds = lib->settings->get_int(lib->settings,
-								"libimcv.plugins.imv-test.rounds", 0);
-			state = imv_test_state_create(connection_id, rounds);
+			state = imv_test_state_create(connection_id);
 			return imv_test->create_state(imv_test, state);
 		case TNC_CONNECTION_STATE_DELETE:
 			return imv_test->delete_state(imv_test, connection_id);
+		case TNC_CONNECTION_STATE_HANDSHAKE:
+			if (!imv_test->get_state(imv_test, connection_id, &state))
+			{
+				return TNC_RESULT_FATAL;
+			}
+			state->change_state(state, new_state);
+			rounds = lib->settings->get_int(lib->settings,
+								"libimcv.plugins.imv-test.rounds", 0);
+			test_state = (imv_test_state_t*)state;
+			test_state->set_rounds(test_state, rounds);
+			return TNC_RESULT_SUCCESS;
 		default:
 			return imv_test->change_state(imv_test, connection_id, new_state);
 	}
diff --git a/src/libimcv/plugins/imv_test/imv_test_state.c b/src/libimcv/plugins/imv_test/imv_test_state.c
index 21d0c4c58..85c9ba05f 100644
--- a/src/libimcv/plugins/imv_test/imv_test_state.c
+++ b/src/libimcv/plugins/imv_test/imv_test_state.c
@@ -89,6 +89,12 @@ METHOD(imv_state_t, destroy, void,
 	free(this);
 }
 
+METHOD(imv_test_state_t, set_rounds, void,
+	private_imv_test_state_t *this, int rounds)
+{
+	this->rounds = rounds;
+}
+
 METHOD(imv_test_state_t, another_round, bool,
 	private_imv_test_state_t *this)
 {
@@ -98,7 +104,7 @@ METHOD(imv_test_state_t, another_round, bool,
 /**
  * Described in header.
  */
-imv_state_t *imv_test_state_create(TNC_ConnectionID connection_id, int rounds)
+imv_state_t *imv_test_state_create(TNC_ConnectionID connection_id)
 {
 	private_imv_test_state_t *this;
 
@@ -111,13 +117,13 @@ imv_state_t *imv_test_state_create(TNC_ConnectionID connection_id, int rounds)
 				.set_recommendation = _set_recommendation,
 				.destroy = _destroy,
 			},
+			.set_rounds = _set_rounds,
 			.another_round = _another_round,
 		},
 		.state = TNC_CONNECTION_STATE_CREATE,
 		.rec = TNC_IMV_ACTION_RECOMMENDATION_NO_RECOMMENDATION,
 		.eval = TNC_IMV_EVALUATION_RESULT_DONT_KNOW,
 		.connection_id = connection_id,
-		.rounds = rounds,
 	);
 	
 	return &this->public.interface;
diff --git a/src/libimcv/plugins/imv_test/imv_test_state.h b/src/libimcv/plugins/imv_test/imv_test_state.h
index 164aca9b1..98d3dd92a 100644
--- a/src/libimcv/plugins/imv_test/imv_test_state.h
+++ b/src/libimcv/plugins/imv_test/imv_test_state.h
@@ -37,20 +37,25 @@ struct imv_test_state_t {
 	imv_state_t interface;
 
 	/**
+	 * Set the IMC-IMV round-trip count
+	 *
+	 * @param		number of re-measurement rounds
+	 */
+	void (*set_rounds)(imv_test_state_t *this, int rounds);
+
+	/**
 	 * Check and decrease IMC-IMV round-trip count
 	 *
 	 * @return		new connection state
 	 */
 	bool (*another_round)(imv_test_state_t *this);
-
 };
 
 /**
  * Create an imv_test_state_t instance
  *
  * @param id		connection ID
- * @param rounds	total number of IMC re-measurements
  */
-imv_state_t* imv_test_state_create(TNC_ConnectionID id, int rounds);
+imv_state_t* imv_test_state_create(TNC_ConnectionID id);
 
 #endif /** IMV_TEST_STATE_H_ @}*/