psa: free RNG implementation before checking for remaining open key slots

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index ef29b77..8005dcb 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -194,6 +194,15 @@
 /**@}*/
 
 /**
+ * \brief PSA random deinitialization.
+ *
+ * This function frees the RNG implementation used by PSA.
+ *
+ * This is an Mbed TLS extension.
+ */
+void mbedtls_psa_random_free(void);
+
+/**
  * \brief Library deinitialization.
  *
  * This function clears all data associated with the PSA layer,
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 1149940..6caab03 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7327,14 +7327,16 @@
 
 /** Deinitialize the PSA random generator.
  */
-static void mbedtls_psa_random_free(mbedtls_psa_random_context_t *rng)
+void mbedtls_psa_random_free(void)
 {
+    if (global_data.rng_state != RNG_NOT_INITIALIZED) {
 #if defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG)
-    memset(rng, 0, sizeof(*rng));
+        memset(&global_data.rng, 0, sizeof(global_data.rng));
 #else /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
-    mbedtls_psa_drbg_free(MBEDTLS_PSA_RANDOM_STATE);
-    rng->entropy_free(&rng->entropy);
+        mbedtls_psa_drbg_free(MBEDTLS_PSA_RANDOM_STATE);
+        global_data.rng.entropy_free(&global_data.rng.entropy);
 #endif /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
+    }
 }
 
 /** Seed the PSA random generator.
@@ -7661,9 +7663,7 @@
 void mbedtls_psa_crypto_free(void)
 {
     psa_wipe_all_key_slots();
-    if (global_data.rng_state != RNG_NOT_INITIALIZED) {
-        mbedtls_psa_random_free(&global_data.rng);
-    }
+    mbedtls_psa_random_free();
     /* Wipe all remaining data, including configuration.
      * In particular, this sets all state indicator to the value
      * indicating "uninitialized". */
@@ -7714,6 +7714,11 @@
     }
     global_data.drivers_initialized = 1;
 
+    status = psa_initialize_key_slots();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
     /* Initialize and seed the random generator. */
     mbedtls_psa_random_init(&global_data.rng);
     global_data.rng_state = RNG_INITIALIZED;
@@ -7723,11 +7728,6 @@
     }
     global_data.rng_state = RNG_SEEDED;
 
-    status = psa_initialize_key_slots();
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
 #if defined(PSA_CRYPTO_STORAGE_HAS_TRANSACTIONS)
     status = psa_crypto_load_transaction();
     if (status == PSA_SUCCESS) {
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 04b90b9..f4c49fb 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -34,6 +34,7 @@
 #define PSA_DONE()                                                      \
     do                                                                  \
     {                                                                   \
+        mbedtls_psa_random_free();                                      \
         mbedtls_test_fail_if_psa_leaking(__LINE__, __FILE__);           \
         mbedtls_test_psa_purge_key_storage();                           \
         mbedtls_psa_crypto_free();                                      \
@@ -125,17 +126,21 @@
 
 /** Shut down the PSA Crypto subsystem, allowing persistent keys to survive.
  * Expect a clean shutdown, with no slots in use.
+ * mbedtls_psa_random_free() is called before any check for remaining open
+ * keys because when AES_C is not defined, CTR_DRBG relies on PSA to perform
+ * AES-ECB so it holds an open AES key for that since psa_crypto_init().
  *
  * If some key slots are still in use, record the test case as failed and
  * jump to the `exit` label.
  */
 #define PSA_SESSION_DONE()                                             \
-    do                                                                  \
-    {                                                                   \
+    do                                                                 \
+    {                                                                  \
+        mbedtls_psa_random_free();                                     \
         mbedtls_test_psa_purge_key_cache();                            \
         ASSERT_PSA_PRISTINE();                                         \
         mbedtls_psa_crypto_free();                                     \
-    }                                                                   \
+    }                                                                  \
     while (0)