Merge pull request #132 from gilles-peskine-arm/mbedtls-stats-api

Test that closing a handle reclaims its resources
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 56e0536..b08f46d 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -116,6 +116,43 @@
  */
 void mbedtls_psa_crypto_free( void );
 
+/** \brief Statistics about
+ * resource consumption related to the PSA keystore.
+ *
+ * \note The content of this structure is not part of the stable API and ABI
+ *       of Mbed Crypto and may change arbitrarily from version to version.
+ */
+typedef struct mbedtls_psa_stats_s
+{
+    /** Number of slots containing key material for a volatile key. */
+    size_t volatile_slots;
+    /** Number of slots containing key material for a key which is in
+     * internal persistent storage. */
+    size_t persistent_slots;
+    /** Number of slots containing a reference to a key in a
+     * secure element. */
+    size_t external_slots;
+    /** Number of slots which are occupied, but do not contain
+     * key material yet. */
+    size_t half_filled_slots;
+    /** Number of slots that contain cache data. */
+    size_t cache_slots;
+    /** Number of slots that are not used for anything. */
+    size_t empty_slots;
+    /** Largest key id value among open keys in internal persistent storage. */
+    psa_key_id_t max_open_internal_key_id;
+    /** Largest key id value among open keys in secure elements. */
+    psa_key_id_t max_open_external_key_id;
+} mbedtls_psa_stats_t;
+
+/** \brief Get statistics about
+ * resource consumption related to the PSA keystore.
+ *
+ * \note When Mbed Crypto is built as part of a service, with isolation
+ *       between the application and the keystore, the service may or
+ *       may not expose this function.
+ */
+void mbedtls_psa_get_stats( mbedtls_psa_stats_t *stats );
 
 /**
  * \brief Inject an initial entropy seed for the random generator into
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 0ffc2aa..900aa41 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -232,4 +232,36 @@
     return( psa_wipe_key_slot( slot ) );
 }
 
+void mbedtls_psa_get_stats( mbedtls_psa_stats_t *stats )
+{
+    psa_key_handle_t key;
+    memset( stats, 0, sizeof( *stats ) );
+    for( key = 1; key <= PSA_KEY_SLOT_COUNT; key++ )
+    {
+        psa_key_slot_t *slot = &global_data.key_slots[key - 1];
+        if( slot->type == PSA_KEY_TYPE_NONE )
+        {
+            if( slot->allocated )
+                ++stats->half_filled_slots;
+            else
+                ++stats->empty_slots;
+            continue;
+        }
+        if( slot->lifetime == PSA_KEY_LIFETIME_VOLATILE )
+            ++stats->volatile_slots;
+        else if( slot->lifetime == PSA_KEY_LIFETIME_PERSISTENT )
+        {
+            ++stats->persistent_slots;
+            if( slot->persistent_storage_id > stats->max_open_internal_key_id )
+                stats->max_open_internal_key_id = slot->persistent_storage_id;
+        }
+        else
+        {
+            ++stats->external_slots;
+            if( slot->persistent_storage_id > stats->max_open_external_key_id )
+                stats->max_open_external_key_id = slot->persistent_storage_id;
+        }
+    }
+}
+
 #endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/tests/Makefile b/tests/Makefile
index aba002b..94f0bc4 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -104,6 +104,11 @@
 	echo "  CC    $<"
 	$(CC) $(LOCAL_CFLAGS) $(CFLAGS) $<	$(LOCAL_LDFLAGS) $(LDFLAGS) -o $@
 
+# Some test suites require additional header files.
+$(filter test_suite_psa_crypto%, $(BINARIES)): psa_crypto_helpers.h
+$(addprefix embedded_,$(filter test_suite_psa_crypto%, $(APPS))): embedded_%: TESTS/mbedtls/%/psa_crypto_helpers.h
+$(filter test_suite_psa_%, $(BINARIES)): psa_helpers.h
+$(addprefix embedded_,$(filter test_suite_psa_%, $(APPS))): embedded_%: TESTS/mbedtls/%/psa_helpers.h
 
 clean:
 ifndef WINDOWS
@@ -141,3 +146,17 @@
 
 generate-target-tests: $(EMBEDDED_TESTS)
 
+define copy_header_to_target
+TESTS/mbedtls/$(1)/$(2): $(2)
+	echo "  Copy ./$$@"
+ifndef WINDOWS
+	mkdir -p $$(@D)
+	cp $$< $$@
+else
+	mkdir $$(@D)
+	copy $$< $$@
+endif
+
+endef
+$(foreach app, $(APPS), $(foreach file, $(wildcard *.h), \
+	$(eval $(call copy_header_to_target,$(app),$(file)))))
diff --git a/tests/psa_crypto_helpers.h b/tests/psa_crypto_helpers.h
new file mode 100644
index 0000000..26d5623
--- /dev/null
+++ b/tests/psa_crypto_helpers.h
@@ -0,0 +1,75 @@
+/*
+ * Helper functions for tests that use the PSA Crypto API.
+ */
+/*  Copyright (C) 2019, ARM Limited, All Rights Reserved
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *  This file is part of mbed TLS (https://tls.mbed.org)
+ */
+
+#ifndef PSA_CRYPTO_HELPERS_H
+#define PSA_CRYPTO_HELPERS_H
+
+#include "psa_helpers.h"
+
+#include <psa/crypto.h>
+
+static int test_helper_is_psa_pristine( int line, const char *file )
+{
+    mbedtls_psa_stats_t stats;
+    const char *msg = NULL;
+
+    mbedtls_psa_get_stats( &stats );
+
+    if( stats.volatile_slots != 0 )
+        msg = "A volatile slot has not been closed properly.";
+    else if( stats.persistent_slots != 0 )
+        msg = "A persistent slot has not been closed properly.";
+    else if( stats.external_slots != 0 )
+        msg = "An external slot has not been closed properly.";
+    else if( stats.half_filled_slots != 0 )
+        msg = "A half-filled slot has not been cleared properly.";
+
+    /* If the test has already failed, don't overwrite the failure
+     * information. Do keep the stats lookup above, because it can be
+     * convenient to break on it when debugging a failure. */
+    if( msg != NULL && test_info.failed == 0 )
+        test_fail( msg, line, file );
+
+    return( msg == NULL );
+}
+
+/** Check that no PSA Crypto key slots are in use.
+ */
+#define ASSERT_PSA_PRISTINE( )                                    \
+    do                                                            \
+    {                                                             \
+        if( ! test_helper_is_psa_pristine( __LINE__, __FILE__ ) ) \
+            goto exit;                                            \
+    }                                                             \
+    while( 0 )
+
+static void test_helper_psa_done( int line, const char *file )
+{
+    (void) test_helper_is_psa_pristine( line, file );
+    mbedtls_psa_crypto_free( );
+}
+
+/** Shut down the PSA Crypto subsystem. Expect a clean shutdown, with no slots
+ * in use.
+ */
+#define PSA_DONE( ) test_helper_psa_done( __LINE__, __FILE__ )
+
+#endif /* PSA_CRYPTO_HELPERS_H */
diff --git a/tests/psa_helpers.h b/tests/psa_helpers.h
new file mode 100644
index 0000000..79f6837
--- /dev/null
+++ b/tests/psa_helpers.h
@@ -0,0 +1,37 @@
+/*
+ * Helper functions for tests that use any PSA API.
+ */
+/*  Copyright (C) 2019, ARM Limited, All Rights Reserved
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *  This file is part of mbed TLS (https://tls.mbed.org)
+ */
+
+#ifndef PSA_HELPERS_H
+#define PSA_HELPERS_H
+
+#if defined(MBEDTLS_PSA_CRYPTO_SPM)
+#include "spm/psa_defs.h"
+#endif
+
+/** Evaluate an expression and fail the test case if it returns an error.
+ *
+ * \param expr      The expression to evaluate. This is typically a call
+ *                  to a \c psa_xxx function that returns a value of type
+ *                  #psa_status_t.
+ */
+#define PSA_ASSERT( expr ) TEST_EQUAL( ( expr ), PSA_SUCCESS )
+
+#endif /* PSA_HELPERS_H */
diff --git a/tests/suites/helpers.function b/tests/suites/helpers.function
index 122a17d..e065272 100644
--- a/tests/suites/helpers.function
+++ b/tests/suites/helpers.function
@@ -126,14 +126,6 @@
 #define TEST_EQUAL( expr1, expr2 )              \
     TEST_ASSERT( ( expr1 ) == ( expr2 ) )
 
