Init PSA in pkey programs

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/programs/pkey/gen_key.c b/programs/pkey/gen_key.c
index 1a6463d..20fc17f 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,7 @@
     mbedtls_pk_free(&key);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+    mbedtls_psa_crypto_free();
 
 #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..82a4de1 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,7 @@
 #endif
 
     mbedtls_pk_free(&pk);
+    mbedtls_psa_crypto_free();
     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..4da4f65 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,7 @@
     mbedtls_mpi_free(&DQ); mbedtls_mpi_free(&QP);
 
     mbedtls_pk_free(&key);
+    mbedtls_psa_crypto_free();
 
 #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..9ab01ec 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,7 @@
     mbedtls_pk_free(&pk);
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
+    mbedtls_psa_crypto_free();
 
 #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..f3c3274 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,7 @@
     mbedtls_pk_free(&pk);
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
+    mbedtls_psa_crypto_free();
 
 #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..a52584a 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,7 @@
     mbedtls_pk_free(&pk);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+    mbedtls_psa_crypto_free();
 
 #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..49b5204 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,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    mbedtls_psa_crypto_free();
 
 #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..63bc755 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,7 @@
     mbedtls_pk_free(&pk);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+    mbedtls_psa_crypto_free();
 
 #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..9fb0401 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,7 @@
 
 exit:
     mbedtls_pk_free(&pk);
+    mbedtls_psa_crypto_free();
 
 #if defined(_WIN32)
     mbedtls_printf("  + Press Enter to exit this program.\n");