Don't call TEST_ASSERT in PSA_DONE

TEST_ASSERT jumps to the exit label, so it must not be called from
cleanup code executed after the exit label. It's legitimate (and
indeed very common) to call PSA_DONE in cleanup code, so PSA_DONE must
not jump to exit.

Define an auxiliary function test_fail_if_psa_leaking() that calls
test_fail() with the error message provided by
mbedtls_test_helper_is_psa_leaking(). This function currently needs to
be in helpers.function rather than in a PSA-specific helper file
because it calls test_fail which is defined in helpers.function.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 3e60a9b..b8eb4aa 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -36,19 +36,29 @@
 const char *mbedtls_test_helper_is_psa_leaking( void );
 
 /** Check that no PSA Crypto key slots are in use.
+ *
+ * If any slots are in use, mark the current test as failed and jump to
+ * the exit label. This is equivalent to
+ * `TEST_ASSERT( ! mbedtls_test_helper_is_psa_leaking( ) )`
+ * but with a more informative message.
  */
-#define ASSERT_PSA_PRISTINE( )                                  \
-    TEST_ASSERT( ! mbedtls_test_helper_is_psa_leaking( ) )
+#define ASSERT_PSA_PRISTINE( )                                          \
+    do                                                                  \
+    {                                                                   \
+        if( test_fail_if_psa_leaking( __LINE__, __FILE__ ) )            \
+            goto exit;                                                  \
+    }                                                                   \
+    while( 0 )
 
 /** Shut down the PSA Crypto subsystem. Expect a clean shutdown, with no slots
  * in use.
  */
-#define PSA_DONE( )                                     \
-    do                                                  \
-    {                                                   \
-        ASSERT_PSA_PRISTINE( );                         \
-        mbedtls_psa_crypto_free( );                     \
-    }                                                   \
+#define PSA_DONE( )                                                     \
+    do                                                                  \
+    {                                                                   \
+        test_fail_if_psa_leaking( __LINE__, __FILE__ );                 \
+        mbedtls_psa_crypto_free( );                                     \
+    }                                                                   \
     while( 0 )
 
 
@@ -60,8 +70,8 @@
  * disabled by default.
  *
  * When MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG is enabled and the test
- * helpers are linked into a program, you must enable this before any code
- * that uses the PSA subsystem to generate random data (including internal
+ * helpers are linked into a program, you must enable this before running any
+ * code that uses the PSA subsystem to generate random data (including internal
  * random generation for purposes such as blinding when the random generation
  * is routed through PSA).
  *
diff --git a/tests/suites/helpers.function b/tests/suites/helpers.function
index 3a9c426..1dc6721 100644
--- a/tests/suites/helpers.function
+++ b/tests/suites/helpers.function
@@ -421,6 +421,26 @@
     test_info.filename = filename;
 }
 
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+/** Check that no PSA Crypto key slots are in use.
+ *
+ * If any slots are in use, mark the current test as failed.
+ *
+ * \return 0 if the key store is empty, 1 otherwise.
+ */
+int test_fail_if_psa_leaking( int line_no, const char *filename )
+{
+    const char *msg = mbedtls_test_helper_is_psa_leaking( );
+    if( msg == NULL )
+        return 0;
+    else
+    {
+        test_fail( msg, line_no, filename );
+        return 1;
+    }
+}
+#endif /* defined(MBEDTLS_PSA_CRYPTO_C) */
+
 #if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
 static int redirect_output( FILE* out_stream, const char* path )
 {