From 1153e7bd574aee4f1727c3c9a2dc7c0221ec4e83 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Tue, 28 May 2019 15:10:21 +0200
Subject: [PATCH] Replace all calls to mbedtls_psa_crypto_free by PSA_DONE

Replace all calls to mbedtls_psa_crypto_free in tests by PSA_DONE.
This is correct for most tests, because most tests close open keys. A
few tests now fail; these tests need to be reviewed and switched back
to mbedtls_psa_crypto_free if they genuinely expected to end with some
slots still in use.
---
 tests/suites/test_suite_psa_crypto.function   | 126 +++++++++---------
 .../test_suite_psa_crypto_entropy.function    |   6 +-
 .../test_suite_psa_crypto_hash.function       |   6 +-
 .../test_suite_psa_crypto_init.function       |  16 +--
 ...t_suite_psa_crypto_persistent_key.function |  16 +--
 ..._suite_psa_crypto_slot_management.function |  24 ++--
 6 files changed, 97 insertions(+), 97 deletions(-)

diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 2e2606f21..acc2f8c18 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1259,7 +1259,7 @@ void import( data_t *data, int type_arg,
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1296,7 +1296,7 @@ void import_rsa_made_up( int bits_arg, int keypair, int expected_status_arg )
 
 exit:
     mbedtls_free( buffer );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1389,7 +1389,7 @@ exit:
     mbedtls_free( exported );
     mbedtls_free( reexported );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1400,7 +1400,7 @@ void invalid_handle( int handle )
     test_operations_on_invalid_handle( handle );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1453,7 +1453,7 @@ exit:
     mbedtls_free( exported );
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1495,7 +1495,7 @@ void import_and_exercise_key( data_t *data,
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1527,7 +1527,7 @@ void key_policy( int usage_arg, int alg_arg )
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1607,7 +1607,7 @@ void mac_key_policy( int policy_usage,
 exit:
     psa_mac_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1650,7 +1650,7 @@ void cipher_key_policy( int policy_usage,
 exit:
     psa_cipher_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1711,7 +1711,7 @@ void aead_key_policy( int policy_usage,
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1772,7 +1772,7 @@ void asymmetric_encryption_key_policy( int policy_usage,
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( buffer );
 }
 /* END_CASE */
@@ -1827,7 +1827,7 @@ void asymmetric_signature_key_policy( int policy_usage,
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1866,7 +1866,7 @@ void derive_key_policy( int policy_usage,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1904,7 +1904,7 @@ void agreement_key_policy( int policy_usage,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1941,7 +1941,7 @@ void key_policy_alg2( int key_type_arg, data_t *key_data,
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1978,7 +1978,7 @@ void raw_agreement_key_policy( int policy_usage,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2059,7 +2059,7 @@ void copy_success( int source_usage_arg,
 exit:
     psa_reset_key_attributes( &source_attributes );
     psa_reset_key_attributes( &target_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 }
 /* END_CASE */
@@ -2103,7 +2103,7 @@ void copy_fail( int source_usage_arg,
 exit:
     psa_reset_key_attributes( &source_attributes );
     psa_reset_key_attributes( &target_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2165,7 +2165,7 @@ void hash_setup( int alg_arg,
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2253,7 +2253,7 @@ void hash_bad_order( )
     PSA_ASSERT( psa_hash_abort( &operation ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2288,7 +2288,7 @@ void hash_verify_bad_args( )
                 PSA_ERROR_INVALID_SIGNATURE );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2310,7 +2310,7 @@ void hash_finish_bad_args( )
                 PSA_ERROR_BUFFER_TOO_SMALL );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2355,7 +2355,7 @@ exit:
     psa_hash_abort( &op_setup );
     psa_hash_abort( &op_finished );
     psa_hash_abort( &op_aborted );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2396,7 +2396,7 @@ exit:
     psa_hash_abort( &op_setup );
     psa_hash_abort( &op_finished );
     psa_hash_abort( &op_aborted );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2467,7 +2467,7 @@ void mac_setup( int key_type_arg,
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2588,7 +2588,7 @@ void mac_bad_order( )
     PSA_ASSERT( psa_mac_abort( &operation ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2643,7 +2643,7 @@ void mac_sign( int key_type_arg,
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2681,7 +2681,7 @@ void mac_verify( int key_type_arg,
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2759,7 +2759,7 @@ void cipher_setup( int key_type_arg,
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2909,7 +2909,7 @@ void cipher_bad_order( )
     PSA_ASSERT( psa_cipher_abort( &operation ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2969,7 +2969,7 @@ void cipher_encrypt( int alg_arg, int key_type_arg,
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3037,7 +3037,7 @@ void cipher_encrypt_multipart( int alg_arg, int key_type_arg,
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3108,7 +3108,7 @@ void cipher_decrypt_multipart( int alg_arg, int key_type_arg,
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3169,7 +3169,7 @@ void cipher_decrypt( int alg_arg, int key_type_arg,
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3251,7 +3251,7 @@ exit:
     mbedtls_free( output1 );
     mbedtls_free( output2 );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3354,7 +3354,7 @@ exit:
     mbedtls_free( output1 );
     mbedtls_free( output2 );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3430,7 +3430,7 @@ exit:
     psa_destroy_key( handle );
     mbedtls_free( output_data );
     mbedtls_free( output_data2 );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3480,7 +3480,7 @@ void aead_encrypt( int key_type_arg, data_t *key_data,
 exit:
     psa_destroy_key( handle );
     mbedtls_free( output_data );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3536,7 +3536,7 @@ void aead_decrypt( int key_type_arg, data_t *key_data,
 exit:
     psa_destroy_key( handle );
     mbedtls_free( output_data );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3601,7 +3601,7 @@ exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3646,7 +3646,7 @@ exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3713,7 +3713,7 @@ exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3745,7 +3745,7 @@ void asymmetric_verify( int key_type_arg, data_t *key_data,
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3781,7 +3781,7 @@ void asymmetric_verify_fail( int key_type_arg, data_t *key_data,
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3850,7 +3850,7 @@ exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3915,7 +3915,7 @@ exit:
     psa_destroy_key( handle );
     mbedtls_free( output );
     mbedtls_free( output2 );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3977,7 +3977,7 @@ exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4039,7 +4039,7 @@ exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4107,7 +4107,7 @@ void derive_setup( int key_type_arg,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4156,7 +4156,7 @@ void test_derive_invalid_key_derivation_state( )
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4298,7 +4298,7 @@ exit:
     mbedtls_free( output_buffer );
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4379,7 +4379,7 @@ void derive_full( int alg_arg,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4439,7 +4439,7 @@ exit:
     psa_reset_key_attributes( &got_attributes );
     psa_destroy_key( base_handle );
     psa_destroy_key( derived_handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4518,7 +4518,7 @@ exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( base_handle );
     psa_destroy_key( derived_handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4566,7 +4566,7 @@ void key_agreement_setup( int alg_arg,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4603,7 +4603,7 @@ void raw_key_agreement( int alg_arg,
 exit:
     mbedtls_free( output );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4663,7 +4663,7 @@ void key_agreement_capacity( int alg_arg,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4722,7 +4722,7 @@ void key_agreement_output( int alg_arg,
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( actual_output );
 }
 /* END_CASE */
@@ -4772,7 +4772,7 @@ void generate_random( int bytes_arg )
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( output );
     mbedtls_free( changed );
 }
@@ -4818,7 +4818,7 @@ void generate_key( int type_arg,
 exit:
     psa_reset_key_attributes( &got_attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4919,7 +4919,7 @@ void generate_key_rsa( int bits_arg,
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( e_read_buffer );
     mbedtls_free( exported );
 }
@@ -5016,7 +5016,7 @@ void persistent_key_load_key_from_storage( data_t *data,
     }
 
     /* Shutdown and restart */
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     PSA_ASSERT( psa_crypto_init() );
 
     /* Check key slot still contains key data */
@@ -5058,6 +5058,6 @@ exit:
         psa_open_key( key_id, &handle );
     }
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_entropy.function b/tests/suites/test_suite_psa_crypto_entropy.function
index 8576c7d95..cd1b81f9e 100644
--- a/tests/suites/test_suite_psa_crypto_entropy.function
+++ b/tests/suites/test_suite_psa_crypto_entropy.function
@@ -77,7 +77,7 @@ void validate_entropy_seed_injection( int seed_length_a,
 exit:
     mbedtls_free( seed );
     remove_seed_file( );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -105,12 +105,12 @@ void run_entropy_inject_with_crypto_init( )
     PSA_ASSERT( status );
     status = psa_crypto_init( );
     PSA_ASSERT( status );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     /* The seed is written by nv_seed callback functions therefore the injection will fail */
     status = mbedtls_psa_inject_entropy( seed, sizeof( seed ) );
     TEST_EQUAL( status, PSA_ERROR_NOT_PERMITTED );
 exit:
     remove_seed_file( );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_hash.function b/tests/suites/test_suite_psa_crypto_hash.function
index 90636b97d..e15f335e8 100644
--- a/tests/suites/test_suite_psa_crypto_hash.function
+++ b/tests/suites/test_suite_psa_crypto_hash.function
@@ -31,7 +31,7 @@ void hash_finish( int alg_arg, data_t *input, data_t *expected_hash )
                     actual_hash, actual_hash_length );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -52,7 +52,7 @@ void hash_verify( int alg_arg, data_t *input, data_t *expected_hash )
                                  expected_hash->len ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -95,6 +95,6 @@ void hash_multi_part( int alg_arg, data_t *input, data_t *expected_hash )
     } while( len++ != input->len );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_init.function b/tests/suites/test_suite_psa_crypto_init.function
index 79131587c..eaf1b8b1e 100644
--- a/tests/suites/test_suite_psa_crypto_init.function
+++ b/tests/suites/test_suite_psa_crypto_init.function
@@ -138,7 +138,7 @@ void init_deinit( int count )
         PSA_ASSERT( status );
         status = psa_crypto_init( );
         PSA_ASSERT( status );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
 }
 /* END_CASE */
@@ -150,9 +150,9 @@ void deinit_without_init( int count )
     for( i = 0; i < count; i++ )
     {
         PSA_ASSERT( psa_crypto_init( ) );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -166,7 +166,7 @@ void validate_module_init_generate_random( int count )
     {
         status = psa_crypto_init( );
         PSA_ASSERT( status );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
     status = psa_generate_random( random, sizeof( random ) );
     TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
@@ -186,7 +186,7 @@ void validate_module_init_key_based( int count )
     {
         status = psa_crypto_init( );
         PSA_ASSERT( status );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
     psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
     status = psa_import_key( &attributes, data, sizeof( data ), &handle );
@@ -212,7 +212,7 @@ void custom_entropy_sources( int sources_arg, int expected_init_status_arg )
     PSA_ASSERT( psa_generate_random( random, sizeof( random ) ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -252,7 +252,7 @@ void fake_entropy_source( int threshold,
     PSA_ASSERT( psa_generate_random( random, sizeof( random ) ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -280,6 +280,6 @@ void entropy_from_nv_seed( int seed_size_arg,
 
 exit:
     mbedtls_free( seed );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_persistent_key.function b/tests/suites/test_suite_psa_crypto_persistent_key.function
index 7e98fae87..e4ab1633c 100644
--- a/tests/suites/test_suite_psa_crypto_persistent_key.function
+++ b/tests/suites/test_suite_psa_crypto_persistent_key.function
@@ -110,7 +110,7 @@ void save_large_persistent_key( int data_too_large, int expected_status )
 
 exit:
     mbedtls_free( data );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
@@ -137,7 +137,7 @@ void persistent_key_destroy( int key_id_arg, int restart,
     if( restart )
     {
         psa_close_key( handle );
-        mbedtls_psa_crypto_free();
+        PSA_DONE();
         PSA_ASSERT( psa_crypto_init() );
         PSA_ASSERT( psa_open_key( key_id, &handle ) );
     }
@@ -152,7 +152,7 @@ void persistent_key_destroy( int key_id_arg, int restart,
     TEST_EQUAL( handle, 0 );
 
     /* Shutdown and restart */
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     PSA_ASSERT( psa_crypto_init() );
 
     /* Create another key in the same slot */
@@ -162,7 +162,7 @@ void persistent_key_destroy( int key_id_arg, int restart,
                                 &handle ) );
 
 exit:
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
@@ -192,7 +192,7 @@ void persistent_key_import( int key_id_arg, int type_arg, data_t *data,
     if( restart )
     {
         psa_close_key( handle );
-        mbedtls_psa_crypto_free();
+        PSA_DONE();
         PSA_ASSERT( psa_crypto_init() );
         PSA_ASSERT( psa_open_key( key_id, &handle ) );
     }
@@ -209,7 +209,7 @@ void persistent_key_import( int key_id_arg, int type_arg, data_t *data,
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_persistent_key( key_id );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
 }
 /* END_CASE */
 
@@ -241,7 +241,7 @@ void import_export_persistent_key( data_t *data, int type_arg,
     if( restart )
     {
         psa_close_key( handle );
-        mbedtls_psa_crypto_free();
+        PSA_DONE();
         PSA_ASSERT( psa_crypto_init() );
         PSA_ASSERT( psa_open_key( key_id, &handle ) );
     }
@@ -276,7 +276,7 @@ void import_export_persistent_key( data_t *data, int type_arg,
 exit:
     psa_reset_key_attributes( &attributes );
     mbedtls_free( exported );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index a7bb59673..fde3b4dfe 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -103,7 +103,7 @@ void transient_slot_lifecycle( int usage_arg, int alg_arg,
             PSA_ASSERT( psa_destroy_key( handle ) );
             break;
         case CLOSE_BY_SHUTDOWN:
-            mbedtls_psa_crypto_free( );
+            PSA_DONE( );
             PSA_ASSERT( psa_crypto_init( ) );
             break;
     }
@@ -114,7 +114,7 @@ void transient_slot_lifecycle( int usage_arg, int alg_arg,
     TEST_EQUAL( psa_close_key( handle ), PSA_ERROR_INVALID_HANDLE );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -180,7 +180,7 @@ void persistent_slot_lifecycle( int lifetime_arg, int id_arg,
             PSA_ASSERT( psa_destroy_key( handle ) );
             break;
         case CLOSE_BY_SHUTDOWN:
-            mbedtls_psa_crypto_free( );
+            PSA_DONE( );
             PSA_ASSERT( psa_crypto_init( ) );
             break;
     }
@@ -236,7 +236,7 @@ void persistent_slot_lifecycle( int lifetime_arg, int id_arg,
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_purge_key_storage( );
     mbedtls_free( reexported );
 }
@@ -303,7 +303,7 @@ void create_existent( int lifetime_arg, int id_arg,
                     reexported, reexported_length );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_purge_key_storage( );
 }
 /* END_CASE */
@@ -322,7 +322,7 @@ void open_fail( int id_arg,
     TEST_EQUAL( handle, 0 );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -350,7 +350,7 @@ void create_fail( int lifetime_arg, int id_arg,
     TEST_EQUAL( handle, 0 );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
 #endif
@@ -428,7 +428,7 @@ void copy_across_lifetimes( int source_lifetime_arg, int source_id_arg,
      * sure that the material is still alive. */
     if( target_lifetime != PSA_KEY_LIFETIME_VOLATILE )
     {
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
         PSA_ASSERT( psa_crypto_init( ) );
         PSA_ASSERT( psa_open_key( target_id, &target_handle ) );
     }
@@ -464,7 +464,7 @@ void copy_across_lifetimes( int source_lifetime_arg, int source_id_arg,
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
@@ -567,7 +567,7 @@ void copy_to_occupied( int source_lifetime_arg, int source_id_arg,
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
@@ -609,7 +609,7 @@ void invalid_handle( )
     PSA_ASSERT( psa_close_key( handle1 ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -657,7 +657,7 @@ void many_transient_handles( int max_handles_arg )
     PSA_ASSERT( psa_close_key( handles[i - 1] ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( handles );
 }
 /* END_CASE */