-/** Evaluate an expression and fail the test case if it returns an error.
- *
- * \param expr      The expression to evaluate. This is typically a call
- *                  to a \c psa_xxx function that returns a value of type
- *                  #psa_status_t.
- */
-#define PSA_ASSERT( expr ) TEST_EQUAL( ( expr ), PSA_SUCCESS )
-
 /** Allocate memory dynamically and fail the test case if this fails.
  *
  * You must set \p pointer to \c NULL before calling this macro and
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index fd923c2..3d38535 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -10,6 +10,11 @@
 #include <limits.h>
 #include <stdint.h>
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#include "mbedtls/psa_util.h"
+#include "psa_crypto_helpers.h"
+#endif
+
 static int rnd_std_rand( void *rng_state, unsigned char *output, size_t len );
 
 #define RSA_KEY_SIZE 512
@@ -67,8 +72,6 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 
-#include "mbedtls/psa_util.h"
-
 /*
  * Generate a key using PSA and return a handle to that key,
  * or 0 if the key generation failed.
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index aaa3189..4441e9b 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1,15 +1,11 @@
 /* BEGIN_HEADER */
 #include <stdint.h>
 
-#if defined(MBEDTLS_PSA_CRYPTO_SPM)
-#include "spm/psa_defs.h"
-#endif
-
 #include "mbedtls/asn1.h"
 #include "mbedtls/asn1write.h"
 #include "mbedtls/oid.h"
 
