Merge pull request #7503 from gilles-peskine-arm/test-argument-types-union-2.28

Backport 2.28: Support larger integer test arguments
diff --git a/ChangeLog.d/programs_psa_fix.txt b/ChangeLog.d/programs_psa_fix.txt
new file mode 100644
index 0000000..fe2099e
--- /dev/null
+++ b/ChangeLog.d/programs_psa_fix.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Fix missing PSA initialization in sample programs when
+     MBEDTLS_USE_PSA_CRYPTO is enabled.
diff --git a/library/pk.c b/library/pk.c
index d46a934..12f4120 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -646,6 +646,7 @@
     psa_key_type_t key_type;
     size_t bits;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    psa_status_t status;
 
     /* export the private key material in the format PSA wants */
     if (mbedtls_pk_get_type(pk) != MBEDTLS_PK_ECKEY) {
@@ -668,7 +669,9 @@
     psa_set_key_algorithm(&attributes, PSA_ALG_ECDSA(hash_alg));
 
     /* import private key into PSA */
-    if (PSA_SUCCESS != psa_import_key(&attributes, d, d_len, key)) {
+    status = psa_import_key(&attributes, d, d_len, key);
+    mbedtls_platform_zeroize(d, sizeof(d));
+    if (status != PSA_SUCCESS) {
         return MBEDTLS_ERR_PK_HW_ACCEL_FAILED;
     }
 
diff --git a/programs/fuzz/fuzz_client.c b/programs/fuzz/fuzz_client.c
index a415874..2de51a6 100644
--- a/programs/fuzz/fuzz_client.c
+++ b/programs/fuzz/fuzz_client.c
@@ -77,6 +77,13 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_entropy_init(&entropy);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
         goto exit;
@@ -184,6 +191,9 @@
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ssl_free(&ssl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_dtlsclient.c b/programs/fuzz/fuzz_dtlsclient.c
index 1fcbc92..d414bb3 100644
--- a/programs/fuzz/fuzz_dtlsclient.c
+++ b/programs/fuzz/fuzz_dtlsclient.c
@@ -61,6 +61,13 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_entropy_init(&entropy);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     srand(1);
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
@@ -119,6 +126,9 @@
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ssl_free(&ssl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_dtlsserver.c b/programs/fuzz/fuzz_dtlsserver.c
index 529fbbf..df4087a 100644
--- a/programs/fuzz/fuzz_dtlsserver.c
+++ b/programs/fuzz/fuzz_dtlsserver.c
@@ -74,6 +74,13 @@
     mbedtls_entropy_init(&entropy);
     mbedtls_ssl_cookie_init(&cookie_ctx);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
         goto exit;
@@ -152,9 +159,16 @@
 exit:
     mbedtls_ssl_cookie_free(&cookie_ctx);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
+    mbedtls_pk_free(&pkey);
+    mbedtls_x509_crt_free(&srvcert);
+#endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ssl_free(&ssl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_privkey.c b/programs/fuzz/fuzz_privkey.c
index c24f275..d1da589 100644
--- a/programs/fuzz/fuzz_privkey.c
+++ b/programs/fuzz/fuzz_privkey.c
@@ -18,6 +18,14 @@
     }
 
     mbedtls_pk_init(&pk);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     ret = mbedtls_pk_parse_key(&pk, Data, Size, NULL, 0);
     if (ret == 0) {
 #if defined(MBEDTLS_RSA_C)
@@ -63,6 +71,10 @@
             abort();
         }
     }
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_pk_free(&pk);
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_pubkey.c b/programs/fuzz/fuzz_pubkey.c
index 388b4c5..daca2b3 100644
--- a/programs/fuzz/fuzz_pubkey.c
+++ b/programs/fuzz/fuzz_pubkey.c
@@ -9,6 +9,12 @@
     mbedtls_pk_context pk;
 
     mbedtls_pk_init(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_pk_parse_public_key(&pk, Data, Size);
     if (ret == 0) {
 #if defined(MBEDTLS_RSA_C)
@@ -64,6 +70,10 @@
             abort();
         }
     }
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_pk_free(&pk);
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_server.c b/programs/fuzz/fuzz_server.c
index e161d7e..06aeb5e 100644
--- a/programs/fuzz/fuzz_server.c
+++ b/programs/fuzz/fuzz_server.c
@@ -89,6 +89,13 @@
     mbedtls_ssl_ticket_init(&ticket_ctx);
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
         goto exit;
@@ -195,8 +202,14 @@
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
+#if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
+    mbedtls_x509_crt_free(&srvcert);
+    mbedtls_pk_free(&pkey);
+#endif
     mbedtls_ssl_free(&ssl);
-
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif
 #else
     (void) Data;
     (void) Size;
diff --git a/programs/fuzz/fuzz_x509crl.c b/programs/fuzz/fuzz_x509crl.c
index 3aaa8e5..1140c3d 100644
--- a/programs/fuzz/fuzz_x509crl.c
+++ b/programs/fuzz/fuzz_x509crl.c
@@ -9,10 +9,20 @@
     unsigned char buf[4096];
 
     mbedtls_x509_crl_init(&crl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_x509_crl_parse(&crl, Data, Size);
     if (ret == 0) {
         ret = mbedtls_x509_crl_info((char *) buf, sizeof(buf) - 1, " ", &crl);
     }
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_x509_crl_free(&crl);
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_x509crt.c b/programs/fuzz/fuzz_x509crt.c
index a5cb7ec..3593236 100644
--- a/programs/fuzz/fuzz_x509crt.c
+++ b/programs/fuzz/fuzz_x509crt.c
@@ -9,10 +9,20 @@
     unsigned char buf[4096];
 
     mbedtls_x509_crt_init(&crt);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_x509_crt_parse(&crt, Data, Size);
     if (ret == 0) {
         ret = mbedtls_x509_crt_info((char *) buf, sizeof(buf) - 1, " ", &crt);
     }
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_x509_crt_free(&crt);
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_x509csr.c b/programs/fuzz/fuzz_x509csr.c
index afd2031..0ca9b87 100644
--- a/programs/fuzz/fuzz_x509csr.c
+++ b/programs/fuzz/fuzz_x509csr.c
@@ -9,10 +9,20 @@
     unsigned char buf[4096];
 
     mbedtls_x509_csr_init(&csr);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_x509_csr_parse(&csr, Data, Size);
     if (ret == 0) {
         ret = mbedtls_x509_csr_info((char *) buf, sizeof(buf) - 1, " ", &csr);
     }
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_x509_csr_free(&csr);
 #else
     (void) Data;
diff --git a/programs/pkey/gen_key.c b/programs/pkey/gen_key.c
index 1a6463d..cd21743 100644
--- a/programs/pkey/gen_key.c
+++ b/programs/pkey/gen_key.c
@@ -204,6 +204,15 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     memset(buf, 0, sizeof(buf));
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -411,6 +420,9 @@
     mbedtls_pk_free(&key);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/pkey/key_app.c b/programs/pkey/key_app.c
index a757cb3..2f30830 100644
--- a/programs/pkey/key_app.c
+++ b/programs/pkey/key_app.c
@@ -91,6 +91,15 @@
     mbedtls_pk_init(&pk);
     memset(buf, 0, sizeof(buf));
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto cleanup;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P); mbedtls_mpi_init(&Q);
     mbedtls_mpi_init(&D); mbedtls_mpi_init(&E); mbedtls_mpi_init(&DP);
     mbedtls_mpi_init(&DQ); mbedtls_mpi_init(&QP);
@@ -275,6 +284,9 @@
 #endif
 
     mbedtls_pk_free(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_mpi_free(&N); mbedtls_mpi_free(&P); mbedtls_mpi_free(&Q);
     mbedtls_mpi_free(&D); mbedtls_mpi_free(&E); mbedtls_mpi_free(&DP);
     mbedtls_mpi_free(&DQ); mbedtls_mpi_free(&QP);
diff --git a/programs/pkey/key_app_writer.c b/programs/pkey/key_app_writer.c
index 0009d91..e986ada 100644
--- a/programs/pkey/key_app_writer.c
+++ b/programs/pkey/key_app_writer.c
@@ -205,6 +205,15 @@
     memset(buf, 0, sizeof(buf));
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P); mbedtls_mpi_init(&Q);
     mbedtls_mpi_init(&D); mbedtls_mpi_init(&E); mbedtls_mpi_init(&DP);
     mbedtls_mpi_init(&DQ); mbedtls_mpi_init(&QP);
@@ -400,6 +409,9 @@
     mbedtls_mpi_free(&DQ); mbedtls_mpi_free(&QP);
 
     mbedtls_pk_free(&key);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/pkey/pk_decrypt.c b/programs/pkey/pk_decrypt.c
index 1dff75c..c3ff53d 100644
--- a/programs/pkey/pk_decrypt.c
+++ b/programs/pkey/pk_decrypt.c
@@ -71,6 +71,15 @@
 
     memset(result, 0, sizeof(result));
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 2) {
         mbedtls_printf("usage: mbedtls_pk_decrypt <key_file>\n");
 
@@ -142,6 +151,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/pk_encrypt.c b/programs/pkey/pk_encrypt.c
index 9a2549a..5f5a424 100644
--- a/programs/pkey/pk_encrypt.c
+++ b/programs/pkey/pk_encrypt.c
@@ -67,6 +67,15 @@
     mbedtls_entropy_init(&entropy);
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: mbedtls_pk_encrypt <key_file> <string of max 100 characters>\n");
 
@@ -144,6 +153,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/pk_sign.c b/programs/pkey/pk_sign.c
index 19a855b..2a8b7a4 100644
--- a/programs/pkey/pk_sign.c
+++ b/programs/pkey/pk_sign.c
@@ -66,6 +66,15 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: mbedtls_pk_sign <key_file> <filename>\n");
 
@@ -141,6 +150,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/pk_verify.c b/programs/pkey/pk_verify.c
index f816e92..96a5d28 100644
--- a/programs/pkey/pk_verify.c
+++ b/programs/pkey/pk_verify.c
@@ -58,6 +58,15 @@
 
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: mbedtls_pk_verify <key_file> <filename>\n");
 
@@ -117,6 +126,9 @@
 
 exit:
     mbedtls_pk_free(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/rsa_sign_pss.c b/programs/pkey/rsa_sign_pss.c
index d1afdee..effff25 100644
--- a/programs/pkey/rsa_sign_pss.c
+++ b/programs/pkey/rsa_sign_pss.c
@@ -67,6 +67,15 @@
     mbedtls_pk_init(&pk);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: rsa_sign_pss <key_file> <filename>\n");
 
@@ -149,6 +158,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/pkey/rsa_verify_pss.c b/programs/pkey/rsa_verify_pss.c
index 1718872..a9c75ef 100644
--- a/programs/pkey/rsa_verify_pss.c
+++ b/programs/pkey/rsa_verify_pss.c
@@ -61,6 +61,15 @@
 
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: rsa_verify_pss <key_file> <filename>\n");
 
@@ -127,6 +136,9 @@
 
 exit:
     mbedtls_pk_free(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/ssl/dtls_client.c b/programs/ssl/dtls_client.c
index ad51cbe..beac5d5 100644
--- a/programs/ssl/dtls_client.c
+++ b/programs/ssl/dtls_client.c
@@ -113,11 +113,21 @@
     mbedtls_ssl_config_init(&conf);
     mbedtls_x509_crt_init(&cacert);
     mbedtls_ctr_drbg_init(&ctr_drbg);
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_printf("\n  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -324,12 +334,14 @@
 #endif
 
     mbedtls_net_free(&server_fd);
-
     mbedtls_x509_crt_free(&cacert);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/ssl/dtls_server.c b/programs/ssl/dtls_server.c
index 4310f4e..2128d02 100644
--- a/programs/ssl/dtls_server.c
+++ b/programs/ssl/dtls_server.c
@@ -124,6 +124,16 @@
     mbedtls_entropy_init(&entropy);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_debug_set_threshold(DEBUG_LEVEL);
 #endif
@@ -394,6 +404,9 @@
 #endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     printf("  Press Enter to exit this program.\n");
diff --git a/programs/ssl/mini_client.c b/programs/ssl/mini_client.c
index 688c9fc..27154d8 100644
--- a/programs/ssl/mini_client.c
+++ b/programs/ssl/mini_client.c
@@ -179,8 +179,16 @@
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     mbedtls_x509_crt_init(&ca);
 #endif
-
     mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
         ret = ctr_drbg_seed_failed;
@@ -266,7 +274,6 @@
 
 exit:
     mbedtls_net_free(&server_fd);
-
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
@@ -274,6 +281,9 @@
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     mbedtls_x509_crt_free(&ca);
 #endif
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(ret);
 }
diff --git a/programs/ssl/ssl_client1.c b/programs/ssl/ssl_client1.c
index ffdef3b..933ae75 100644
--- a/programs/ssl/ssl_client1.c
+++ b/programs/ssl/ssl_client1.c
@@ -95,11 +95,21 @@
     mbedtls_ssl_config_init(&conf);
     mbedtls_x509_crt_init(&cacert);
     mbedtls_ctr_drbg_init(&ctr_drbg);
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_printf("\n  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
+
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -274,12 +284,14 @@
 #endif
 
     mbedtls_net_free(&server_fd);
-
     mbedtls_x509_crt_free(&cacert);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/ssl/ssl_context_info.c b/programs/ssl/ssl_context_info.c
index a8b2b47..d503fab 100644
--- a/programs/ssl/ssl_context_info.c
+++ b/programs/ssl/ssl_context_info.c
@@ -23,6 +23,7 @@
 #include MBEDTLS_CONFIG_FILE
 #endif
 #include "mbedtls/debug.h"
+#include "mbedtls/platform.h"
 
 #include <stdio.h>
 #include <stdlib.h>
@@ -939,6 +940,15 @@
     size_t ssl_max_len = SSL_INIT_LEN;
     size_t ssl_len = 0;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        return MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     /* The 'b64_file' is opened when parsing arguments to check that the
      * file name is correct */
     parse_arguments(argc, argv);
@@ -1007,6 +1017,10 @@
         printf("Finished. No valid base64 code found\n");
     }
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     return 0;
 }
 
diff --git a/programs/ssl/ssl_fork_server.c b/programs/ssl/ssl_fork_server.c
index 5a4ac3e..adba12a 100644
--- a/programs/ssl/ssl_fork_server.c
+++ b/programs/ssl/ssl_fork_server.c
@@ -109,6 +109,15 @@
     mbedtls_x509_crt_init(&srvcert);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     signal(SIGCHLD, SIG_IGN);
 
     /*
@@ -369,13 +378,15 @@
 exit:
     mbedtls_net_free(&client_fd);
     mbedtls_net_free(&listen_fd);
-
     mbedtls_x509_crt_free(&srvcert);
     mbedtls_pk_free(&pkey);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  Press Enter to exit this program.\n");
diff --git a/programs/ssl/ssl_mail_client.c b/programs/ssl/ssl_mail_client.c
index 6f1dc1c..89a26fc 100644
--- a/programs/ssl/ssl_mail_client.c
+++ b/programs/ssl/ssl_mail_client.c
@@ -366,6 +366,16 @@
     mbedtls_x509_crt_init(&clicert);
     mbedtls_pk_init(&pkey);
     mbedtls_ctr_drbg_init(&ctr_drbg);
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if (argc < 2) {
 usage:
@@ -455,7 +465,6 @@
     mbedtls_printf("\n  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -792,6 +801,9 @@
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/ssl/ssl_pthread_server.c b/programs/ssl/ssl_pthread_server.c
index 4d7e648..b4a718d 100644
--- a/programs/ssl/ssl_pthread_server.c
+++ b/programs/ssl/ssl_pthread_server.c
@@ -337,6 +337,16 @@
      */
     mbedtls_entropy_init(&entropy);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     /*
      * 1. Load the certificates and private RSA key
      */
@@ -477,14 +487,14 @@
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
     mbedtls_ssl_config_free(&conf);
-
     mbedtls_net_free(&listen_fd);
-
     mbedtls_mutex_free(&debug_mutex);
-
 #if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C)
     mbedtls_memory_buffer_alloc_free();
 #endif
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  Press Enter to exit this program.\n");
diff --git a/programs/ssl/ssl_server.c b/programs/ssl/ssl_server.c
index 8f6a573..69fd0bb 100644
--- a/programs/ssl/ssl_server.c
+++ b/programs/ssl/ssl_server.c
@@ -109,6 +109,16 @@
     mbedtls_entropy_init(&entropy);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_debug_set_threshold(DEBUG_LEVEL);
 #endif
@@ -347,7 +357,6 @@
 
     mbedtls_net_free(&client_fd);
     mbedtls_net_free(&listen_fd);
-
     mbedtls_x509_crt_free(&srvcert);
     mbedtls_pk_free(&pkey);
     mbedtls_ssl_free(&ssl);
@@ -357,6 +366,9 @@
 #endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  Press Enter to exit this program.\n");
diff --git a/programs/x509/cert_app.c b/programs/x509/cert_app.c
index b14b084..294e994 100644
--- a/programs/x509/cert_app.c
+++ b/programs/x509/cert_app.c
@@ -157,6 +157,7 @@
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
     mbedtls_x509_crt_init(&cacert);
+    mbedtls_entropy_init(&entropy);
 #if defined(MBEDTLS_X509_CRL_PARSE_C)
     mbedtls_x509_crl_init(&cacrl);
 #else
@@ -165,6 +166,15 @@
     memset(&cacrl, 0, sizeof(mbedtls_x509_crl));
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -342,7 +352,6 @@
         mbedtls_printf("\n  . Seeding the random number generator...");
         fflush(stdout);
 
-        mbedtls_entropy_init(&entropy);
         if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                          (const unsigned char *) pers,
                                          strlen(pers))) != 0) {
@@ -452,6 +461,9 @@
 #endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/x509/cert_req.c b/programs/x509/cert_req.c
index d7818d7..db200d9 100644
--- a/programs/x509/cert_req.c
+++ b/programs/x509/cert_req.c
@@ -162,6 +162,16 @@
     mbedtls_pk_init(&key);
     mbedtls_ctr_drbg_init(&ctr_drbg);
     memset(buf, 0, sizeof(buf));
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if (argc < 2) {
 usage:
@@ -294,7 +304,6 @@
     mbedtls_printf("  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -365,6 +374,9 @@
     mbedtls_pk_free(&key);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/x509/cert_write.c b/programs/x509/cert_write.c
index ea20144..02ff836 100644
--- a/programs/x509/cert_write.c
+++ b/programs/x509/cert_write.c
@@ -241,6 +241,15 @@
     mbedtls_x509_crt_init(&issuer_crt);
     memset(buf, 0, 1024);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -717,6 +726,9 @@
     mbedtls_mpi_free(&serial);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/x509/crl_app.c b/programs/x509/crl_app.c
index b00f9f3..e3e0577 100644
--- a/programs/x509/crl_app.c
+++ b/programs/x509/crl_app.c
@@ -72,6 +72,15 @@
      */
     mbedtls_x509_crl_init(&crl);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -127,6 +136,9 @@
 
 exit:
     mbedtls_x509_crl_free(&crl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/programs/x509/load_roots.c b/programs/x509/load_roots.c
index faf4ba9..e28f35a 100644
--- a/programs/x509/load_roots.c
+++ b/programs/x509/load_roots.c
@@ -127,6 +127,15 @@
     struct mbedtls_timing_hr_time timer;
     unsigned long ms;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc <= 1) {
         mbedtls_printf(USAGE);
         goto exit;
@@ -191,6 +200,9 @@
     exit_code = MBEDTLS_EXIT_SUCCESS;
 
 exit:
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_exit(exit_code);
 }
 #endif /* necessary configuration */
diff --git a/programs/x509/req_app.c b/programs/x509/req_app.c
index dd7fac7..b447c6a 100644
--- a/programs/x509/req_app.c
+++ b/programs/x509/req_app.c
@@ -72,6 +72,15 @@
      */
     mbedtls_x509_csr_init(&csr);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -127,6 +136,9 @@
 
 exit:
     mbedtls_x509_csr_free(&csr);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");
diff --git a/tests/data_files/Makefile b/tests/data_files/Makefile
index 80172e2..c492558 100644
--- a/tests/data_files/Makefile
+++ b/tests/data_files/Makefile
@@ -888,6 +888,10 @@
 	$(OPENSSL) pkey -in $< -inform DER -out $@
 all_final += ec_prv.pk8param.pem
 
+ec_pub.pem: ec_prv.sec1.der
+	$(OPENSSL) pkey -in $< -inform DER -outform PEM -pubout -out $@
+all_final += ec_pub.pem
+
 ################################################################
 #### Convert PEM keys to DER format
 ################################################################
diff --git a/tests/data_files/ec_pub.der b/tests/data_files/ec_pub.der
index 74c5951..e4e5915 100644
--- a/tests/data_files/ec_pub.der
+++ b/tests/data_files/ec_pub.der
Binary files differ
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 4f24a46..da79062 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -133,12 +133,10 @@
     size_t len;
     mbedtls_pk_debug_item dbg;
 
-    PSA_ASSERT(psa_crypto_init());
-
     mbedtls_pk_init(&pk);
     mbedtls_pk_init(&pk2);
 
-    TEST_ASSERT(psa_crypto_init() == PSA_SUCCESS);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup_opaque(&pk, MBEDTLS_SVC_KEY_ID_INIT) ==
                 MBEDTLS_ERR_PK_BAD_INPUT_DATA);
@@ -212,6 +210,7 @@
     void *options = NULL;
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_VALID_PARAM(mbedtls_pk_free(NULL));
 
@@ -292,6 +291,9 @@
     TEST_ASSERT(mbedtls_pk_parse_public_key(&pk, NULL, 0) ==
                 MBEDTLS_ERR_PK_KEY_INVALID_FORMAT);
 #endif /* MBEDTLS_PK_PARSE_C */
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -302,6 +304,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) == 0);
@@ -322,6 +326,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -342,6 +347,7 @@
     (void) str;
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_INVALID_PARAM(mbedtls_pk_init(NULL));
 
