Merge pull request #7461 from valeriosetti/issue7460-part1

Fixing USE_PSA_INIT/DONE in SSL/X509/PK test suites
diff --git a/docs/use-psa-crypto.md b/docs/use-psa-crypto.md
index 9d783d5..92d0985 100644
--- a/docs/use-psa-crypto.md
+++ b/docs/use-psa-crypto.md
@@ -13,7 +13,8 @@
 
 **Application code:** when this option is enabled, you need to call
 `psa_crypto_init()` before calling any function from the SSL/TLS, X.509 or PK
-module.
+modules, except for the various mbedtls_xxx_init() functions which can be called
+at any time.
 
 **Why enable this option:** to fully take advantage of PSA drivers in PK,
 X.509 and TLS. For example, enabling this option is what allows use of drivers
diff --git a/include/mbedtls/mbedtls_config.h b/include/mbedtls/mbedtls_config.h
index 6158850..bc0f502 100644
--- a/include/mbedtls/mbedtls_config.h
+++ b/include/mbedtls/mbedtls_config.h
@@ -1931,7 +1931,8 @@
  * break backwards compatibility.
  *
  * \warning If you enable this option, you need to call `psa_crypto_init()`
- * before calling any function from the SSL/TLS, X.509 or PK modules.
+ * before calling any function from the SSL/TLS, X.509 or PK modules, except
+ * for the various mbedtls_xxx_init() functions which can be called at any time.
  *
  * \note An important and desirable effect of this option is that it allows
  * PK, X.509 and TLS to take advantage of PSA drivers. For example, enabling
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 5a4e77f..0591029 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -216,10 +216,9 @@
     size_t len;
     mbedtls_pk_debug_item dbg;
 
-    PSA_ASSERT(psa_crypto_init());
-
     mbedtls_pk_init(&pk);
     mbedtls_pk_init(&pk2);
+    USE_PSA_INIT();
 
     TEST_ASSERT(psa_crypto_init() == PSA_SUCCESS);
 
@@ -314,9 +313,8 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    PSA_ASSERT(psa_crypto_init());
-
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     if (opaque_key == 1) {
         psa_set_key_usage_flags(&attributes, key_usage);
@@ -362,6 +360,7 @@
     size_t buf_size = sizeof(buf);
 
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     TEST_EQUAL(MBEDTLS_ERR_PK_BAD_INPUT_DATA,
                mbedtls_pk_verify_restartable(&ctx, MBEDTLS_MD_NONE,
@@ -397,6 +396,7 @@
                                            NULL));
 exit:
     mbedtls_pk_free(&ctx);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -409,6 +409,7 @@
     void *options = NULL;
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, NULL) ==
                 MBEDTLS_ERR_PK_BAD_INPUT_DATA);
@@ -484,6 +485,7 @@
     TEST_ASSERT(mbedtls_pk_parse_public_key(&pk, NULL, 0) ==
                 MBEDTLS_ERR_PK_KEY_INVALID_FORMAT);
 #endif /* MBEDTLS_PK_PARSE_C */
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -494,6 +496,8 @@
 
     /* For the write tests to be effective, we need a valid key pair. */
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_pk_parse_key(&pk,
                                      key_data->x, key_data->len, NULL, 0,
                                      mbedtls_test_rnd_std_rand, NULL) == 0);
@@ -514,6 +518,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -522,8 +527,8 @@
 {
     mbedtls_pk_context pk;
 
-    USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(type)) == 0);
     TEST_ASSERT(pk_genkey(&pk, parameter) == 0);
@@ -545,11 +550,10 @@
 {
     mbedtls_pk_context pub, prv, alt;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pub);
     mbedtls_pk_init(&prv);
     mbedtls_pk_init(&alt);
+    USE_PSA_INIT();
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     /* mbedtls_pk_check_pair() returns either PK or ECP error codes depending
@@ -604,10 +608,8 @@
     mbedtls_ecp_set_max_ops(1);
 #endif
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
-
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
     rsa = mbedtls_pk_rsa(pk);
@@ -646,8 +648,8 @@
     void *options;
     int ret;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_pk_init(&pk);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
     rsa = mbedtls_pk_rsa(pk);
@@ -735,11 +737,11 @@
     unsigned char sig[MBEDTLS_ECDSA_MAX_LEN];
     size_t slen;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_restart_init(&rs_ctx);
     mbedtls_pk_init(&prv);
     mbedtls_pk_init(&pub);
+    USE_PSA_INIT();
+
     memset(sig, 0, sizeof(sig));
 
     TEST_ASSERT(mbedtls_pk_setup(&prv, mbedtls_pk_info_from_type(pk_type)) == 0);
@@ -915,12 +917,11 @@
     mbedtls_pk_init(&pk);
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P);
     mbedtls_mpi_init(&Q); mbedtls_mpi_init(&E);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
     memset(output,     0, sizeof(output));
 
-    USE_PSA_INIT();
-
     /* encryption test */
 
     /* init pk-rsa context */
@@ -988,15 +989,13 @@
     mbedtls_pk_context pk;
     size_t olen;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P);
     mbedtls_mpi_init(&Q); mbedtls_mpi_init(&E);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
 
-
     /* init pk-rsa context */
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
     rsa = mbedtls_pk_rsa(pk);