-#include "psa/crypto.h"
+#include "psa_crypto_helpers.h"
 
 /** An invalid export length that will never be set by psa_export_key(). */
 static const size_t INVALID_EXPORT_LENGTH = ~0U;
@@ -686,43 +682,6 @@
     return( ok );
 }
 
-static int is_oid_of_key_type( psa_key_type_t type,
-                               const uint8_t *oid, size_t oid_length )
-{
-    const uint8_t *expected_oid = NULL;
-    size_t expected_oid_length = 0;
-#if defined(MBEDTLS_RSA_C)
-    if( PSA_KEY_TYPE_IS_RSA( type ) )
-    {
-        expected_oid = (uint8_t *) MBEDTLS_OID_PKCS1_RSA;
-        expected_oid_length = sizeof( MBEDTLS_OID_PKCS1_RSA ) - 1;
-    }
-    else
-#endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_C)
-    if( PSA_KEY_TYPE_IS_ECC( type ) )
-    {
-        expected_oid = (uint8_t *) MBEDTLS_OID_EC_ALG_UNRESTRICTED;
-        expected_oid_length = sizeof( MBEDTLS_OID_EC_ALG_UNRESTRICTED ) - 1;
-    }
-    else
-#endif /* MBEDTLS_ECP_C */
-    {
-        char message[40];
-        mbedtls_snprintf( message, sizeof( message ),
-                          "OID not known for key type=0x%08lx",
-                          (unsigned long) type );
-        test_fail( message, __LINE__, __FILE__ );
-        return( 0 );
-    }
-
-    ASSERT_COMPARE( expected_oid, expected_oid_length, oid, oid_length );
-    return( 1 );
-
-exit:
-    return( 0 );
-}
-
 static int asn1_skip_integer( unsigned char **p, const unsigned char *end,
                               size_t min_bits, size_t max_bits,
                               int must_be_odd )
@@ -762,25 +721,6 @@
     return( 0 );
 }
 