@@ -591,6 +597,8 @@
 #endif /* MBEDTLS_PEM_WRITE_C */
 
 #endif /* MBEDTLS_PK_WRITE_C */
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -600,6 +608,7 @@
     mbedtls_pk_context pk;
 
     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);
@@ -612,6 +621,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -623,6 +633,7 @@
     mbedtls_pk_init(&pub);
     mbedtls_pk_init(&prv);
     mbedtls_pk_init(&alt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_parse_public_keyfile(&pub, pub_file) == 0);
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&prv, prv_file, NULL) == 0);
@@ -638,9 +649,11 @@
     }
 #endif
 
+exit:
     mbedtls_pk_free(&pub);
     mbedtls_pk_free(&prv);
     mbedtls_pk_free(&alt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -663,6 +676,7 @@
 #endif
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     memset(hash_result, 0x00, MBEDTLS_MD_MAX_SIZE);
 
@@ -691,6 +705,7 @@
     mbedtls_pk_restart_free(rs_ctx);
 #endif
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -709,6 +724,7 @@
     size_t hash_len;
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     memset(hash_result, 0x00, sizeof(hash_result));
 
@@ -744,6 +760,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -795,6 +812,7 @@
     mbedtls_pk_init(&pub);
     memset(hash, 0, sizeof(hash));
     memset(sig, 0, sizeof(sig));
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&prv, mbedtls_pk_info_from_type(pk_type)) == 0);
     TEST_ASSERT(mbedtls_ecp_group_load(&mbedtls_pk_ec(prv)->grp, grp_id) == 0);
