From 9ba7691bf721117eb133b33a1f5d032696698abc Mon Sep 17 00:00:00 2001
From: Ronald Cron <ronald.cron@arm.com>
Date: Sat, 10 Apr 2021 16:57:30 +0200
Subject: [PATCH] psa: Add driver initialization and termination

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
---
 library/psa_crypto.c                          | 13 +++----
 library/psa_crypto_driver_wrappers.c          | 38 +++++++++++++++++++
 library/psa_crypto_driver_wrappers.h          |  6 +++
 tests/include/test/drivers/key_management.h   |  5 +++
 .../src/drivers/test_driver_key_management.c  | 19 ++++++++++
 5 files changed, 73 insertions(+), 8 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c84c17955..ffe659a40 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5818,11 +5818,9 @@ void mbedtls_psa_crypto_free( void )
      * In particular, this sets all state indicator to the value
      * indicating "uninitialized". */
     mbedtls_platform_zeroize( &global_data, sizeof( global_data ) );
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    /* Unregister all secure element drivers, so that we restart from
-     * a pristine state. */
-    psa_unregister_all_se_drivers( );
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
+    /* Terminate drivers */
+    psa_driver_wrapper_free( );
 }
 
 #if defined(PSA_CRYPTO_STORAGE_HAS_TRANSACTIONS)
@@ -5871,11 +5869,10 @@ psa_status_t psa_crypto_init( void )
     if( status != PSA_SUCCESS )
         goto exit;
 
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    status = psa_init_all_se_drivers( );
+    /* Init drivers */
+    status = psa_driver_wrapper_init( );
     if( status != PSA_SUCCESS )
         goto exit;
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
 #if defined(PSA_CRYPTO_STORAGE_HAS_TRANSACTIONS)
     status = psa_crypto_load_transaction( );
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index d6e5778df..96797159a 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -67,6 +67,44 @@
 #include "psa_crypto_se.h"
 #endif
 
+psa_status_t psa_driver_wrapper_init( void )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    status = psa_init_all_se_drivers( );
+    if( status != PSA_SUCCESS )
+        return( status );
+#endif
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    status = mbedtls_test_transparent_init( );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    status = mbedtls_test_opaque_init( );
+    if( status != PSA_SUCCESS )
+        return( status );
+#endif
+
+    (void) status;
+    return( PSA_SUCCESS );
+}
+
+void psa_driver_wrapper_free( void )
+{
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    /* Unregister all secure element drivers, so that we restart from
+     * a pristine state. */
+    psa_unregister_all_se_drivers( );
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    mbedtls_test_transparent_free( );
+    mbedtls_test_opaque_free( );
+#endif
+}
+
 /* Start delegation functions */
 psa_status_t psa_driver_wrapper_sign_message(
     const psa_key_attributes_t *attributes,
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index 0873b738d..6026b8220 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -25,6 +25,12 @@
 #include "psa/crypto.h"
 #include "psa/crypto_driver_common.h"
 
+/*
+ * Initialization and termination functions
+ */
+psa_status_t psa_driver_wrapper_init( void );
+void psa_driver_wrapper_free( void );
+
 /*
  * Signature functions
  */
diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index d147568cd..5bba61174 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -66,6 +66,11 @@ size_t mbedtls_test_opaque_size_function(
 extern mbedtls_test_driver_key_management_hooks_t
     mbedtls_test_driver_key_management_hooks;
 
+psa_status_t mbedtls_test_transparent_init( void );
+void mbedtls_test_transparent_free( void );
+psa_status_t mbedtls_test_opaque_init( void );
+void mbedtls_test_opaque_free( void );
+
 psa_status_t mbedtls_test_transparent_generate_key(
     const psa_key_attributes_t *attributes,
     uint8_t *key, size_t key_size, size_t *key_length );
diff --git a/tests/src/drivers/test_driver_key_management.c b/tests/src/drivers/test_driver_key_management.c
index a2e637aea..59a1ce4ef 100644
--- a/tests/src/drivers/test_driver_key_management.c
+++ b/tests/src/drivers/test_driver_key_management.c
@@ -56,6 +56,25 @@ const uint8_t mbedtls_test_driver_ecdsa_pubkey[65] =
       0xbc, 0x25, 0x16, 0xc3, 0xd2, 0x70, 0x2d, 0x79,
       0x2f, 0x13, 0x1a, 0x92, 0x20, 0x95, 0xfd, 0x6c };
 
+psa_status_t mbedtls_test_transparent_init( void )
+{
+    return( PSA_SUCCESS );
+}
+
+void mbedtls_test_transparent_free( void )
+{
+    return;
+}
+
+psa_status_t mbedtls_test_opaque_init( void )
+{
+    return( PSA_SUCCESS );
+}
+
+void mbedtls_test_opaque_free( void )
+{
+    return;
+}
 
 /*
  * This macro returns the base size for the key context when SE does not