-static int asn1_get_implicit_tag( unsigned char **p, const unsigned char *end,
-                                  size_t *len,
-                                  unsigned char n, unsigned char tag )
-{
-    int ret;
-    ret = mbedtls_asn1_get_tag( p, end, len,
-                                MBEDTLS_ASN1_CONTEXT_SPECIFIC |
-                                MBEDTLS_ASN1_CONSTRUCTED | ( n ) );
-    if( ret != 0 )
-        return( ret );
-    end = *p + *len;
-    ret = mbedtls_asn1_get_tag( p, end, len, tag );
-    if( ret != 0 )
-        return( ret );
-    if( *p + *len != end )
-        return( MBEDTLS_ERR_ASN1_LENGTH_MISMATCH );
-    return( 0 );
-}
-
 static int exported_key_sanity_check( psa_key_type_t type, size_t bits,
                                       uint8_t *exported, size_t exported_length )
 {
@@ -1263,7 +1203,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1295,12 +1235,13 @@
     psa_set_key_type( &attributes, type );
     status = psa_import_key( &attributes, p, length, &handle );
     TEST_EQUAL( status, expected_status );
+
     if( status == PSA_SUCCESS )
         PSA_ASSERT( psa_destroy_key( handle ) );
 
 exit:
     mbedtls_free( buffer );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1393,7 +1334,7 @@
     mbedtls_free( exported );
     mbedtls_free( reexported );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1404,7 +1345,7 @@
     test_operations_on_invalid_handle( handle );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1457,7 +1398,7 @@
     mbedtls_free( exported );
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1499,7 +1440,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &got_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1531,7 +1472,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1611,7 +1552,7 @@
 exit:
     psa_mac_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1654,7 +1595,7 @@
 exit:
     psa_cipher_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1715,7 +1656,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1776,7 +1717,7 @@
 exit:
     psa_destroy_key( handle );
     psa_reset_key_attributes( &attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( buffer );
 }
 /* END_CASE */
@@ -1831,7 +1772,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1870,7 +1811,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1908,7 +1849,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1945,7 +1886,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1982,7 +1923,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2063,7 +2004,7 @@
 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 */
@@ -2104,10 +2045,13 @@
     TEST_EQUAL( psa_copy_key( source_handle,
                               &target_attributes, &target_handle ),
                 expected_status_arg );
+
+    PSA_ASSERT( psa_destroy_key( source_handle ) );
+
 exit:
     psa_reset_key_attributes( &source_attributes );
     psa_reset_key_attributes( &target_attributes );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2169,7 +2113,7 @@
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2257,7 +2201,7 @@
     PSA_ASSERT( psa_hash_abort( &operation ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2292,7 +2236,7 @@
                 PSA_ERROR_INVALID_SIGNATURE );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2314,7 +2258,7 @@
                 PSA_ERROR_BUFFER_TOO_SMALL );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2359,7 +2303,7 @@
     psa_hash_abort( &op_setup );
     psa_hash_abort( &op_finished );
     psa_hash_abort( &op_aborted );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2400,7 +2344,7 @@
     psa_hash_abort( &op_setup );
     psa_hash_abort( &op_finished );
     psa_hash_abort( &op_aborted );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2471,7 +2415,7 @@
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2591,8 +2535,10 @@
                 PSA_ERROR_BAD_STATE );
     PSA_ASSERT( psa_mac_abort( &operation ) );
 
+    PSA_ASSERT( psa_destroy_key( handle ) );
+
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2647,7 +2593,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2685,7 +2631,7 @@
 
 exit:
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2763,7 +2709,7 @@
 #endif
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2912,8 +2858,10 @@
                 PSA_ERROR_BAD_STATE );
     PSA_ASSERT( psa_cipher_abort( &operation ) );
 
+    PSA_ASSERT( psa_destroy_key( handle ) );
+
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -2973,7 +2921,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3041,7 +2989,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3112,7 +3060,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3173,7 +3121,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3255,7 +3203,7 @@
     mbedtls_free( output1 );
     mbedtls_free( output2 );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3358,7 +3306,7 @@
     mbedtls_free( output1 );
     mbedtls_free( output2 );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3434,7 +3382,7 @@
     psa_destroy_key( handle );
     mbedtls_free( output_data );
     mbedtls_free( output_data2 );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3484,7 +3432,7 @@
 exit:
     psa_destroy_key( handle );
     mbedtls_free( output_data );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3540,7 +3488,7 @@
 exit:
     psa_destroy_key( handle );
     mbedtls_free( output_data );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3605,7 +3553,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3650,7 +3598,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3717,7 +3665,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( signature );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3749,7 +3697,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3785,7 +3733,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3854,7 +3802,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3919,7 +3867,7 @@
     psa_destroy_key( handle );
     mbedtls_free( output );
     mbedtls_free( output2 );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3981,7 +3929,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4043,7 +3991,7 @@
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
     mbedtls_free( output );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4111,7 +4059,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4160,7 +4108,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4302,7 +4250,7 @@
     mbedtls_free( output_buffer );
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4383,7 +4331,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4443,7 +4391,7 @@
     psa_reset_key_attributes( &got_attributes );
     psa_destroy_key( base_handle );
     psa_destroy_key( derived_handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4522,7 +4470,7 @@
     psa_key_derivation_abort( &operation );
     psa_destroy_key( base_handle );
     psa_destroy_key( derived_handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4570,7 +4518,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4607,7 +4555,7 @@
 exit:
     mbedtls_free( output );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4667,7 +4615,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4726,7 +4674,7 @@
 exit:
     psa_key_derivation_abort( &operation );
     psa_destroy_key( our_key );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( actual_output );
 }
 /* END_CASE */
@@ -4776,7 +4724,7 @@
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( output );
     mbedtls_free( changed );
 }