@@ -872,6 +890,7 @@
     mbedtls_pk_restart_free(&rs_ctx);
     mbedtls_pk_free(&prv);
     mbedtls_pk_free(&pub);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -980,8 +999,8 @@
     memset(&rnd_info,  0, sizeof(mbedtls_test_rnd_pseudo_info));
     memset(output,     0, sizeof(output));
 
-
     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);
 
@@ -997,6 +1016,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1016,6 +1036,7 @@
     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));
 
@@ -1050,6 +1071,7 @@
     mbedtls_mpi_free(&N); mbedtls_mpi_free(&P);
     mbedtls_mpi_free(&Q); mbedtls_mpi_free(&E);
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1064,6 +1086,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));
@@ -1081,6 +1104,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1088,7 +1112,7 @@
 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];
 
     if (SIZE_MAX <= UINT_MAX) {
@@ -1099,6 +1123,7 @@
     memset(sig, 0, sizeof(sig));
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_setup(&pk,
                                  mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
@@ -1118,6 +1143,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1138,7 +1164,9 @@
     int ret = MBEDTLS_ERR_PK_TYPE_MISMATCH;
 
     mbedtls_rsa_init(&raw, MBEDTLS_RSA_PKCS_V15, MBEDTLS_MD_NONE);
-    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));
@@ -1199,7 +1227,9 @@
 
 exit:
     mbedtls_rsa_free(&raw);
