2 files changed, 30 insertions, 18 deletions

diff --git a/src/charon-tkm/src/tkm/tkm_nonceg.c b/src/charon-tkm/src/tkm/tkm_nonceg.c
index 89baab7ce..336f16ecd 100644
--- a/src/charon-tkm/src/tkm/tkm_nonceg.c
+++ b/src/charon-tkm/src/tkm/tkm_nonceg.c
@@ -33,23 +33,32 @@ struct private_tkm_nonceg_t {
 	tkm_nonceg_t public;
 
 	/**
-	 * Context id.
+	 * Nonce chunk.
 	 */
-	nc_id_type context_id;
-
+	chunk_t nonce;
 };
 
 METHOD(nonce_gen_t, get_nonce, bool,
 	private_tkm_nonceg_t *this, size_t size, u_int8_t *buffer)
 {
 	nonce_type nonce;
+	uint64_t nc_id;
 
-	if (ike_nc_create(this->context_id, size, &nonce) != TKM_OK)
+	nc_id = tkm->idmgr->acquire_id(tkm->idmgr, TKM_CTX_NONCE);
+	if (!nc_id)
 	{
 		return FALSE;
 	}
 
+	if (ike_nc_create(nc_id, size, &nonce) != TKM_OK)
+	{
+		tkm->idmgr->release_id(tkm->idmgr, TKM_CTX_NONCE, nc_id);
+		return FALSE;
+	}
+
 	memcpy(buffer, &nonce.data, size);
+	this->nonce = chunk_clone(chunk_create(buffer, size));
+	tkm->chunk_map->insert(tkm->chunk_map, &this->nonce, nc_id);
 	return TRUE;
 }
 
@@ -57,17 +66,27 @@ METHOD(nonce_gen_t, allocate_nonce, bool,
 	private_tkm_nonceg_t *this, size_t size, chunk_t *chunk)
 {
 	*chunk = chunk_alloc(size);
-	if (get_nonce(this, chunk->len, chunk->ptr))
-	{
-		tkm->chunk_map->insert(tkm->chunk_map, chunk, this->context_id);
-		return TRUE;
-	}
-	return FALSE;
+	return get_nonce(this, chunk->len, chunk->ptr);
 }
 
 METHOD(nonce_gen_t, destroy, void,
 	private_tkm_nonceg_t *this)
 {
+	uint64_t nc_id;
+
+	nc_id = tkm->chunk_map->get_id(tkm->chunk_map, &this->nonce);
+	if (nc_id)
+	{
+		DBG1(DBG_IKE, "resetting stale nonce context %llu", nc_id);
+
+		if (ike_nc_reset(nc_id) != TKM_OK)
+		{
+			DBG1(DBG_IKE, "failed to reset nonce context %llu", nc_id);
+		}
+		tkm->idmgr->release_id(tkm->idmgr, TKM_CTX_NONCE, nc_id);
+		tkm->chunk_map->remove(tkm->chunk_map, &this->nonce);
+	}
+	chunk_free(&this->nonce);
 	free(this);
 }
 
@@ -86,14 +105,7 @@ tkm_nonceg_t *tkm_nonceg_create()
 				.destroy = _destroy,
 			},
 		},
-		.context_id = tkm->idmgr->acquire_id(tkm->idmgr, TKM_CTX_NONCE),
 	);
 
-	if (!this->context_id)
-	{
-		free(this);
-		return NULL;
-	}
-
 	return &this->public;
 }
diff --git a/src/charon-tkm/tests/keymat_tests.c b/src/charon-tkm/tests/keymat_tests.c
index 889965a78..d087bee3f 100644
--- a/src/charon-tkm/tests/keymat_tests.c
+++ b/src/charon-tkm/tests/keymat_tests.c
@@ -46,7 +46,6 @@ START_TEST(test_derive_ike_keys)
 	fail_if(!ng, "Unable to create nonce generator");
 	fail_unless(ng->nonce_gen.allocate_nonce(&ng->nonce_gen, 32, &nonce),
 			"Unable to allocate nonce");
-	ng->nonce_gen.destroy(&ng->nonce_gen);
 
 	tkm_diffie_hellman_t *dh = tkm_diffie_hellman_create(MODP_4096_BIT);
 	fail_if(!dh, "Unable to create DH");
@@ -69,6 +68,7 @@ START_TEST(test_derive_ike_keys)
 	fail_if(aead->get_block_size(aead) != 16, "Block size mismatch %d",
 			aead->get_block_size(aead));
 
+	ng->nonce_gen.destroy(&ng->nonce_gen);
 	proposal->destroy(proposal);
 	dh->dh.destroy(&dh->dh);
 	ike_sa_id->destroy(ike_sa_id);