@@ -4822,7 +4770,7 @@
 exit:
     psa_reset_key_attributes( &got_attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -4923,7 +4871,7 @@
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_key( handle );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( e_read_buffer );
     mbedtls_free( exported );
 }
@@ -5020,7 +4968,8 @@
     }
 
     /* Shutdown and restart */
-    mbedtls_psa_crypto_free();
+    PSA_ASSERT( psa_close_key( handle ) );
+    PSA_DONE();
     PSA_ASSERT( psa_crypto_init() );
 
     /* Check key slot still contains key data */
@@ -5062,6 +5011,6 @@
         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 91e210e..8538d6d 100644
--- a/tests/suites/test_suite_psa_crypto_entropy.function
+++ b/tests/suites/test_suite_psa_crypto_entropy.function
@@ -1,10 +1,10 @@
 /* BEGIN_HEADER */
 #include <stdint.h>
 
-#include "psa/crypto.h"
 #include "mbedtls/entropy.h"
 #include "mbedtls/entropy_poll.h"
 
+#include "psa_crypto_helpers.h"
 #if defined(MBEDTLS_PSA_ITS_FILE_C)
 #include <stdio.h>
 #else
@@ -77,7 +77,7 @@
 exit:
     mbedtls_free( seed );
     remove_seed_file( );
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -105,12 +105,12 @@
     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 8abd4e2..d50ff5a 100644
--- a/tests/suites/test_suite_psa_crypto_hash.function
+++ b/tests/suites/test_suite_psa_crypto_hash.function
@@ -2,11 +2,7 @@
 
 #include <stdint.h>
 
-#if defined(MBEDTLS_PSA_CRYPTO_SPM)
-#include "spm/psa_defs.h"
-#endif
-
-#include "psa/crypto.h"
+#include "psa_crypto_helpers.h"
 
 /* END_HEADER */
 
@@ -35,7 +31,7 @@
                     actual_hash, actual_hash_length );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -56,7 +52,7 @@
                                  expected_hash->len ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -99,6 +95,6 @@
     } 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 f10a4b2..3c4b42e 100644
--- a/tests/suites/test_suite_psa_crypto_init.function
+++ b/tests/suites/test_suite_psa_crypto_init.function
@@ -1,11 +1,7 @@
 /* BEGIN_HEADER */
 #include <stdint.h>
 
-#if defined(MBEDTLS_PSA_CRYPTO_SPM)
-#include "spm/psa_defs.h"
-#endif
-#include "psa/crypto.h"
-
+#include "psa_crypto_helpers.h"
 /* Some tests in this module configure entropy sources. */
 #include "psa_crypto_invasive.h"
 
@@ -142,7 +138,7 @@
         PSA_ASSERT( status );
         status = psa_crypto_init( );
         PSA_ASSERT( status );
-        mbedtls_psa_crypto_free( );
+        PSA_DONE( );
     }
 }
 /* END_CASE */
@@ -154,9 +150,9 @@
     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 */
 
@@ -170,7 +166,7 @@
     {
         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 );
@@ -190,7 +186,7 @@
     {
         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 );