-    mbedtls_pk_free(&rsa); mbedtls_pk_free(&alt);
+    mbedtls_pk_free(&rsa);
+    mbedtls_pk_free(&alt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1229,10 +1259,10 @@
      * - 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);
+    USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_pk_setup(&pk,
                                  mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
     TEST_ASSERT(mbedtls_ecp_gen_key(grpid,
diff --git a/tests/suites/test_suite_pkparse.function b/tests/suites/test_suite_pkparse.function
index f2f5e97..d6b6984 100644
--- a/tests/suites/test_suite_pkparse.function
+++ b/tests/suites/test_suite_pkparse.function
@@ -17,6 +17,7 @@
     char *pwd = password;
 
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     if (strcmp(pwd, "NULL") == 0) {
         pwd = NULL;
@@ -35,6 +36,7 @@
 
 exit:
     mbedtls_pk_free(&ctx);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -45,6 +47,7 @@
     int res;
 
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     res = mbedtls_pk_parse_public_keyfile(&ctx, key_file);
 
@@ -59,6 +62,7 @@
 
 exit:
     mbedtls_pk_free(&ctx);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -69,6 +73,7 @@
     int res;
 
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     res = mbedtls_pk_parse_public_keyfile(&ctx, key_file);
 
@@ -83,6 +88,7 @@
 
 exit:
     mbedtls_pk_free(&ctx);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -93,6 +99,7 @@
     int res;
 
     mbedtls_pk_init(&ctx);
+    USE_PSA_INIT();
 
     res = mbedtls_pk_parse_keyfile(&ctx, key_file, password);
 
@@ -107,6 +114,7 @@
 
 exit:
     mbedtls_pk_free(&ctx);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -116,10 +124,12 @@
     mbedtls_pk_context pk;
 
     mbedtls_pk_init(&pk);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_parse_key(&pk, buf->x, buf->len, NULL, 0) == result);
 
 exit:
     mbedtls_pk_free(&pk);
+    USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_pkwrite.data b/tests/suites/test_suite_pkwrite.data
index 83bfdcb..f10bdd6 100644
--- a/tests/suites/test_suite_pkwrite.data
+++ b/tests/suites/test_suite_pkwrite.data
@@ -93,3 +93,15 @@
 Private key write check EC Brainpool 512 bits (DER)
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_BP512R1_ENABLED
 pk_write_key_check:"data_files/ec_bp512_prv.der":TEST_DER
+
+Derive public key EC 192 bits
+depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP192R1_ENABLED
+pk_write_public_from_private:"data_files/ec_prv.sec1.der":"data_files/ec_pub.der"
+
+Derive public key EC 521 bits
+depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP521R1_ENABLED
+pk_write_public_from_private:"data_files/ec_521_prv.der":"data_files/ec_521_pub.der"
+
+Derive public key EC Brainpool 512 bits
+depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_BP512R1_ENABLED
+pk_write_public_from_private:"data_files/ec_bp512_prv.der":"data_files/ec_bp512_pub.der"
diff --git a/tests/suites/test_suite_pkwrite.function b/tests/suites/test_suite_pkwrite.function
index 08ec835..c5391ba 100644
--- a/tests/suites/test_suite_pkwrite.function
+++ b/tests/suites/test_suite_pkwrite.function
@@ -2,6 +2,7 @@
 #include "mbedtls/pk.h"
 #include "mbedtls/pem.h"
 #include "mbedtls/oid.h"
+#include "psa/crypto_sizes.h"
 
 typedef enum {
     TEST_PEM,
@@ -36,6 +37,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 +60,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) {
@@ -97,6 +100,7 @@
     mbedtls_free(buf);
     mbedtls_free(check_buf);
     mbedtls_pk_free(&key);
+    USE_PSA_DONE();
 }
 /* END_HEADER */
 
@@ -120,3 +124,55 @@
     goto exit; /* make the compiler happy */
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void pk_write_public_from_private(char *priv_key_file, char *pub_key_file)
+{
+    mbedtls_pk_context priv_key;
+    uint8_t *derived_key_raw = NULL;
+    size_t derived_key_len = 0;
+    uint8_t *pub_key_raw = NULL;
+    size_t pub_key_len = 0;
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_svc_key_id_t opaque_key_id = MBEDTLS_SVC_KEY_ID_INIT;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+    mbedtls_pk_init(&priv_key);
+    USE_PSA_INIT();
+
+    TEST_EQUAL(mbedtls_pk_parse_keyfile(&priv_key, priv_key_file, NULL), 0);
+    TEST_EQUAL(mbedtls_pk_load_file(pub_key_file, &pub_key_raw,
+                                    &pub_key_len), 0);
+
+    derived_key_len = pub_key_len;
+    ASSERT_ALLOC(derived_key_raw, derived_key_len);
+
+    TEST_EQUAL(mbedtls_pk_write_pubkey_der(&priv_key, derived_key_raw,
+                                           derived_key_len), pub_key_len);
+
+    ASSERT_COMPARE(derived_key_raw, derived_key_len,
+                   pub_key_raw, pub_key_len);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_platform_zeroize(derived_key_raw, sizeof(derived_key_raw));
+
+    TEST_EQUAL(mbedtls_pk_wrap_as_opaque(&priv_key, &opaque_key_id,
+                                         PSA_ALG_NONE), 0);
+
+    TEST_EQUAL(mbedtls_pk_write_pubkey_der(&priv_key, derived_key_raw,
+                                           derived_key_len), pub_key_len);
+
+    ASSERT_COMPARE(derived_key_raw, derived_key_len,
+                   pub_key_raw, pub_key_len);
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+exit:
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_destroy_key(opaque_key_id);
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+    mbedtls_free(derived_key_raw);
+    mbedtls_free(pub_key_raw);
+    mbedtls_pk_free(&priv_key);
+    USE_PSA_DONE();
+}
+/* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index aad1bf2..03cc2ff 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -2428,7 +2428,7 @@
 aead_encrypt_decrypt:PSA_KEY_TYPE_AES:"D7828D13B2B0BDC325A76236DF93CC6B":PSA_ALG_CTR:"000102030405060708090A0B0C0D0E0F":"":"":PSA_ERROR_NOT_SUPPORTED
 
 PSA AEAD encrypt/decrypt: invalid algorithm (ChaCha20)
-depends_on:MBEDTLS_CHACHA20_C
+depends_on:PSA_WANT_KEY_TYPE_CHACHA20
 aead_encrypt_decrypt:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_STREAM_CIPHER:"":"":"":PSA_ERROR_NOT_SUPPORTED
 
 PSA signature size: RSA keypair, 1024 bits, PKCS#1 v1.5 raw
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 0b03beb..bc999f1 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -5,6 +5,8 @@
 
 #include <test/constant_flow.h>
 
+#define SSL_MESSAGE_QUEUE_INIT      { NULL, 0, 0, 0 }
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -20,6 +22,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. */
@@ -72,8 +75,8 @@
 
 
 exit:
-
     mbedtls_test_ssl_buffer_free(&buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -108,6 +111,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
@@ -182,10 +186,10 @@
     }
 
 exit:
-
     mbedtls_free(input);
     mbedtls_free(output);
     mbedtls_test_ssl_buffer_free(&buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -203,6 +207,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);
@@ -217,8 +222,8 @@
     mbedtls_test_mock_socket_close(&socket);
 
 exit:
-
     mbedtls_test_mock_socket_close(&socket);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -252,6 +257,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 */
@@ -309,9 +315,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 */
 
@@ -349,6 +355,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 */
@@ -437,17 +444,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);
@@ -460,14 +468,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 */
@@ -487,14 +497,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) */
@@ -513,14 +525,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
@@ -547,16 +561,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 */
@@ -566,6 +582,7 @@
                 == (int) buffer_len);
 exit:
     mbedtls_test_ssl_message_queue_free(&queue);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -580,6 +597,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);
