diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c index 62fa4550fa..7b082a341b 100644 --- a/tests/src/psa_exercise_key.c +++ b/tests/src/psa_exercise_key.c @@ -693,7 +693,13 @@ psa_status_t mbedtls_test_psa_raw_key_agreement_with_self( size_t public_key_length; uint8_t output[1024]; size_t output_length; + + uint8_t *exported = NULL; + size_t exported_size = 0; + size_t exported_length = 0; + psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT; + psa_key_attributes_t export_attributes = PSA_KEY_ATTRIBUTES_INIT; mbedtls_svc_key_id_t shared_secret_id = MBEDTLS_SVC_KEY_ID_INIT; psa_key_attributes_t shared_secret_attributes = PSA_KEY_ATTRIBUTES_INIT; @@ -746,6 +752,28 @@ psa_status_t mbedtls_test_psa_raw_key_agreement_with_self( if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) { /* The key has been destroyed. */ status = PSA_SUCCESS; + goto exit; + } else if (status == PSA_SUCCESS) { + + status = psa_get_key_attributes(shared_secret_id, &export_attributes); + if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) { + /* The key has been destroyed. */ + status = PSA_SUCCESS; + goto exit; + } + + exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(psa_get_key_type(&export_attributes), + psa_get_key_bits(&export_attributes)); + TEST_CALLOC(exported, exported_size); + + status = psa_export_key(shared_secret_id, exported, exported_size, &exported_length); + + if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) { + /* The key has been destroyed. */ + status = PSA_SUCCESS; + } + + PSA_ASSERT(status); } exit: @@ -754,12 +782,14 @@ exit: * thus reset them as required. */ psa_reset_key_attributes(&attributes); + psa_reset_key_attributes(&export_attributes); /* Make sure to reset and free derived key attributes and slot. */ psa_reset_key_attributes(&shared_secret_attributes); psa_destroy_key(shared_secret_id); mbedtls_free(public_key); + mbedtls_free(exported); return status; }