@@ -216,7 +212,7 @@
     PSA_ASSERT( psa_generate_random( random, sizeof( random ) ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -256,7 +252,7 @@
     PSA_ASSERT( psa_generate_random( random, sizeof( random ) ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -284,6 +280,6 @@
 
 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 0417d84..fc19248 100644
--- a/tests/suites/test_suite_psa_crypto_persistent_key.function
+++ b/tests/suites/test_suite_psa_crypto_persistent_key.function
@@ -1,7 +1,9 @@
 /* BEGIN_HEADER */
 #include <stdint.h>
-#include "psa/crypto.h"
+
+#include "psa_crypto_helpers.h"
 #include "psa_crypto_storage.h"
+
 #include "mbedtls/md.h"
 
 #define PSA_KEY_STORAGE_MAGIC_HEADER "PSA\0KEY"
@@ -106,9 +108,12 @@
     TEST_EQUAL( psa_import_key( &attributes, data, data_length, &handle ),
                 expected_status );
 
+    if( expected_status == PSA_SUCCESS )
+        PSA_ASSERT( psa_destroy_key( handle ) );
+
 exit:
     mbedtls_free( data );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
@@ -135,7 +140,7 @@
     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 ) );
     }
@@ -150,7 +155,7 @@
     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 */
@@ -159,8 +164,10 @@
     PSA_ASSERT( psa_import_key( &attributes, second_data->x, second_data->len,
                                 &handle ) );
 
+    PSA_ASSERT( psa_destroy_key( handle ) );
+
 exit:
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
     psa_destroy_persistent_key( key_id );
 }
 /* END_CASE */
@@ -190,7 +197,7 @@
     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 ) );
     }
@@ -204,10 +211,12 @@
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
 
+    PSA_ASSERT( psa_destroy_key( handle ) );
+
 exit:
     psa_reset_key_attributes( &attributes );
     psa_destroy_persistent_key( key_id );
-    mbedtls_psa_crypto_free();
+    PSA_DONE();
 }
 /* END_CASE */
 
@@ -239,7 +248,7 @@
     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 ) );
     }
@@ -274,7 +283,7 @@
 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.data b/tests/suites/test_suite_psa_crypto_slot_management.data
index e65befe..233b166 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.data
+++ b/tests/suites/test_suite_psa_crypto_slot_management.data
@@ -1,19 +1,31 @@
 Transient slot, check after closing
 transient_slot_lifecycle:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_CLOSE
 
+Transient slot, check after closing and restarting
+transient_slot_lifecycle:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_CLOSE_WITH_SHUTDOWN
+
 Transient slot, check after destroying
 transient_slot_lifecycle:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_DESTROY
 
-Transient slot, check after restart
+Transient slot, check after destroying and restarting
+transient_slot_lifecycle:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_DESTROY_WITH_SHUTDOWN
+
+Transient slot, check after restart with live handles
 transient_slot_lifecycle:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_SHUTDOWN
 
 Persistent slot, check after closing, id=min
 persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_ID_USER_MIN:0:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_CLOSE
 
+Persistent slot, check after closing and restarting, id=min
+persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_ID_USER_MIN:0:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_CLOSE
+
 Persistent slot, check after destroying, id=min
 persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_ID_USER_MIN:0:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_DESTROY
 
-Persistent slot, check after restart, id=min
+Persistent slot, check after destroying and restarting, id=min
+persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_ID_USER_MIN:0:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_DESTROY
+
+Persistent slot, check after restart with live handle, id=min
 persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_ID_USER_MIN:0:0:0:PSA_KEY_TYPE_RAW_DATA:"0123456789abcdef0123456789abcdef":CLOSE_BY_SHUTDOWN
 
 Persistent slot, check after closing, id=max
@@ -29,6 +41,10 @@
 depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
 persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_ECDSA_ANY:0:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":CLOSE_BY_CLOSE
 
+Persistent slot: ECP keypair (ECDSA, exportable); close+restart
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_ECDSA_ANY:0:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":CLOSE_BY_CLOSE_WITH_SHUTDOWN
+
 Persistent slot: ECP keypair (ECDSA, exportable); restart
 depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
 persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_ECDSA_ANY:0:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":CLOSE_BY_SHUTDOWN