@@ -1045,11 +1044,10 @@
     mbedtls_svc_key_id_t key_id;
     size_t olen;
 
-    USE_PSA_INIT();
-
     mbedtls_pk_init(&pk);
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P);
     mbedtls_mpi_init(&Q); mbedtls_mpi_init(&E);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
 
@@ -1107,6 +1105,7 @@
     int ret = MBEDTLS_ERR_PK_TYPE_MISMATCH;
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
     memset(output,     0, sizeof(output));
@@ -1124,6 +1123,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1131,14 +1131,15 @@
 void pk_rsa_overflow()
 {
     mbedtls_pk_context pk;
-    size_t hash_len = SIZE_MAX, sig_len = SIZE_MAX;
+    size_t hash_len = UINT_MAX + 1, sig_len = UINT_MAX + 1;
     unsigned char hash[50], sig[100];
 
+    mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
+
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
 
-    mbedtls_pk_init(&pk);
-
     TEST_ASSERT(mbedtls_pk_setup(&pk,
                                  mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
 
@@ -1158,6 +1159,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1177,10 +1179,10 @@
     size_t sig_len, ciph_len, test_len;
     int ret = MBEDTLS_ERR_PK_TYPE_MISMATCH;
 
-    USE_PSA_INIT();
-
     mbedtls_rsa_init(&raw);
-    mbedtls_pk_init(&rsa); mbedtls_pk_init(&alt);
+    mbedtls_pk_init(&rsa);
+    mbedtls_pk_init(&alt);
+    USE_PSA_INIT();
 
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
@@ -1274,12 +1276,12 @@
      * - parse it to a PK context and verify the signature this way
      */
 
-    PSA_ASSERT(psa_crypto_init());
+    mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
 #if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
     if (PSA_KEY_TYPE_IS_RSA(psa_type_arg)) {
         /* Create legacy RSA public/private key in PK context. */
-        mbedtls_pk_init(&pk);
         TEST_ASSERT(mbedtls_pk_setup(&pk,
                                      mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
         TEST_ASSERT(mbedtls_rsa_gen_key(mbedtls_pk_rsa(pk),
@@ -1293,7 +1295,6 @@
         mbedtls_ecp_group_id grpid = parameter_arg;
 
         /* Create legacy EC public/private key in PK context. */
-        mbedtls_pk_init(&pk);
         TEST_ASSERT(mbedtls_pk_setup(&pk,
                                      mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
         TEST_ASSERT(pk_genkey(&pk, grpid) == 0);
@@ -1434,8 +1435,8 @@
     TEST_ASSERT(mbedtls_pk_verify_ext(key_pk_type, options, &pk, md_alg,
                                       hash, hash_len, sig, sig_len) == 0);
 exit:
-    PSA_DONE();
     mbedtls_pk_free(&pk);
+    PSA_DONE();
 }
 /* END_CASE */
 
diff --git a/tests/suites/test_suite_pkparse.function b/tests/suites/test_suite_pkparse.function
index 751482a..838a7db 100644
--- a/tests/suites/test_suite_pkparse.function
+++ b/tests/suites/test_suite_pkparse.function
@@ -17,9 +17,8 @@
     int res;
     char *pwd = password;
 
-    MD_PSA_INIT();
-
     mbedtls_pk_init(&ctx);
+    MD_PSA_INIT();
 
     if (strcmp(pwd, "NULL") == 0) {
         pwd = NULL;
@@ -50,9 +49,8 @@
     mbedtls_pk_context ctx;
     int res;
 
-    MD_PSA_INIT();
-
     mbedtls_pk_init(&ctx);
+    MD_PSA_INIT();
 
     res = mbedtls_pk_parse_public_keyfile(&ctx, key_file);
 
@@ -78,6 +76,7 @@
     int res;
 
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     res = mbedtls_pk_parse_public_keyfile(&ctx, key_file);
 
@@ -92,6 +91,7 @@
 
 exit:
     mbedtls_pk_free(&ctx);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -101,8 +101,8 @@
     mbedtls_pk_context ctx;
     int res;
 
-    USE_PSA_INIT();
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     res = mbedtls_pk_parse_keyfile(&ctx, key_file, password,
                                    mbedtls_test_rnd_std_rand, NULL);
@@ -128,11 +128,13 @@
     mbedtls_pk_context pk;
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_parse_key(&pk, buf->x, buf->len, NULL, 0,
                                      mbedtls_test_rnd_std_rand, NULL) == result);
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_pkwrite.function b/tests/suites/test_suite_pkwrite.function
index c0c5ad0..804e9a7 100644
--- a/tests/suites/test_suite_pkwrite.function
+++ b/tests/suites/test_suite_pkwrite.function
@@ -36,6 +36,9 @@
     size_t buf_len, check_buf_len;
     int ret;
 
+    mbedtls_pk_init(&key);
+    USE_PSA_INIT();
+
     /* Note: if mbedtls_pk_load_file() successfully reads the file, then
        it also allocates check_buf, which should be freed on exit */
     TEST_EQUAL(mbedtls_pk_load_file(key_file, &check_buf, &check_buf_len), 0);
@@ -56,7 +59,6 @@
 
     ASSERT_ALLOC(buf, check_buf_len);
 
-    mbedtls_pk_init(&key);
     if (is_public_key) {
         TEST_EQUAL(mbedtls_pk_parse_public_keyfile(&key, key_file), 0);
         if (is_der) {
@@ -98,6 +100,7 @@
     mbedtls_free(buf);
     mbedtls_free(check_buf);
     mbedtls_pk_free(&key);
+    USE_PSA_DONE();
 }
 /* END_HEADER */
 
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 6bda6ca..6f9e544 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -12,6 +12,8 @@
 #include <constant_time_internal.h>
 #include <test/constant_flow.h>
 
+#define SSL_MESSAGE_QUEUE_INIT      { NULL, 0, 0, 0 }
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -27,6 +29,7 @@
     unsigned char input[MSGLEN];
     unsigned char output[MSGLEN];
 
+    USE_PSA_INIT();
     memset(input, 0, sizeof(input));
 
     /* Make sure calling put and get on NULL buffer results in error. */
@@ -79,8 +82,8 @@
 
 
 exit:
-
     mbedtls_test_ssl_buffer_free(&buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -115,6 +118,7 @@
     size_t i, j, written, read;
 
     mbedtls_test_ssl_buffer_init(&buf);
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_test_ssl_buffer_setup(&buf, size) == 0);
 
     /* Check the sanity of input parameters and initialise local variables. That
@@ -189,10 +193,10 @@
     }
 
 exit:
-
     mbedtls_free(input);
     mbedtls_free(output);
     mbedtls_test_ssl_buffer_free(&buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -210,6 +214,7 @@
     mbedtls_test_mock_socket socket;
 
     mbedtls_test_mock_socket_init(&socket);
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_test_mock_tcp_send_b(&socket, message, MSGLEN) < 0);
     mbedtls_test_mock_socket_close(&socket);
     mbedtls_test_mock_socket_init(&socket);
@@ -224,8 +229,8 @@
     mbedtls_test_mock_socket_close(&socket);
 
 exit:
-
     mbedtls_test_mock_socket_close(&socket);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -259,6 +264,7 @@
 
     mbedtls_test_mock_socket_init(&client);
     mbedtls_test_mock_socket_init(&server);
+    USE_PSA_INIT();
 
     /* Fill up the buffer with structured data so that unwanted changes
      * can be detected */
@@ -317,9 +323,9 @@
     TEST_ASSERT(memcmp(message, received, MSGLEN) == 0);
 
 exit:
-
     mbedtls_test_mock_socket_close(&client);
     mbedtls_test_mock_socket_close(&server);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -357,6 +363,7 @@
 
     mbedtls_test_mock_socket_init(&client);
     mbedtls_test_mock_socket_init(&server);
+    USE_PSA_INIT();
 
     /* Fill up the buffers with structured data so that unwanted changes
      * can be detected */
@@ -445,17 +452,18 @@
     }
 
 exit:
-
     mbedtls_test_mock_socket_close(&client);
     mbedtls_test_mock_socket_close(&server);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void ssl_message_queue_sanity()
 {
-    mbedtls_test_ssl_message_queue queue;
+    mbedtls_test_ssl_message_queue queue = SSL_MESSAGE_QUEUE_INIT;
 
+    USE_PSA_INIT();
     /* Trying to push/pull to an empty queue */
     TEST_ASSERT(mbedtls_test_ssl_message_queue_push_info(NULL, 1)
                 == MBEDTLS_TEST_ERROR_ARG_NULL);
@@ -468,14 +476,16 @@
 
 exit:
     mbedtls_test_ssl_message_queue_free(&queue);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void ssl_message_queue_basic()
 {
-    mbedtls_test_ssl_message_queue queue;
+    mbedtls_test_ssl_message_queue queue = SSL_MESSAGE_QUEUE_INIT;
 
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_test_ssl_message_queue_setup(&queue, 3) == 0);
 
     /* Sanity test - 3 pushes and 3 pops with sufficient space */
@@ -495,14 +505,16 @@
 
 exit:
     mbedtls_test_ssl_message_queue_free(&queue);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void ssl_message_queue_overflow_underflow()
 {
-    mbedtls_test_ssl_message_queue queue;
+    mbedtls_test_ssl_message_queue queue = SSL_MESSAGE_QUEUE_INIT;
 
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_test_ssl_message_queue_setup(&queue, 3) == 0);
 
     /* 4 pushes (last one with an error), 4 pops (last one with an error) */
@@ -521,14 +533,16 @@
 
 exit:
     mbedtls_test_ssl_message_queue_free(&queue);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void ssl_message_queue_interleaved()
 {
-    mbedtls_test_ssl_message_queue queue;
+    mbedtls_test_ssl_message_queue queue = SSL_MESSAGE_QUEUE_INIT;
 
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_test_ssl_message_queue_setup(&queue, 3) == 0);
 
     /* Interleaved test - [2 pushes, 1 pop] twice, and then two pops
@@ -555,16 +569,18 @@
 
 exit:
     mbedtls_test_ssl_message_queue_free(&queue);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void ssl_message_queue_insufficient_buffer()
 {
-    mbedtls_test_ssl_message_queue queue;
+    mbedtls_test_ssl_message_queue queue = SSL_MESSAGE_QUEUE_INIT;
     size_t message_len = 10;
     size_t buffer_len = 5;
 
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_test_ssl_message_queue_setup(&queue, 1) == 0);
 
     /* Popping without a sufficient buffer */
@@ -574,6 +590,7 @@
                 == (int) buffer_len);
 exit:
     mbedtls_test_ssl_message_queue_free(&queue);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -588,6 +605,7 @@
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
 
+    USE_PSA_INIT();
     /* Send with a NULL context */
     TEST_ASSERT(mbedtls_test_mock_tcp_send_msg(NULL, message, MSGLEN)
                 == MBEDTLS_TEST_ERROR_CONTEXT_ERROR);
@@ -626,6 +644,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -638,8 +657,10 @@
     unsigned i;
     mbedtls_test_ssl_message_queue server_queue, client_queue;
     mbedtls_test_message_socket_context server_context, client_context;
+
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_message_socket_setup(&server_queue,
                                                   &client_queue, 1,
@@ -685,6 +706,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -697,8 +719,10 @@
     unsigned i;
     mbedtls_test_ssl_message_queue server_queue, client_queue;
     mbedtls_test_message_socket_context server_context, client_context;
+
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_message_socket_setup(&server_queue,
                                                   &client_queue, 2,
@@ -749,6 +773,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -761,8 +786,10 @@
     unsigned i;
     mbedtls_test_ssl_message_queue server_queue, client_queue;
     mbedtls_test_message_socket_context server_context, client_context;
+
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_message_socket_setup(&server_queue,
                                                   &client_queue, 2,
@@ -801,6 +828,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -813,8 +841,10 @@
     unsigned i;
     mbedtls_test_ssl_message_queue server_queue, client_queue;
     mbedtls_test_message_socket_context server_context, client_context;
+
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_message_socket_setup(&server_queue,
                                                   &client_queue, 2,
@@ -865,6 +895,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -877,8 +908,10 @@
     unsigned i;
     mbedtls_test_ssl_message_queue server_queue, client_queue;
     mbedtls_test_message_socket_context server_context, client_context;
+
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_message_socket_setup(&server_queue,
                                                   &client_queue, 1,
@@ -923,6 +956,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -935,8 +969,10 @@
     unsigned i;
     mbedtls_test_ssl_message_queue server_queue, client_queue;
     mbedtls_test_message_socket_context server_context, client_context;
+
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_message_socket_setup(&server_queue,
                                                   &client_queue, 3,
@@ -983,6 +1019,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -995,8 +1032,10 @@
     unsigned i;
     mbedtls_test_ssl_message_queue server_queue, client_queue;
     mbedtls_test_message_socket_context server_context, client_context;
+
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_message_socket_setup(&server_queue,
                                                   &client_queue, 3,
@@ -1070,6 +1109,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1080,10 +1120,9 @@
     mbedtls_ssl_context ssl;
     mbedtls_ssl_config conf;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_config_defaults(&conf,
                                             MBEDTLS_SSL_IS_CLIENT,
@@ -1112,12 +1151,16 @@
 void ssl_set_hostname_twice(char *hostname0, char *hostname1)
 {
     mbedtls_ssl_context ssl;
+
     mbedtls_ssl_init(&ssl);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_set_hostname(&ssl, hostname0) == 0);
     TEST_ASSERT(mbedtls_ssl_set_hostname(&ssl, hostname1) == 0);
 
+exit:
     mbedtls_ssl_free(&ssl);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1141,11 +1184,11 @@
     size_t const buflen = 512;
     mbedtls_record rec, rec_backup;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    MD_OR_USE_PSA_INIT();
+
     ret = mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
                                             etm, tag_mode, ver,
                                             (size_t) cid0_len,
@@ -1295,11 +1338,11 @@
     int seen_success; /* Indicates if in the current mode we've
                        * already seen a successful test. */
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    MD_OR_USE_PSA_INIT();
+
     ret = mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
                                             etm, tag_mode, ver,
                                             (size_t) cid0_len,
@@ -1456,11 +1499,10 @@
     int ret;
     const unsigned char pad_max_len = 255; /* Per the standard */
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    MD_OR_USE_PSA_INIT();
 
     /* Set up transforms with dummy keys */
     ret = mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
@@ -1685,6 +1727,7 @@
     ASSERT_COMPARE(dst, (size_t) desired_length,
                    expected->x, (size_t) expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1738,6 +1781,7 @@
                    expected_server_write_iv->x,
                    (size_t) desired_iv_len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1782,6 +1826,7 @@
     ASSERT_COMPARE(dst, desired_length,
                    expected->x, desired_length);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1815,6 +1860,7 @@
     ASSERT_COMPARE(secrets.early_exporter_master_secret, hash_len,
                    exporter_expected->x, exporter_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1848,6 +1894,7 @@
     ASSERT_COMPARE(secrets.server_handshake_traffic_secret, hash_len,
                    server_expected->x, server_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1885,6 +1932,7 @@
     ASSERT_COMPARE(secrets.exporter_master_secret, hash_len,
                    exporter_expected->x, exporter_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1914,6 +1962,7 @@
     ASSERT_COMPARE(secrets.resumption_master_secret, hash_len,
                    resumption_expected->x, resumption_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1947,6 +1996,7 @@
     ASSERT_COMPARE(binder, hash_len,
                    binder_expected->x, binder_expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -1971,8 +2021,6 @@
     size_t buf_len;
     int other_endpoint;
 
-    MD_OR_USE_PSA_INIT();
-
     TEST_ASSERT(endpoint == MBEDTLS_SSL_IS_CLIENT ||
                 endpoint == MBEDTLS_SSL_IS_SERVER);
 
@@ -2000,6 +2048,7 @@
 
     mbedtls_ssl_transform_init(&transform_recv);
     mbedtls_ssl_transform_init(&transform_send);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_tls13_populate_transform(
                     &transform_send, endpoint,
@@ -2045,6 +2094,7 @@
     ASSERT_COMPARE(rec.buf + rec.data_offset, rec.data_len,
                    plaintext->x, plaintext->len);
 
+exit:
     mbedtls_free(buf);
     mbedtls_ssl_transform_free(&transform_send);
     mbedtls_ssl_transform_free(&transform_recv);
@@ -2071,6 +2121,7 @@
     ASSERT_COMPARE(secret_new, (size_t) expected->len,
                    expected->x, (size_t) expected->len);
 
+exit:
     PSA_DONE();
 }
 /* END_CASE */
@@ -2114,9 +2165,9 @@
     /*
      * Test that a save-load pair is the identity
      */
-
     mbedtls_ssl_session_init(&original);
     mbedtls_ssl_session_init(&restored);
+    USE_PSA_INIT();
 
     /* Prepare a dummy session to work on */
     ((void) endpoint_type);
@@ -2248,6 +2299,7 @@
     mbedtls_ssl_session_free(&original);
     mbedtls_ssl_session_free(&restored);
     mbedtls_free(buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2262,8 +2314,8 @@
     /*
      * Test that a load-save pair is the identity
      */
-
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Prepare a dummy session to work on */
     ((void) endpoint_type);
@@ -2310,6 +2362,7 @@
     mbedtls_ssl_session_free(&session);
     mbedtls_free(buf1);
     mbedtls_free(buf2);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2324,8 +2377,8 @@
     /*
      * Test that session_save() fails cleanly on small buffers
      */
-
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Prepare dummy session and get serialized size */
     ((void) endpoint_type);
@@ -2357,6 +2410,7 @@
 exit:
     mbedtls_ssl_session_free(&session);
     mbedtls_free(buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2371,8 +2425,8 @@
     /*
      * Test that session_load() fails cleanly on small buffers
      */
-
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Prepare serialized session data */
     ((void) endpoint_type);
@@ -2410,6 +2464,7 @@
     mbedtls_ssl_session_free(&session);
     mbedtls_free(good_buf);
     mbedtls_free(bad_buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2432,6 +2487,7 @@
                                       corrupt_config == 1 };
 
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
     ((void) endpoint_type);
     ((void) tls_version);
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
@@ -2484,7 +2540,8 @@
             *byte ^= corrupted_bit;
         }
     }
-
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2645,13 +2702,10 @@
 /* BEGIN_CASE depends_on:MBEDTLS_SSL_HANDSHAKE_WITH_CERT_ENABLED:MBEDTLS_PKCS1_V15:MBEDTLS_SSL_PROTO_TLS1_2:MBEDTLS_MD_CAN_SHA256 */
 void handshake_cipher(char *cipher, int pk_alg, int dtls)
 {
-    MD_OR_USE_PSA_INIT();
-
     test_handshake_psk_cipher(cipher, pk_alg, NULL, dtls);
 
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
-    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2721,15 +2775,10 @@
                   int expected_cli_fragments,
                   int expected_srv_fragments)
 {
-    MD_OR_USE_PSA_INIT();
-
     test_app_data(mfl, cli_msg_len, srv_msg_len, expected_cli_fragments,
                   expected_srv_fragments, 0);
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
-
-exit:
-    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2738,15 +2787,10 @@
                    int expected_cli_fragments,
                    int expected_srv_fragments)
 {
-    MD_OR_USE_PSA_INIT();
-
     test_app_data(mfl, cli_msg_len, srv_msg_len, expected_cli_fragments,
                   expected_srv_fragments, 1);
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
-
-exit:
-    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2859,7 +2903,6 @@
 {
     test_resize_buffers(mfl, 0, MBEDTLS_SSL_LEGACY_NO_RENEGOTIATION, 1, 1,
                         (char *) "");
-
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
 }
@@ -2870,7 +2913,6 @@
                                     char *cipher)
 {
     test_resize_buffers(mfl, 1, legacy_renegotiation, 0, 1, cipher);
-
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
 }
@@ -2887,8 +2929,8 @@
 
     mbedtls_ssl_config conf;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_conf_psk(&conf,
                                      psk0, sizeof(psk0),
@@ -2899,9 +2941,7 @@
                 MBEDTLS_ERR_SSL_FEATURE_UNAVAILABLE);
 
 exit:
-
     mbedtls_ssl_config_free(&conf);
-
     MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
@@ -2931,8 +2971,8 @@
 
     mbedtls_ssl_config conf;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     switch (mode) {
         case 0:
@@ -2983,7 +3023,6 @@
     }
 
 exit:
-
     mbedtls_ssl_config_free(&conf);
     MD_OR_USE_PSA_DONE();
 
@@ -2998,10 +3037,9 @@
     mbedtls_ssl_config conf;
     mbedtls_ssl_context ssl;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_config_init(&conf);
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
 
     mbedtls_ssl_conf_endpoint(&conf, endpoint);
     mbedtls_ssl_conf_transport(&conf, transport);
@@ -3042,10 +3080,10 @@
 #endif
     mbedtls_ssl_conf_curves(&conf, curve_list);
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_context ssl;
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_ssl_setup(&ssl, &conf) == 0);
 
     TEST_ASSERT(ssl.handshake != NULL && ssl.handshake->group_list != NULL);
@@ -3059,7 +3097,6 @@
         TEST_EQUAL(iana_tls_group_list[i], ssl.handshake->group_list[i]);
     }
 
-
 exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
@@ -3083,10 +3120,10 @@
 
     mbedtls_ssl_conf_groups(&conf, iana_tls_group_list);
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_context ssl;
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_ssl_setup(&ssl, &conf) == 0);
 
     TEST_ASSERT(ssl.conf != NULL && ssl.conf->group_list != NULL);
@@ -3122,12 +3159,12 @@
     options.srv_log_obj = &srv_pattern;
     options.srv_log_fun = mbedtls_test_ssl_log_analyzer;
 
-    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&client, sizeof(client));
     mbedtls_platform_zeroize(&server, sizeof(server));
 
     mbedtls_test_message_socket_init(&server_context);
     mbedtls_test_message_socket_init(&client_context);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_test_ssl_endpoint_init(&client, MBEDTLS_SSL_IS_CLIENT,
                                                &options, NULL, NULL,
@@ -3181,6 +3218,8 @@
 
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
+    USE_PSA_INIT();
+
     TEST_EQUAL(mbedtls_ssl_config_defaults(&conf, MBEDTLS_SSL_IS_SERVER,
                                            MBEDTLS_SSL_TRANSPORT_DATAGRAM,
                                            MBEDTLS_SSL_PRESET_DEFAULT),
@@ -3195,8 +3234,10 @@
                                                     &len),
                exp_ret);
 
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3205,9 +3246,13 @@
 {
     mbedtls_timing_delay_context    delay_context;
 
+    USE_PSA_INIT();
     mbedtls_timing_set_delay(&delay_context, 50, 100);
 
     TEST_ASSERT(mbedtls_timing_get_final_delay(&delay_context) == 100);
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3224,10 +3269,9 @@
 
     mbedtls_test_rnd_std_rand(NULL, own_cid, sizeof(own_cid));
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_config_defaults(&conf,
                                             MBEDTLS_SSL_IS_CLIENT,
@@ -3383,11 +3427,12 @@
     /*
      * Test set-up
      */
-    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&client_ep, sizeof(client_ep));
     mbedtls_platform_zeroize(&server_ep, sizeof(server_ep));
 
     mbedtls_test_init_handshake_options(&client_options);
+    MD_OR_USE_PSA_INIT();
+
     client_options.pk_alg = MBEDTLS_PK_ECDSA;
     ret = mbedtls_test_ssl_endpoint_init(&client_ep, MBEDTLS_SSL_IS_CLIENT,
                                          &client_options, NULL, NULL, NULL,
@@ -3478,9 +3523,8 @@
     size_t pwd_len = 0;
     int ret;
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_ssl_init(&ssl);
+    MD_OR_USE_PSA_INIT();
 
     /* test with uninitalized SSL context */
     ECJPAKE_TEST_SET_PASSWORD(MBEDTLS_ERR_SSL_BAD_INPUT_DATA);
@@ -3617,7 +3661,8 @@
 #else
     TEST_UNAVAILABLE_ECC(30, MBEDTLS_ECP_DP_CURVE448, PSA_ECC_FAMILY_MONTGOMERY, 448);
 #endif
-
+    goto exit;
+exit:
     MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_x509parse.function b/tests/suites/test_suite_x509parse.function
index 905d62f..73e6803 100644
--- a/tests/suites/test_suite_x509parse.function
+++ b/tests/suites/test_suite_x509parse.function
@@ -429,12 +429,15 @@
     int expected_result = ext_type & has_ext_type;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     crt.ext_types = ext_type;
 
     TEST_ASSERT(mbedtls_x509_crt_has_ext_type(&crt, has_ext_type) == expected_result);
 
+exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -463,6 +466,7 @@
     size_t n = sizeof(buf);
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_EQUAL(mbedtls_x509_crt_parse_file(&crt, crt_file), parse_result);
@@ -490,8 +494,8 @@
     TEST_ASSERT(strcmp(buf, result_str) == 0);
 
 exit:
-
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -503,6 +507,7 @@
     int res;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -515,6 +520,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -526,6 +532,7 @@
     int res;
 
     mbedtls_x509_crl_init(&crl);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crl_parse_file(&crl, crl_file) == 0);
@@ -538,6 +545,7 @@
 
 exit:
     mbedtls_x509_crl_free(&crl);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -548,12 +556,14 @@
     char buf[2000];
 
     mbedtls_x509_crl_init(&crl);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crl_parse_file(&crl, crl_file) == result);
 
 exit:
     mbedtls_x509_crl_free(&crl);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -565,6 +575,7 @@
     int res;
 
     mbedtls_x509_csr_init(&csr);
+    USE_PSA_INIT();
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_csr_parse_file(&csr, csr_file) == 0);
@@ -577,6 +588,7 @@
 
 exit:
     mbedtls_x509_csr_free(&csr);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -586,6 +598,7 @@
     char buf[2000];
     int res;
 
+    USE_PSA_INIT();
     memset(buf, 0, sizeof(buf));
 
     res = mbedtls_x509_crt_verify_info(buf, sizeof(buf), prefix, flags);
@@ -593,6 +606,9 @@
     TEST_ASSERT(res >= 0);
 
     TEST_ASSERT(strcmp(buf, result_str) == 0);
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -616,11 +632,9 @@
      * - x509_verify() for server5 -> test-ca2:             ~ 18800
      * - x509_verify() for server10 -> int-ca3 -> int-ca2:  ~ 25500
      */
-
     mbedtls_x509_crt_restart_init(&rs_ctx);
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
-
     MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -673,7 +687,6 @@
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
     mbedtls_x509_crl_init(&crl);
-
     MD_OR_USE_PSA_INIT();
 
     if (strcmp(cn_name_str, "NULL") != 0) {
@@ -758,6 +771,7 @@
 
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&ca, ca_file) == 0);
@@ -775,6 +789,7 @@
 exit:
     mbedtls_x509_crt_free(&crt);
     mbedtls_x509_crt_free(&ca);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -790,10 +805,10 @@
 
     mbedtls_x509_crt_init(&crt);
     mbedtls_x509_crt_init(&ca);
-    verify_print_init(&vrfy_ctx);
-
     MD_OR_USE_PSA_INIT();
 
+    verify_print_init(&vrfy_ctx);
+
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&ca, ca_file) == 0);
 
@@ -827,6 +842,8 @@
     int res = 0;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
+
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -844,6 +861,7 @@
     }
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -855,6 +873,8 @@
     int res = 0;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
+
     memset(buf, 0, 2000);
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
@@ -873,18 +893,20 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_X509_CRT_PARSE_C */
 void mbedtls_x509_get_name(char *rdn_sequence, int exp_ret)
 {
-    unsigned char *name;
+    unsigned char *name = NULL;
     unsigned char *p;
     size_t name_len;
     mbedtls_x509_name head;
     int ret;
 
+    USE_PSA_INIT();
     memset(&head, 0, sizeof(head));
 
     name = mbedtls_test_unhexify_alloc(rdn_sequence, &name_len);
@@ -897,7 +919,9 @@
 
     TEST_EQUAL(ret, exp_ret);
 
+exit:
     mbedtls_free(name);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -916,6 +940,7 @@
     unsigned char buf[80], *out = NULL, *c;
     const char *short_name;
 
+    USE_PSA_INIT();
     memset(&parsed, 0, sizeof(parsed));
     memset(buf, 0, sizeof(buf));
     c = buf + sizeof(buf);
@@ -964,6 +989,7 @@
     mbedtls_free(out);
     mbedtls_asn1_free_named_data_list(&names);
     mbedtls_asn1_free_named_data_list_shallow(parsed.next);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -973,6 +999,7 @@
     mbedtls_x509_crt   crt;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -986,6 +1013,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -995,6 +1023,7 @@
     mbedtls_x509_crt   crt;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1008,6 +1037,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1017,11 +1047,13 @@
     mbedtls_x509_crt crt;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == result);
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1037,6 +1069,7 @@
 #endif
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_der(&crt, buf->x, buf->len) == (result));
 #if !defined(MBEDTLS_X509_REMOVE_INFO)
@@ -1103,6 +1136,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1124,6 +1158,7 @@
     oid.p = (unsigned char *) MBEDTLS_OID_PKIX "\x01\x1F";
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_der_with_ext_cb(&crt, buf->x, buf->len, 0, parse_crt_ext_cb,
                                                        &oid) == (result));
@@ -1157,6 +1192,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1168,6 +1204,8 @@
     int res;
 
     mbedtls_x509_crl_init(&crl);
+    USE_PSA_INIT();
+
     memset(output, 0, 2000);
 
 
@@ -1183,6 +1221,7 @@
 
 exit:
     mbedtls_x509_crl_free(&crl);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1194,6 +1233,8 @@
     int my_ret;
 
     mbedtls_x509_csr_init(&csr);
+    USE_PSA_INIT();
+
     memset(my_out, 0, sizeof(my_out));
 
     my_ret = mbedtls_x509_csr_parse_der(&csr, csr_der->x, csr_der->len);
@@ -1207,6 +1248,7 @@
 
 exit:
     mbedtls_x509_csr_free(&csr);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1218,6 +1260,8 @@
     int my_ret;
 
     mbedtls_x509_csr_init(&csr);
+    USE_PSA_INIT();
+
     memset(my_out, 0, sizeof(my_out));
 
     my_ret = mbedtls_x509_csr_parse_file(&csr, csr_file);
@@ -1231,6 +1275,7 @@
 
 exit:
     mbedtls_x509_csr_free(&csr);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1241,6 +1286,7 @@
     int i;
 
     mbedtls_x509_crt_init(&chain);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_path(&chain, crt_path) == ret);
 
@@ -1255,6 +1301,7 @@
 
 exit:
     mbedtls_x509_crt_free(&chain);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1271,10 +1318,8 @@
      * We expect chain_dir to contain certificates 00.crt, 01.crt, etc.
      * with NN.crt signed by NN-1.crt
      */
-
     mbedtls_x509_crt_init(&trusted);
     mbedtls_x509_crt_init(&chain);
-
     MD_OR_USE_PSA_INIT();
 
     /* Load trusted root */
@@ -1313,7 +1358,6 @@
 
     mbedtls_x509_crt_init(&chain);
     mbedtls_x509_crt_init(&trusted);
-
     MD_OR_USE_PSA_INIT();
 
     while ((act = mystrsep(&chain_paths, " ")) != NULL) {
@@ -1353,6 +1397,7 @@
     const char *desc = NULL;
     int ret;
 
+    USE_PSA_INIT();
 
     oid.tag = MBEDTLS_ASN1_OID;
     oid.p   = buf->x;
@@ -1368,6 +1413,9 @@
         TEST_ASSERT(desc != NULL);
         TEST_ASSERT(strcmp(desc, ref_desc) == 0);
     }
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1377,6 +1425,8 @@
     mbedtls_x509_buf oid;
     char num_buf[100];
 
+    USE_PSA_INIT();
+
     memset(num_buf, 0x2a, sizeof(num_buf));
 
     oid.tag = MBEDTLS_ASN1_OID;
@@ -1391,6 +1441,9 @@
         TEST_ASSERT(num_buf[ret] == 0);
         TEST_ASSERT(strcmp(num_buf, numstr) == 0);
     }
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1400,6 +1453,7 @@
     mbedtls_x509_crt crt;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1407,6 +1461,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1417,7 +1472,7 @@
     mbedtls_x509_crt crt;
 
     mbedtls_x509_crt_init(&crt);
-
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1426,6 +1481,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1438,6 +1494,7 @@
     unsigned char *start = buf;
     unsigned char *end = buf;
 
+    USE_PSA_INIT();
     memset(&time, 0x00, sizeof(time));
     *end = (unsigned char) tag; end++;
     *end = strlen(time_str);
@@ -1455,6 +1512,8 @@
         TEST_ASSERT(min  == time.min);
         TEST_ASSERT(sec  == time.sec);
     }
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1468,6 +1527,8 @@
     mbedtls_md_type_t my_msg_md, my_mgf_md;
     int my_salt_len;
 
+    USE_PSA_INIT();
+
     buf.p = params->x;
     buf.len = params->len;
     buf.tag = params_tag;
@@ -1484,6 +1545,6 @@
     }
 
 exit:
-    ;;
+    USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_x509write.function b/tests/suites/test_suite_x509write.function
index 0e4062e..b08555c 100644
--- a/tests/suites/test_suite_x509write.function
+++ b/tests/suites/test_suite_x509write.function
@@ -176,10 +176,9 @@
     memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
 
     mbedtls_x509write_csr_init(&req);
-
+    mbedtls_pk_init(&key);
     MD_OR_USE_PSA_INIT();
 
-    mbedtls_pk_init(&key);
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&key, key_file, NULL,
                                          mbedtls_test_rnd_std_rand, NULL) == 0);
 
@@ -266,12 +265,11 @@
     const char *subject_name = "C=NL,O=PolarSSL,CN=PolarSSL Server 1";
     mbedtls_test_rnd_pseudo_info rnd_info;
 
-    memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
-
     mbedtls_x509write_csr_init(&req);
-
     MD_OR_USE_PSA_INIT();
 
+    memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
+
     md_alg_psa = mbedtls_hash_info_psa_from_md((mbedtls_md_type_t) md_type);
     TEST_ASSERT(md_alg_psa != MBEDTLS_MD_NONE);
 
@@ -315,7 +313,7 @@
     mbedtls_x509write_csr_free(&req);
     mbedtls_pk_free(&key);
     psa_destroy_key(key_id);
-    PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -356,13 +354,11 @@
     mbedtls_mpi_init(&serial_mpi);
 #endif
 
-    MD_OR_USE_PSA_INIT();
-
     mbedtls_pk_init(&subject_key);
     mbedtls_pk_init(&issuer_key);
     mbedtls_pk_init(&issuer_key_alt);
-
     mbedtls_x509write_crt_init(&crt);
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&subject_key, subject_key_file,
                                          subject_pwd, mbedtls_test_rnd_std_rand, NULL) == 0);
@@ -597,6 +593,7 @@
     mbedtls_x509write_cert ctx;
     uint8_t invalid_serial[MBEDTLS_X509_RFC5280_MAX_SERIAL_LEN + 1];
 
+    USE_PSA_INIT();
     memset(invalid_serial, 0x01, sizeof(invalid_serial));
 
 #if defined(MBEDTLS_TEST_DEPRECATED) && defined(MBEDTLS_BIGNUM_C)
@@ -619,6 +616,7 @@
 #else
     ;
 #endif
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -632,6 +630,8 @@
     mbedtls_x509_name parsed, *parsed_cur, *parsed_prv;
     unsigned char buf[1024], out[1024], *c;
 
+    USE_PSA_INIT();
+
     memset(&parsed, 0, sizeof(parsed));
     memset(out, 0, sizeof(out));
     memset(buf, 0, sizeof(buf));
@@ -665,5 +665,6 @@
         parsed_cur = parsed_cur->next;
         mbedtls_free(parsed_prv);
     }
+    USE_PSA_DONE();
 }
 /* END_CASE */