@@ -618,6 +636,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -630,8 +649,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,
@@ -677,6 +698,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -689,8 +711,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,
@@ -741,6 +765,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -753,8 +778,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,
@@ -793,6 +820,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -805,8 +833,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,
@@ -857,6 +887,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -869,8 +900,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,
@@ -915,6 +948,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -927,8 +961,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,
@@ -975,6 +1011,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -987,8 +1024,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,
@@ -1062,6 +1101,7 @@
 exit:
     mbedtls_test_message_socket_close(&server_context);
     mbedtls_test_message_socket_close(&client_context);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1074,6 +1114,7 @@
 
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_config_defaults(&conf,
                                             MBEDTLS_SSL_IS_CLIENT,
@@ -1091,8 +1132,10 @@
     memcpy(ssl.in_ctr + 2, new->x, 6);
     TEST_ASSERT(mbedtls_ssl_dtls_replay_check(&ssl) == ret);
 
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1100,12 +1143,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 */
 
@@ -1130,6 +1177,8 @@
     mbedtls_record rec, rec_backup;
 
     mbedtls_ssl_init(&ssl);
+    USE_PSA_INIT();
+
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
     TEST_ASSERT(mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
@@ -1230,8 +1279,8 @@
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_transform_free(&t0);
     mbedtls_ssl_transform_free(&t1);
-
     mbedtls_free(buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1281,6 +1330,7 @@
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
                                                   etm, tag_mode, ver,
                                                   (size_t) cid0_len,
@@ -1390,8 +1440,8 @@
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_transform_free(&t0);
     mbedtls_ssl_transform_free(&t1);
-
     mbedtls_free(buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1431,6 +1481,7 @@
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
     mbedtls_ssl_transform_init(&t1);
+    USE_PSA_INIT();
 
     /* Set up transforms with dummy keys */
     TEST_ASSERT(mbedtls_test_ssl_build_transforms(&t0, &t1, cipher_type, hash_id,
@@ -1601,6 +1652,7 @@
     mbedtls_ssl_transform_free(&t1);
     mbedtls_free(buf);
     mbedtls_free(buf_save);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1788,9 +1840,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 */
     TEST_ASSERT(mbedtls_test_ssl_populate_session(
@@ -1875,6 +1927,7 @@
     mbedtls_ssl_session_free(&original);
     mbedtls_ssl_session_free(&restored);
     mbedtls_free(buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1888,8 +1941,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 */
     TEST_ASSERT(mbedtls_test_ssl_populate_session(
@@ -1926,6 +1979,7 @@
     mbedtls_ssl_session_free(&session);
     mbedtls_free(buf1);
     mbedtls_free(buf2);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1939,8 +1993,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 */
     TEST_ASSERT(mbedtls_test_ssl_populate_session(
@@ -1962,6 +2016,7 @@
 exit:
     mbedtls_ssl_session_free(&session);
     mbedtls_free(buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1975,8 +2030,8 @@
     /*
      * Test that session_load() fails cleanly on small buffers
      */
-
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Prepare serialized session data */
     TEST_ASSERT(mbedtls_test_ssl_populate_session(
@@ -2004,6 +2059,7 @@
     mbedtls_ssl_session_free(&session);
     mbedtls_free(good_buf);
     mbedtls_free(bad_buf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2024,6 +2080,7 @@
                                       corrupt_config == 1 };
 
     mbedtls_ssl_session_init(&session);
+    USE_PSA_INIT();
 
     /* Infer length of serialized session. */
     TEST_ASSERT(mbedtls_ssl_session_save(&session,
@@ -2065,7 +2122,7 @@
             *byte ^= corrupted_bit;
         }
     }
-
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2076,6 +2133,7 @@
     mbedtls_test_ssl_endpoint ep;
     int ret = -1;
 
+    USE_PSA_INIT();
     ret = mbedtls_test_ssl_endpoint_init(NULL, endpoint_type, MBEDTLS_PK_RSA,
                                          NULL, NULL, NULL, NULL);
     TEST_ASSERT(MBEDTLS_ERR_SSL_BAD_INPUT_DATA == ret);
@@ -2089,6 +2147,7 @@
 
 exit:
     mbedtls_test_ssl_endpoint_free(&ep, NULL);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2114,6 +2173,8 @@
         MBEDTLS_PK_RSA, NULL, NULL, NULL, NULL);
     TEST_ASSERT(ret == 0);
 
+    USE_PSA_INIT();
+
     ret = mbedtls_test_mock_socket_connect(&(base_ep.socket),
                                            &(second_ep.socket),
                                            BUFFSIZE);
@@ -2133,6 +2194,7 @@
 exit:
     mbedtls_test_ssl_endpoint_free(&base_ep, NULL);
     mbedtls_test_ssl_endpoint_free(&second_ep, NULL);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2331,7 +2393,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;
 }
@@ -2342,7 +2403,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;
 }
@@ -2358,7 +2418,7 @@
 
     mbedtls_ecp_group_id curve_list[] = { MBEDTLS_ECP_DP_SECP256R1,
                                           MBEDTLS_ECP_DP_NONE };
-    USE_PSA_INIT();
+
     mbedtls_platform_zeroize(&client, sizeof(client));
     mbedtls_platform_zeroize(&server, sizeof(server));
 
@@ -2374,6 +2434,8 @@
                                               MBEDTLS_PK_ECDSA, NULL, NULL,
                                               NULL, NULL), 0);
 
+    USE_PSA_INIT();
+
     TEST_EQUAL(mbedtls_test_mock_socket_connect(&(client.socket),
                                                 &(server.socket),
                                                 BUFFSIZE), 0);
@@ -2422,6 +2484,7 @@
     size_t len;
 
     mbedtls_ssl_init(&ssl);
+    USE_PSA_INIT();
     mbedtls_ssl_config_init(&conf);
     TEST_EQUAL(mbedtls_ssl_config_defaults(&conf, MBEDTLS_SSL_IS_SERVER,
                                            MBEDTLS_SSL_TRANSPORT_DATAGRAM,
@@ -2437,7 +2500,9 @@
                                                     &len),
                exp_ret);
 
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_x509parse.function b/tests/suites/test_suite_x509parse.function
index db7c086..6831b0d 100644
--- a/tests/suites/test_suite_x509parse.function
+++ b/tests/suites/test_suite_x509parse.function
@@ -409,6 +409,7 @@
 
     mbedtls_x509_crt_init(&crt);
     memset(buf, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -432,6 +433,7 @@
 exit:
 
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -444,6 +446,7 @@
 
     mbedtls_x509_crt_init(&crt);
     memset(buf, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
     res = mbedtls_x509_crt_info(buf, 2000, "", &crt);
@@ -455,6 +458,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -467,6 +471,7 @@
 
     mbedtls_x509_crl_init(&crl);
     memset(buf, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crl_parse_file(&crl, crl_file) == 0);
     res = mbedtls_x509_crl_info(buf, 2000, "", &crl);
@@ -478,6 +483,7 @@
 
 exit:
     mbedtls_x509_crl_free(&crl);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -489,11 +495,13 @@
 
     mbedtls_x509_crl_init(&crl);
     memset(buf, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crl_parse_file(&crl, crl_file) == result);
 
 exit:
     mbedtls_x509_crl_free(&crl);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -506,6 +514,7 @@
 
     mbedtls_x509_csr_init(&csr);
     memset(buf, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_csr_parse_file(&csr, csr_file) == 0);
     res = mbedtls_x509_csr_info(buf, 2000, "", &csr);
@@ -517,6 +526,7 @@
 
 exit:
     mbedtls_x509_csr_free(&csr);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -526,6 +536,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);
@@ -533,6 +544,9 @@
     TEST_ASSERT(res >= 0);
 
     TEST_ASSERT(strcmp(buf, result_str) == 0);
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -556,10 +570,10 @@
      * - 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);
+    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);
@@ -589,6 +603,7 @@
     mbedtls_x509_crt_restart_free(&rs_ctx);
     mbedtls_x509_crt_free(&crt);
     mbedtls_x509_crt_free(&ca);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -695,6 +710,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);
@@ -712,6 +728,7 @@
 exit:
     mbedtls_x509_crt_free(&crt);
     mbedtls_x509_crt_free(&ca);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -762,6 +779,7 @@
 
     mbedtls_x509_crt_init(&crt);
     memset(buf, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
     if (strcmp(entity, "subject") == 0) {
@@ -779,6 +797,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -794,6 +813,7 @@
 
     mbedtls_x509_crt_init(&crt);
     memset(buf, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
     crt.subject.next->val.p = (unsigned char *) new_subject_ou;
@@ -810,19 +830,21 @@
     }
 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;
     mbedtls_x509_name *allocated, *prev;
     int ret;
 
+    USE_PSA_INIT();
     memset(&head, 0, sizeof(head));
 
     name = mbedtls_test_unhexify_alloc(rdn_sequence, &name_len);
@@ -843,6 +865,9 @@
     TEST_EQUAL(ret, exp_ret);
 
     mbedtls_free(name);
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -852,6 +877,7 @@
     mbedtls_x509_crt   crt;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -865,6 +891,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -874,6 +901,7 @@
     mbedtls_x509_crt   crt;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -887,6 +915,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -896,11 +925,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 */
 
@@ -913,6 +944,7 @@
 
     mbedtls_x509_crt_init(&crt);
     memset(output, 0, 2000);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_der(&crt, buf->x, buf->len) == (result));
     if ((result) == 0) {
@@ -970,6 +1002,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -987,6 +1020,7 @@
 
     mbedtls_x509_crt_init(&crt);
     memset(output, 0, 2000);
+    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));
@@ -1016,6 +1050,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1028,7 +1063,7 @@
 
     mbedtls_x509_crl_init(&crl);
     memset(output, 0, 2000);
-
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crl_parse(&crl, buf->x, buf->len) == (result));
     if ((result) == 0) {
@@ -1042,6 +1077,7 @@
 
 exit:
     mbedtls_x509_crl_free(&crl);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1054,6 +1090,7 @@
 
     mbedtls_x509_csr_init(&csr);
     memset(my_out, 0, sizeof(my_out));
+    USE_PSA_INIT();
 
     my_ret = mbedtls_x509_csr_parse_der(&csr, csr_der->x, csr_der->len);
     TEST_ASSERT(my_ret == ref_ret);
@@ -1066,6 +1103,7 @@
 
 exit:
     mbedtls_x509_csr_free(&csr);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1076,6 +1114,7 @@
     int i;
 
     mbedtls_x509_crt_init(&chain);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_path(&chain, crt_path) == ret);
 
@@ -1090,6 +1129,7 @@
 
 exit:
     mbedtls_x509_crt_free(&chain);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1106,10 +1146,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);
-
     USE_PSA_INIT();
 
     /* Load trusted root */
@@ -1148,7 +1186,6 @@
 
     mbedtls_x509_crt_init(&chain);
     mbedtls_x509_crt_init(&trusted);
-
     USE_PSA_INIT();
 
     while ((act = mystrsep(&chain_paths, " ")) != NULL) {
@@ -1188,7 +1225,7 @@
     const char *desc = NULL;
     int ret;
 
-
+    USE_PSA_INIT();
     oid.tag = MBEDTLS_ASN1_OID;
     oid.p   = buf->x;
     oid.len   = buf->len;
@@ -1203,6 +1240,9 @@
         TEST_ASSERT(desc != NULL);
         TEST_ASSERT(strcmp(desc, ref_desc) == 0);
     }
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1212,6 +1252,7 @@
     mbedtls_x509_buf oid;
     char num_buf[100];
 
+    USE_PSA_INIT();
     memset(num_buf, 0x2a, sizeof(num_buf));
 
     oid.tag = MBEDTLS_ASN1_OID;
@@ -1226,6 +1267,9 @@
         TEST_ASSERT(num_buf[ret] == 0);
         TEST_ASSERT(strcmp(num_buf, numstr) == 0);
     }
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1235,6 +1279,7 @@
     mbedtls_x509_crt crt;
 
     mbedtls_x509_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1242,6 +1287,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1252,7 +1298,7 @@
     mbedtls_x509_crt crt;
 
     mbedtls_x509_crt_init(&crt);
-
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_x509_crt_parse_file(&crt, crt_file) == 0);
 
@@ -1261,6 +1307,7 @@
 
 exit:
     mbedtls_x509_crt_free(&crt);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1273,6 +1320,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);
@@ -1290,6 +1338,9 @@
         TEST_ASSERT(min  == time.min);
         TEST_ASSERT(sec  == time.sec);
     }
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1303,6 +1354,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;
@@ -1319,13 +1372,17 @@
     }
 
 exit:
-    ;;
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_X509_CRT_PARSE_C:MBEDTLS_SELF_TEST */
 void x509_selftest()
 {
+    USE_PSA_INIT();
     TEST_ASSERT(mbedtls_x509_self_test(1) == 0);
+
+exit:
+    USE_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_x509write.function b/tests/suites/test_suite_x509write.function
index cb3f6a5..b4509e2 100644
--- a/tests/suites/test_suite_x509write.function
+++ b/tests/suites/test_suite_x509write.function
@@ -145,8 +145,9 @@
     memset(&rnd_info, 0x2a, sizeof(mbedtls_test_rnd_pseudo_info));
 
     mbedtls_x509write_csr_init(&req);
-
     mbedtls_pk_init(&key);
+    USE_PSA_INIT();
+
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&key, key_file, NULL) == 0);
 
     mbedtls_x509write_csr_set_md_alg(&req, md_type);
@@ -197,6 +198,7 @@
 exit:
     mbedtls_x509write_csr_free(&req);
     mbedtls_pk_free(&key);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -246,12 +248,11 @@
     buf[pem_len] = '\0';
     TEST_ASSERT(x509_crt_verifycsr(buf, pem_len + 1) == 0);
 
-
 exit:
     mbedtls_x509write_csr_free(&req);
     mbedtls_pk_free(&key);
     psa_destroy_key(key_id);
-    PSA_DONE();
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -287,6 +288,7 @@
     mbedtls_pk_init(&issuer_key_alt);
 
     mbedtls_x509write_crt_init(&crt);
+    USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_pk_parse_keyfile(&subject_key, subject_key_file,
                                          subject_pwd) == 0);
@@ -422,6 +424,7 @@
     mbedtls_pk_free(&subject_key);
     mbedtls_pk_free(&issuer_key);
     mbedtls_mpi_free(&serial);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -432,6 +435,7 @@
     mbedtls_mpi serial_mpi;
     uint8_t invalid_serial[MBEDTLS_X509_RFC5280_MAX_SERIAL_LEN + 1];
 
+    USE_PSA_INIT();
     memset(invalid_serial, 0x01, sizeof(invalid_serial));
 
     mbedtls_mpi_init(&serial_mpi);
@@ -442,6 +446,7 @@
 
 exit:
     mbedtls_mpi_free(&serial_mpi);
+    USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -455,6 +460,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));
@@ -488,5 +495,6 @@
         parsed_cur = parsed_cur->next;
         mbedtls_free(parsed_prv);
     }
+    USE_PSA_DONE();
 }
 /* END_CASE */