@@ -37,6 +53,10 @@
 depends_on:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
 persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_ALG_ECDSA_ANY:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":CLOSE_BY_CLOSE
 
+Persistent slot: ECP keypair (ECDH+ECDSA, exportable); close+restart
+depends_on:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_ALG_ECDSA_ANY:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":CLOSE_BY_CLOSE_WITH_SHUTDOWN
+
 Persistent slot: ECP keypair (ECDH+ECDSA, exportable); restart
 depends_on:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
 persistent_slot_lifecycle:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN | PSA_KEY_USAGE_VERIFY:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_ALG_ECDSA_ANY:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":CLOSE_BY_SHUTDOWN
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index db46328..3b9eada 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -1,18 +1,16 @@
 /* BEGIN_HEADER */
 #include <stdint.h>
 
-#if defined(MBEDTLS_PSA_CRYPTO_SPM)
-#include "spm/psa_defs.h"
-#endif
-#include "psa/crypto.h"
-
+#include "psa_crypto_helpers.h"
 #include "psa_crypto_storage.h"
 
 typedef enum
 {
-    CLOSE_BY_CLOSE,
-    CLOSE_BY_DESTROY,
-    CLOSE_BY_SHUTDOWN,
+    CLOSE_BY_CLOSE, /**< Close the handle(s). */
+    CLOSE_BY_DESTROY, /**< Destroy the handle(s). */
+    CLOSE_BY_SHUTDOWN, /**< Deinit and reinit without closing handles. */
+    CLOSE_BY_CLOSE_WITH_SHUTDOWN, /**< Close handle(s) then deinit/reinit. */
+    CLOSE_BY_DESTROY_WITH_SHUTDOWN, /**< Destroy handle(s) then deinit/reinit. */
 } close_method_t;
 
 typedef enum
@@ -66,6 +64,58 @@
 #define TEST_USES_KEY_ID( key_id ) ( (void) ( key_id ) )
 #endif /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
 
+/** Apply \p close_method to invalidate the specified handles:
+ * close it, destroy it, or do nothing;
+ */
+static int invalidate_handle( close_method_t close_method,
+                              psa_key_handle_t handle )
+{
+    switch( close_method )
+    {
+        case CLOSE_BY_CLOSE:
+        case CLOSE_BY_CLOSE_WITH_SHUTDOWN:
+            PSA_ASSERT( psa_close_key( handle ) );
+            break;
+        case CLOSE_BY_DESTROY:
+        case CLOSE_BY_DESTROY_WITH_SHUTDOWN:
+            PSA_ASSERT( psa_destroy_key( handle ) );
+            break;
+        case CLOSE_BY_SHUTDOWN:
+            break;
+    }
+    return( 1 );
+exit:
+    return( 0 );
+}
+
+/** Restart the PSA subsystem if \p close_method says so. */
+static int invalidate_psa( close_method_t close_method )
+{
+    switch( close_method )
+    {
+        case CLOSE_BY_CLOSE:
+        case CLOSE_BY_DESTROY:
+            return( 1 );
+        case CLOSE_BY_CLOSE_WITH_SHUTDOWN:
+        case CLOSE_BY_DESTROY_WITH_SHUTDOWN:
+            /* All keys must have been closed. */
+            PSA_DONE( );
+            break;
+        case CLOSE_BY_SHUTDOWN:
+            /* Some keys may remain behind, and we're testing that this
+             * properly closes them. */
+            mbedtls_psa_crypto_free( );
+            break;
+    }
+
+    PSA_ASSERT( psa_crypto_init( ) );
+    ASSERT_PSA_PRISTINE( );
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -98,19 +148,10 @@
     TEST_EQUAL( psa_get_key_type( &attributes ), type );
 
     /* Do something that invalidates the handle. */
-    switch( close_method )
-    {
-        case CLOSE_BY_CLOSE:
-            PSA_ASSERT( psa_close_key( handle ) );
-            break;
-        case CLOSE_BY_DESTROY:
-            PSA_ASSERT( psa_destroy_key( handle ) );
-            break;
-        case CLOSE_BY_SHUTDOWN:
-            mbedtls_psa_crypto_free( );
-            PSA_ASSERT( psa_crypto_init( ) );
-            break;
-    }
+    if( ! invalidate_handle( close_method, handle ) )
+        goto exit;
+    if( ! invalidate_psa( close_method ) )
+        goto exit;
 
     /* Test that the handle is now invalid. */
     TEST_EQUAL( psa_get_key_attributes( handle, &attributes ),
@@ -118,7 +159,7 @@
     TEST_EQUAL( psa_close_key( handle ), PSA_ERROR_INVALID_HANDLE );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -175,19 +216,11 @@
     TEST_EQUAL( psa_get_key_type( &attributes ), type );
 
     /* Do something that invalidates the handle. */
-    switch( close_method )
-    {
-        case CLOSE_BY_CLOSE:
-            PSA_ASSERT( psa_close_key( handle ) );
-            break;
-        case CLOSE_BY_DESTROY:
-            PSA_ASSERT( psa_destroy_key( handle ) );
-            break;
-        case CLOSE_BY_SHUTDOWN:
-            mbedtls_psa_crypto_free( );
-            PSA_ASSERT( psa_crypto_init( ) );
-            break;
-    }
+    if( ! invalidate_handle( close_method, handle ) )
+        goto exit;
+    if( ! invalidate_psa( close_method ) )
+        goto exit;
+
     /* Test that the handle is now invalid. */
     TEST_EQUAL( psa_get_key_attributes( handle, &read_attributes ),
                 PSA_ERROR_INVALID_HANDLE );
@@ -200,6 +233,7 @@
     switch( close_method )
     {
         case CLOSE_BY_CLOSE:
+        case CLOSE_BY_CLOSE_WITH_SHUTDOWN:
         case CLOSE_BY_SHUTDOWN:
             PSA_ASSERT( psa_open_key( id, &handle ) );
             PSA_ASSERT( psa_get_key_attributes( handle, &read_attributes ) );
@@ -232,15 +266,18 @@
                                             &reexported_length ),
                             PSA_ERROR_NOT_PERMITTED );
             }
