From 5386f6ba071a36e105b52cd5c6dc3efbc07633c0 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Thu, 1 Aug 2019 12:47:40 +0200
Subject: [PATCH] Fix PSA init/deinit in mbedtls_xxx tests when using PSA

In tests of mbedtls_cipher_xxx and mbedtls_pk_xxx with
MBEDTLS_USE_PSA_CRYPTO enabled, initialize and deinitialize the PSA
subsystem in every function. Before, the tests were only passing
because the first function to be called happened to call
psa_crypto_init() but not mbedtls_psa_crypto_free(). In some
configurations (not tested on CI), psa_crypto_init() was not called so
the tests using PSA failed.

Call PSA_DONE() at the end of each test function. This ensures that no
resources are leaked in the form of PSA crypto slot contents.
Incidentally, this also fixes a build error due to
test_helper_psa_done() being unused in test_suite_pk: the fact that it
wasn't used betrayed the missing calls to PSA_DONE().
---
 tests/suites/test_suite_cipher.function | 17 +++++++++++++----
 tests/suites/test_suite_pk.function     | 18 ++++++++++++++++--
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/tests/suites/test_suite_cipher.function b/tests/suites/test_suite_cipher.function
index 1ea14088b6..f6367f1758 100644
--- a/tests/suites/test_suite_cipher.function
+++ b/tests/suites/test_suite_cipher.function
@@ -4,6 +4,11 @@
 #if defined(MBEDTLS_GCM_C)
 #include "mbedtls/gcm.h"
 #endif
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#include "psa_crypto_helpers.h"
+#endif
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -982,7 +987,7 @@ void auth_crypt_tv( int cipher_id, data_t * key, data_t * iv,
 #else
     if( use_psa == 1 )
     {
-        TEST_ASSERT( psa_crypto_init() == 0 );
+        PSA_ASSERT( psa_crypto_init( ) );
 
         /* PSA requires that the tag immediately follows the ciphertext. */
         tmp_cipher = mbedtls_calloc( 1, cipher->len + tag->len );
@@ -1066,14 +1071,15 @@ void auth_crypt_tv( int cipher_id, data_t * key, data_t * iv,
 
 exit:
 
+    mbedtls_cipher_free( &ctx );
+
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if( use_psa == 1 )
     {
         mbedtls_free( tmp_cipher );
+        PSA_DONE( );
     }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
-
-    mbedtls_cipher_free( &ctx );
 }
 /* END_CASE */
 
@@ -1143,7 +1149,7 @@ void test_vec_crypt( int cipher_id, int operation, char *hex_key,
 #else
     if( use_psa == 1 )
     {
-        TEST_ASSERT( psa_crypto_init() == 0 );
+        PSA_ASSERT( psa_crypto_init( ) );
         TEST_ASSERT( 0 == mbedtls_cipher_setup_psa( &ctx,
                               mbedtls_cipher_info_from_type( cipher_id ), 0 ) );
     }
@@ -1172,6 +1178,9 @@ void test_vec_crypt( int cipher_id, int operation, char *hex_key,
 
 exit:
     mbedtls_cipher_free( &ctx );
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    PSA_DONE( );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 }
 /* END_CASE */
 
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 3d38535e32..fbb69073e5 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -13,6 +13,13 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 #include "mbedtls/psa_util.h"
 #include "psa_crypto_helpers.h"
+#define PSA_INIT( ) PSA_ASSERT( psa_crypto_init( ) )
+#else
+/* Define empty macros so that we can use them in the preamble and teardown
+ * of every test function that uses PSA conditionally based on
+ * MBEDTLS_USE_PSA_CRYPTO. */
+#define PSA_INIT( ) ( (void) 0 )
+#define PSA_DONE( ) ( (void) 0 )
 #endif
 
 static int rnd_std_rand( void *rng_state, unsigned char *output, size_t len );
@@ -117,7 +124,7 @@ void pk_psa_utils(  )
     size_t len;
     mbedtls_pk_debug_item dbg;
 
-    TEST_ASSERT( psa_crypto_init() == 0 );
+    PSA_ASSERT( psa_crypto_init( ) );
 
     mbedtls_pk_init( &pk );
     mbedtls_pk_init( &pk2 );
@@ -173,6 +180,7 @@ void pk_psa_utils(  )
 exit:
     mbedtls_pk_free( &pk ); /* redundant except upon error */
     mbedtls_pk_free( &pk2 );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -763,7 +771,7 @@ void pk_ec_test_vec( int type, int id, data_t * key, data_t * hash,
     mbedtls_ecp_keypair *eckey;
 
     mbedtls_pk_init( &pk );
-
+    PSA_INIT( );
 
     TEST_ASSERT( mbedtls_pk_setup( &pk, mbedtls_pk_info_from_type( type ) ) == 0 );
 
@@ -780,6 +788,7 @@ void pk_ec_test_vec( int type, int id, data_t * key, data_t * hash,
 
 exit:
     mbedtls_pk_free( &pk );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -904,6 +913,7 @@ void pk_sign_verify( int type, int sign_ret, int verify_ret )
 #endif
 
     mbedtls_pk_init( &pk );
+    PSA_INIT( );
 
     memset( hash, 0x2a, sizeof hash );
     memset( sig, 0, sizeof sig );
@@ -955,6 +965,7 @@ exit:
     mbedtls_pk_restart_free( rs_ctx );
 #endif
     mbedtls_pk_free( &pk );
+    PSA_DONE( );
 }
 /* END_CASE */
 
@@ -1210,6 +1221,8 @@ void pk_psa_sign(  )
      * - parse it to a PK context and verify the signature this way
      */
 
+    PSA_ASSERT( psa_crypto_init( ) );
+
     /* Create legacy EC public/private key in PK context. */
     mbedtls_pk_init( &pk );
     TEST_ASSERT( mbedtls_pk_setup( &pk,
@@ -1259,5 +1272,6 @@ void pk_psa_sign(  )
 
 exit:
     mbedtls_pk_free( &pk );
+    PSA_DONE( );
 }
 /* END_CASE */