+            PSA_ASSERT( psa_close_key( handle ) );
             break;
+
         case CLOSE_BY_DESTROY:
+        case CLOSE_BY_DESTROY_WITH_SHUTDOWN:
             TEST_EQUAL( psa_open_key( id, &handle ),
                         PSA_ERROR_DOES_NOT_EXIST );
             break;
     }
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_purge_key_storage( );
     mbedtls_free( reexported );
 }
@@ -306,8 +343,10 @@
     ASSERT_COMPARE( material1, sizeof( material1 ),
                     reexported, reexported_length );
 
+    PSA_ASSERT( psa_close_key( handle1 ) );
+
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     psa_purge_key_storage( );
 }
 /* END_CASE */
@@ -326,7 +365,7 @@
     TEST_EQUAL( handle, 0 );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -354,7 +393,7 @@
     TEST_EQUAL( handle, 0 );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
 #endif
@@ -467,8 +506,10 @@
                     PSA_ERROR_NOT_PERMITTED );
     }
 
+    PSA_ASSERT( psa_destroy_key( target_handle ) );
+
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
@@ -570,8 +611,12 @@
                         export_buffer, length );
     }
 
+    PSA_ASSERT( psa_destroy_key( source_handle ) );
+    if( target_handle != source_handle )
+        PSA_ASSERT( psa_destroy_key( target_handle ) );
+
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( export_buffer );
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_purge_key_storage( );
@@ -613,7 +658,7 @@
     PSA_ASSERT( psa_close_key( handle1 ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -661,7 +706,7 @@
     PSA_ASSERT( psa_close_key( handles[i - 1] ) );
 
 exit:
-    mbedtls_psa_crypto_free( );
+    PSA_DONE( );
     mbedtls_free( handles );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_its.function b/tests/suites/test_suite_psa_its.function
index 867f64f..8b15005 100644
--- a/tests/suites/test_suite_psa_its.function
+++ b/tests/suites/test_suite_psa_its.function
@@ -1,6 +1,8 @@
 /* BEGIN_HEADER */
 #include "../library/psa_crypto_its.h"
 
+#include "psa_helpers.h"
+
 /* Internal definitions of the implementation, copied for the sake of
  * some of the tests and of the cleanup code. */
 #define PSA_ITS_STORAGE_PREFIX ""