From d38480b0e02dacc8c52534cfe0798ff23a358727 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Tue, 19 Mar 2024 13:47:27 +0100
Subject: [PATCH 01/14] test_suite_pk: reshape pk_psa_sign()

The behavior of the functions is kept intact. Changes concern:
- generate the initial PK context using PSA parameters only; this
  allows to remove 1 input parameter for the test function.
- add/fix comments.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.data     | 20 +++++-----
 tests/suites/test_suite_pk.function | 59 +++++++++++++----------------
 2 files changed, 36 insertions(+), 43 deletions(-)

diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index a0dacf0a38..102aee2d5d 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -648,44 +648,44 @@ pk_sign_verify_restart:MBEDTLS_PK_ECKEY:MBEDTLS_ECP_DP_SECP256R1:"C9AFA9D845BA75
 
 PSA wrapped sign: SECP256R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP256R1
-pk_psa_sign:MBEDTLS_ECP_DP_SECP256R1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256
 
 PSA wrapped sign: SECP384R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP384R1
-pk_psa_sign:MBEDTLS_ECP_DP_SECP384R1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):384
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):384
 
 PSA wrapped sign: SECP521R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP521R1
-pk_psa_sign:MBEDTLS_ECP_DP_SECP521R1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):521
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):521
 
 PSA wrapped sign: SECP192K1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP192K1
-pk_psa_sign:MBEDTLS_ECP_DP_SECP192K1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):192
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):192
 
 ## Currently buggy: https://github.com/ARMmbed/mbed-crypto/issues/336
 # PSA wrapped sign: SECP224K1
 # depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP224K1
-# pk_psa_sign:MBEDTLS_ECP_DP_SECP224K1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):224
+# pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):224
 
 PSA wrapped sign: SECP256K1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP256K1
-pk_psa_sign:MBEDTLS_ECP_DP_SECP256K1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):256
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):256
 
 PSA wrapped sign: BP256R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_BP256R1
-pk_psa_sign:MBEDTLS_ECP_DP_BP256R1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):256
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):256
 
 PSA wrapped sign: BP384R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_BP384R1
-pk_psa_sign:MBEDTLS_ECP_DP_BP384R1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):384
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):384
 
 PSA wrapped sign: BP512R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_BP512R1
-pk_psa_sign:MBEDTLS_ECP_DP_BP512R1:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):512
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):512
 
 PSA wrapped sign: RSA PKCS1 v1.5
 depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15:MBEDTLS_GENPRIME:MBEDTLS_PK_WRITE_C
-pk_psa_sign:1024:PSA_KEY_TYPE_RSA_KEY_PAIR:1024
+pk_psa_sign:PSA_KEY_TYPE_RSA_KEY_PAIR:1024
 
 PK sign ext: RSA2048, PK_RSA, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index afc1e342c2..e1a8e1ce6f 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1817,7 +1817,7 @@ exit:
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_MD_CAN_SHA256:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_TEST_PK_PSA_SIGN */
-void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
+void pk_psa_sign(int psa_type, int bits)
 {
     mbedtls_pk_context pk;
     unsigned char hash[32];
@@ -1831,50 +1831,45 @@ void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     /*
-     * This tests making signatures with a wrapped PSA key:
-     * - generate a fresh ECP/RSA legacy PK context
-     * - wrap it in a PK context and make a signature this way
-     * - extract the public key
-     * - parse it to a PK context and verify the signature this way
+     * Following checks are perfomed:
+     * - create an RSA/EC opaque context;
+     * - sign with opaque context for both EC and RSA keys;
+     * - [EC only] verify with opaque context;
+     * - verify that public keys of opaque and non-opaque contexts match;
+     * - verify with non-opaque context.
      */
 
     mbedtls_pk_init(&pk);
     USE_PSA_INIT();
 
+    /* Create the legacy EC/RSA PK context. */
 #if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
     if (PSA_KEY_TYPE_IS_RSA(psa_type)) {
-        /* Create legacy RSA public/private key in PK context. */
         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),
-                                        mbedtls_test_rnd_std_rand, NULL,
-                                        curve_or_keybits, 3) == 0);
+        TEST_ASSERT(mbedtls_rsa_gen_key(mbedtls_pk_rsa(pk), mbedtls_test_rnd_std_rand, NULL,
+                                        bits, 3) == 0);
     } else
 #endif /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
-        mbedtls_ecp_group_id grpid = curve_or_keybits;
-
-        /* Create legacy EC public/private key in PK context. */
-        TEST_ASSERT(mbedtls_pk_setup(&pk,
-                                     mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
+        mbedtls_ecp_group_id grpid = mbedtls_ecc_group_from_psa(psa_type, bits);
+        TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
         TEST_ASSERT(pk_genkey(&pk, grpid) == 0);
     } else
 #endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
     {
-        (void) curve_or_keybits;
         TEST_ASSUME(!"Opaque PK key not supported in this configuration");
     }
 
-    /* Export underlying public key for re-importing in a legacy context. */
+    /* Export public key from the non-opaque PK context we just created. */
 #if defined(MBEDTLS_PK_WRITE_C)
-    ret = mbedtls_pk_write_pubkey_der(&pk, pkey_legacy,
-                                      sizeof(pkey_legacy));
+    ret = mbedtls_pk_write_pubkey_der(&pk, pkey_legacy, sizeof(pkey_legacy));
     TEST_ASSERT(ret >= 0);
     klen_legacy = (size_t) ret;
     /* mbedtls_pk_write_pubkey_der() writes backwards in the data buffer. */
     pkey_legacy_start = pkey_legacy + sizeof(pkey_legacy) - klen_legacy;
-#else
+#else /* MBEDTLS_PK_WRITE_C */
     ret = mbedtls_ecp_point_write_binary(&(mbedtls_pk_ec_ro(pk)->grp),
                                          &(mbedtls_pk_ec_ro(pk)->Q),
                                          MBEDTLS_ECP_PF_UNCOMPRESSED,
@@ -1884,7 +1879,7 @@ void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
     pkey_legacy_start = pkey_legacy;
 #endif /* MBEDTLS_PK_WRITE_C */
 
-    /* Turn PK context into an opaque one. */
+    /* Turn the PK context into an opaque one. */
     TEST_EQUAL(mbedtls_pk_get_psa_attributes(&pk, PSA_KEY_USAGE_SIGN_HASH, &attributes), 0);
     TEST_EQUAL(mbedtls_pk_import_into_psa(&pk, &attributes, &key_id), 0);
     mbedtls_pk_free(&pk);
@@ -1893,13 +1888,12 @@ void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
 
     PSA_ASSERT(psa_get_key_attributes(key_id, &attributes));
     TEST_EQUAL(psa_get_key_type(&attributes), (psa_key_type_t) psa_type);
-    TEST_EQUAL(psa_get_key_bits(&attributes), (size_t) expected_bits);
-    TEST_EQUAL(psa_get_key_lifetime(&attributes),
-               PSA_KEY_LIFETIME_VOLATILE);
+    TEST_EQUAL(psa_get_key_bits(&attributes), (size_t) bits);
+    TEST_EQUAL(psa_get_key_lifetime(&attributes), PSA_KEY_LIFETIME_VOLATILE);
 
+    /* Sign with the opaque context. */
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
-
     TEST_ASSERT(mbedtls_pk_sign(&pk, MBEDTLS_MD_SHA256,
                                 hash, sizeof(hash), sig, sizeof(sig), &sig_len,
                                 NULL, NULL) == 0);
@@ -1909,7 +1903,7 @@ void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
                                       hash, sizeof(hash), sig, sig_len) == 0);
     }
 
-    /* Export underlying public key for re-importing in a psa context. */
+    /* Export public key from the opaque PK context. */
 #if defined(MBEDTLS_PK_WRITE_C)
     ret = mbedtls_pk_write_pubkey_der(&pk, pkey_psa,
                                       sizeof(pkey_psa));
@@ -1926,18 +1920,21 @@ void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
     pkey_psa_start = pkey_psa;
 #endif /* MBEDTLS_PK_WRITE_C */
 
+    /* Check that the public keys of opaque and non-opaque PK contexts match. */
     TEST_ASSERT(klen_psa == klen_legacy);
     TEST_ASSERT(memcmp(pkey_psa_start, pkey_legacy_start, klen_psa) == 0);
 
+    /* Destroy the opaque PK context. */
     mbedtls_pk_free(&pk);
     TEST_ASSERT(PSA_SUCCESS == psa_destroy_key(key_id));
 
-    mbedtls_pk_init(&pk);
-
-    /* If we used "pk_write" previously, then we go for a "pk_parse" here;
+    /* Create a new non-opaque PK context to verify the signature.
+     *
+     * Note: if we used "pk_write" previously, then we go for a "pk_parse" here;
      * otherwise if we went for "ecp_point_write_binary" then we'll go
      * for a "ecp_point_read_binary" here. This allows to drop dependencies
      * on "PK_WRITE" and "PK_PARSE" if required */
+    mbedtls_pk_init(&pk);
 #if defined(MBEDTLS_PK_WRITE_C) && defined(MBEDTLS_PK_PARSE_C)
     TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, pkey_legacy_start,
                                            klen_legacy), 0);
@@ -1955,10 +1952,6 @@ void pk_psa_sign(int curve_or_keybits, int psa_type, int expected_bits)
                                   hash, sizeof(hash), sig, sig_len) == 0);
 
 exit:
-    /*
-     * Key attributes may have been returned by psa_get_key_attributes()
-     * thus reset them as required.
-     */
     psa_reset_key_attributes(&attributes);
 
     mbedtls_pk_free(&pk);

From c262561424609a6b7807ebc927c57d28e39b9f6c Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Tue, 19 Mar 2024 14:37:05 +0100
Subject: [PATCH 02/14] test_suite_pk: rename some variables in pk_psa_sign()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.function | 65 ++++++++++++++---------------
 1 file changed, 31 insertions(+), 34 deletions(-)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index e1a8e1ce6f..8f5e7f6718 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1822,13 +1822,15 @@ void pk_psa_sign(int psa_type, int bits)
     mbedtls_pk_context pk;
     unsigned char hash[32];
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
-    unsigned char pkey_legacy[200];
-    unsigned char pkey_psa[200];
-    unsigned char *pkey_legacy_start, *pkey_psa_start;
-    size_t sig_len, klen_legacy, klen_psa;
+    unsigned char legacy_pub_key[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
+    unsigned char opaque_pub_key[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
+    size_t sig_len, legacy_pub_key_len, opaque_pub_key_len;
     int ret;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+#if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
+    mbedtls_ecp_group_id ecp_grp_id;
+#endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
 
     /*
      * Following checks are perfomed:
@@ -1853,9 +1855,9 @@ void pk_psa_sign(int psa_type, int bits)
 #endif /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
-        mbedtls_ecp_group_id grpid = mbedtls_ecc_group_from_psa(psa_type, bits);
+        ecp_grp_id = mbedtls_ecc_group_from_psa(psa_type, bits);
         TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
-        TEST_ASSERT(pk_genkey(&pk, grpid) == 0);
+        TEST_ASSERT(pk_genkey(&pk, ecp_grp_id) == 0);
     } else
 #endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
     {
@@ -1864,19 +1866,21 @@ void pk_psa_sign(int psa_type, int bits)
 
     /* Export public key from the non-opaque PK context we just created. */
 #if defined(MBEDTLS_PK_WRITE_C)
-    ret = mbedtls_pk_write_pubkey_der(&pk, pkey_legacy, sizeof(pkey_legacy));
+    ret = mbedtls_pk_write_pubkey_der(&pk, legacy_pub_key, sizeof(legacy_pub_key));
     TEST_ASSERT(ret >= 0);
-    klen_legacy = (size_t) ret;
-    /* mbedtls_pk_write_pubkey_der() writes backwards in the data buffer. */
-    pkey_legacy_start = pkey_legacy + sizeof(pkey_legacy) - klen_legacy;
+    legacy_pub_key_len = (size_t) ret;
+    /* mbedtls_pk_write_pubkey_der() writes backwards in the data buffer so we
+     * shift data back to the beginning of the buffer. */
+    memmove(legacy_pub_key,
+            legacy_pub_key + sizeof(legacy_pub_key) - legacy_pub_key_len,
+            legacy_pub_key_len);
 #else /* MBEDTLS_PK_WRITE_C */
     ret = mbedtls_ecp_point_write_binary(&(mbedtls_pk_ec_ro(pk)->grp),
                                          &(mbedtls_pk_ec_ro(pk)->Q),
                                          MBEDTLS_ECP_PF_UNCOMPRESSED,
-                                         &klen_legacy, pkey_legacy,
-                                         sizeof(pkey_legacy));
+                                         &legacy_pub_key_len, legacy_pub_key,
+                                         sizeof(legacy_pub_key));
     TEST_EQUAL(ret, 0);
-    pkey_legacy_start = pkey_legacy;
 #endif /* MBEDTLS_PK_WRITE_C */
 
     /* Turn the PK context into an opaque one. */
@@ -1905,26 +1909,23 @@ void pk_psa_sign(int psa_type, int bits)
 
     /* Export public key from the opaque PK context. */
 #if defined(MBEDTLS_PK_WRITE_C)
-    ret = mbedtls_pk_write_pubkey_der(&pk, pkey_psa,
-                                      sizeof(pkey_psa));
+    ret = mbedtls_pk_write_pubkey_der(&pk, opaque_pub_key, sizeof(opaque_pub_key));
     TEST_ASSERT(ret >= 0);
-    klen_psa = (size_t) ret;
+    opaque_pub_key_len = (size_t) ret;
     /* mbedtls_pk_write_pubkey_der() writes backwards in the data buffer. */
-    pkey_psa_start = pkey_psa + sizeof(pkey_psa) - klen_psa;
+    memmove(opaque_pub_key,
+            opaque_pub_key + sizeof(opaque_pub_key) - opaque_pub_key_len,
+            opaque_pub_key_len);
 #else
-    psa_status_t status;
-
-    status = psa_export_public_key(key_id, pkey_psa, sizeof(pkey_psa),
-                                   &klen_psa);
-    TEST_EQUAL(status, PSA_SUCCESS);
-    pkey_psa_start = pkey_psa;
+    TEST_EQUAL(psa_export_public_key(key_id, opaque_pub_key, sizeof(opaque_pub_key),
+                                     &opaque_pub_key_len), PSA_SUCCESS);
 #endif /* MBEDTLS_PK_WRITE_C */
 
     /* Check that the public keys of opaque and non-opaque PK contexts match. */
-    TEST_ASSERT(klen_psa == klen_legacy);
-    TEST_ASSERT(memcmp(pkey_psa_start, pkey_legacy_start, klen_psa) == 0);
+    TEST_EQUAL(opaque_pub_key_len, legacy_pub_key_len);
+    TEST_MEMORY_COMPARE(opaque_pub_key, opaque_pub_key_len, legacy_pub_key, legacy_pub_key_len);
 
-    /* Destroy the opaque PK context. */
+    /* Destroy the opaque PK context and the wrapped PSA key. */
     mbedtls_pk_free(&pk);
     TEST_ASSERT(PSA_SUCCESS == psa_destroy_key(key_id));
 
@@ -1936,17 +1937,13 @@ void pk_psa_sign(int psa_type, int bits)
      * on "PK_WRITE" and "PK_PARSE" if required */
     mbedtls_pk_init(&pk);
 #if defined(MBEDTLS_PK_WRITE_C) && defined(MBEDTLS_PK_PARSE_C)
-    TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, pkey_legacy_start,
-                                           klen_legacy), 0);
+    TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, legacy_pub_key, legacy_pub_key_len), 0);
 #else
-    TEST_EQUAL(mbedtls_pk_setup(&pk,
-                                mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
-    TEST_EQUAL(mbedtls_ecp_group_load(
-                   &(mbedtls_pk_ec_rw(pk)->grp),
-                   (mbedtls_ecp_group_id) curve_or_keybits), 0);
+    TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
+    TEST_EQUAL(mbedtls_ecp_group_load(&(mbedtls_pk_ec_rw(pk)->grp), ecp_grp_id), 0);
     TEST_EQUAL(mbedtls_ecp_point_read_binary(&(mbedtls_pk_ec_ro(pk)->grp),
                                              &(mbedtls_pk_ec_rw(pk)->Q),
-                                             pkey_legacy_start, klen_legacy), 0);
+                                             legacy_pub_key, legacy_pub_key_len), 0);
 #endif
     TEST_ASSERT(mbedtls_pk_verify(&pk, MBEDTLS_MD_SHA256,
                                   hash, sizeof(hash), sig, sig_len) == 0);

From d971b7834b746b6d2c8d08f513456027835b0b76 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Tue, 19 Mar 2024 15:08:53 +0100
Subject: [PATCH 03/14] test_suite_pk: fix RSA issue in pk_psa_sign() when
 !PK_[PARSE|WRITE]_C are defined

This bug was not found until now because:
- !PK_[WRITE|PARSE]_C is only tested in component_full_no_pkparse_pkwrite()
- the test only case concerning RSA key had MBEDTLS_PK_WRITE_C as dependency
  so it was not executed in that component.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.data     |  2 +-
 tests/suites/test_suite_pk.function | 37 ++++++++++++++++++++---------
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index 102aee2d5d..a979cbd7dd 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -684,7 +684,7 @@ depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_BP512R1
 pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):512
 
 PSA wrapped sign: RSA PKCS1 v1.5
-depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15:MBEDTLS_GENPRIME:MBEDTLS_PK_WRITE_C
+depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15:MBEDTLS_GENPRIME
 pk_psa_sign:PSA_KEY_TYPE_RSA_KEY_PAIR:1024
 
 PK sign ext: RSA2048, PK_RSA, MD_SHA256
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 8f5e7f6718..0a0f1585c2 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1875,12 +1875,21 @@ void pk_psa_sign(int psa_type, int bits)
             legacy_pub_key + sizeof(legacy_pub_key) - legacy_pub_key_len,
             legacy_pub_key_len);
 #else /* MBEDTLS_PK_WRITE_C */
-    ret = mbedtls_ecp_point_write_binary(&(mbedtls_pk_ec_ro(pk)->grp),
-                                         &(mbedtls_pk_ec_ro(pk)->Q),
-                                         MBEDTLS_ECP_PF_UNCOMPRESSED,
-                                         &legacy_pub_key_len, legacy_pub_key,
-                                         sizeof(legacy_pub_key));
-    TEST_EQUAL(ret, 0);
+    if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
+        TEST_EQUAL(mbedtls_ecp_point_write_binary(&(mbedtls_pk_ec_ro(pk)->grp),
+                                                  &(mbedtls_pk_ec_ro(pk)->Q),
+                                                  MBEDTLS_ECP_PF_UNCOMPRESSED,
+                                                  &legacy_pub_key_len, legacy_pub_key,
+                                                  sizeof(legacy_pub_key)), 0);
+    } else {
+        unsigned char *end = legacy_pub_key + sizeof(legacy_pub_key);
+        ret = mbedtls_rsa_write_pubkey(mbedtls_pk_rsa(pk), legacy_pub_key, &end);
+        legacy_pub_key_len = (size_t) ret;
+        TEST_ASSERT(legacy_pub_key_len > 0);
+        /* mbedtls_rsa_write_pubkey() writes data backward in the buffer so
+         * we shift that to the origin of the buffer instead. */
+        memmove(legacy_pub_key, end, legacy_pub_key_len);
+    }
 #endif /* MBEDTLS_PK_WRITE_C */
 
     /* Turn the PK context into an opaque one. */
@@ -1939,11 +1948,17 @@ void pk_psa_sign(int psa_type, int bits)
 #if defined(MBEDTLS_PK_WRITE_C) && defined(MBEDTLS_PK_PARSE_C)
     TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, legacy_pub_key, legacy_pub_key_len), 0);
 #else
-    TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
-    TEST_EQUAL(mbedtls_ecp_group_load(&(mbedtls_pk_ec_rw(pk)->grp), ecp_grp_id), 0);
-    TEST_EQUAL(mbedtls_ecp_point_read_binary(&(mbedtls_pk_ec_ro(pk)->grp),
-                                             &(mbedtls_pk_ec_rw(pk)->Q),
-                                             legacy_pub_key, legacy_pub_key_len), 0);
+    if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
+        TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
+        TEST_EQUAL(mbedtls_ecp_group_load(&(mbedtls_pk_ec_rw(pk)->grp), ecp_grp_id), 0);
+        TEST_EQUAL(mbedtls_ecp_point_read_binary(&(mbedtls_pk_ec_ro(pk)->grp),
+                                                 &(mbedtls_pk_ec_rw(pk)->Q),
+                                                 legacy_pub_key, legacy_pub_key_len), 0);
+    } else {
+        TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)), 0);
+        TEST_EQUAL(mbedtls_rsa_parse_pubkey(mbedtls_pk_rsa(pk), legacy_pub_key,
+                                            legacy_pub_key_len), 0);
+    }
 #endif
     TEST_ASSERT(mbedtls_pk_verify(&pk, MBEDTLS_MD_SHA256,
                                   hash, sizeof(hash), sig, sig_len) == 0);

From 4f3262de2d88358e7115bd40b83295895620364a Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Tue, 19 Mar 2024 15:43:25 +0100
Subject: [PATCH 04/14] pk_wrap: fix algorithm selection in
 rsa_opaque_sign_wrap()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 library/pk_wrap.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 256863a5a5..98b4f9a4e2 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -1508,10 +1508,11 @@ static int rsa_opaque_sign_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg
     }
 
     type = psa_get_key_type(&attributes);
+    alg = psa_get_key_algorithm(&attributes);
     psa_reset_key_attributes(&attributes);
 
     if (PSA_KEY_TYPE_IS_RSA(type)) {
-        alg = PSA_ALG_RSA_PKCS1V15_SIGN(mbedtls_md_psa_alg_from_type(md_alg));
+        alg = (alg & ~PSA_ALG_HASH_MASK) | mbedtls_md_psa_alg_from_type(md_alg);
     } else {
         return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
     }

From aa9cc4987985088e5f3c5ae959d73cf97a107cae Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Tue, 19 Mar 2024 15:43:52 +0100
Subject: [PATCH 05/14] test_suite_pk: test also RSA keys with PKCS1 v2.1
 padding mode in pk_psa_sign()

Previously only only PKCS1 v1.5 was tested.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.data     | 24 ++++++++++++++----------
 tests/suites/test_suite_pk.function |  9 ++++++++-
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index a979cbd7dd..b985ca050b 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -648,44 +648,48 @@ pk_sign_verify_restart:MBEDTLS_PK_ECKEY:MBEDTLS_ECP_DP_SECP256R1:"C9AFA9D845BA75
 
 PSA wrapped sign: SECP256R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP256R1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:0
 
 PSA wrapped sign: SECP384R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP384R1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):384
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):384:0
 
 PSA wrapped sign: SECP521R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP521R1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):521
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):521:0
 
 PSA wrapped sign: SECP192K1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP192K1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):192
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):192:0
 
 ## Currently buggy: https://github.com/ARMmbed/mbed-crypto/issues/336
 # PSA wrapped sign: SECP224K1
 # depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP224K1
-# pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):224
+# pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):224:0
 
 PSA wrapped sign: SECP256K1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP256K1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):256
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):256:0
 
 PSA wrapped sign: BP256R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_BP256R1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):256
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):256:0
 
 PSA wrapped sign: BP384R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_BP384R1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):384
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):384:0
 
 PSA wrapped sign: BP512R1
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_BP512R1
-pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):512
+pk_psa_sign:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):512:0
 
 PSA wrapped sign: RSA PKCS1 v1.5
 depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15:MBEDTLS_GENPRIME
-pk_psa_sign:PSA_KEY_TYPE_RSA_KEY_PAIR:1024
+pk_psa_sign:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:MBEDTLS_RSA_PKCS_V15
+
+PSA wrapped sign: RSA PKCS1 v2.1
+depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V21:MBEDTLS_GENPRIME
+pk_psa_sign:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:MBEDTLS_RSA_PKCS_V21
 
 PK sign ext: RSA2048, PK_RSA, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 0a0f1585c2..f232288661 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1817,7 +1817,7 @@ exit:
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_MD_CAN_SHA256:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_TEST_PK_PSA_SIGN */
-void pk_psa_sign(int psa_type, int bits)
+void pk_psa_sign(int psa_type, int bits, int rsa_padding)
 {
     mbedtls_pk_context pk;
     unsigned char hash[32];
@@ -1851,6 +1851,7 @@ void pk_psa_sign(int psa_type, int bits)
                                      mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
         TEST_ASSERT(mbedtls_rsa_gen_key(mbedtls_pk_rsa(pk), mbedtls_test_rnd_std_rand, NULL,
                                         bits, 3) == 0);
+        TEST_EQUAL(mbedtls_rsa_set_padding(mbedtls_pk_rsa(pk), rsa_padding, MBEDTLS_MD_NONE), 0);
     } else
 #endif /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
@@ -1861,6 +1862,7 @@ void pk_psa_sign(int psa_type, int bits)
     } else
 #endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
     {
+        (void) rsa_padding;
         TEST_ASSUME(!"Opaque PK key not supported in this configuration");
     }
 
@@ -1960,6 +1962,11 @@ void pk_psa_sign(int psa_type, int bits)
                                             legacy_pub_key_len), 0);
     }
 #endif
+#if defined(MBEDTLS_RSA_C)
+    if (PSA_KEY_TYPE_IS_RSA(psa_type)) {
+        TEST_EQUAL(mbedtls_rsa_set_padding(mbedtls_pk_rsa(pk), rsa_padding, MBEDTLS_MD_NONE), 0);
+    }
+#endif /* MBEDTLS_RSA_C */
     TEST_ASSERT(mbedtls_pk_verify(&pk, MBEDTLS_MD_SHA256,
                                   hash, sizeof(hash), sig, sig_len) == 0);
 

From f71c060cb202dca1ac2333d502ad0a793578c8d3 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Tue, 19 Mar 2024 19:35:37 +0100
Subject: [PATCH 06/14] test_suite_pk: properly size buffers for public keys in
 pk_psa_sign()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.function | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index f232288661..f4967644f3 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1822,8 +1822,13 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     mbedtls_pk_context pk;
     unsigned char hash[32];
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
-    unsigned char legacy_pub_key[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
-    unsigned char opaque_pub_key[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
+    /* These buffers will be used to contain the key in DER format. Unfortunately
+     * when only EC is supported on the PSA side (i.e. no RSA or DH)
+     * PSA_EXPORT_PUBLIC_KEY_MAX_SIZE falls to PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(),
+     * but this is not enough to contain the DER representation of the public key.
+     * Therefore we pick the RSA size to be safe. */
+    unsigned char legacy_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
+    unsigned char opaque_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
     size_t sig_len, legacy_pub_key_len, opaque_pub_key_len;
     int ret;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;

From 480dfc7ad77845f310e6e143a6988293ffeffe98 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Wed, 20 Mar 2024 09:35:54 +0100
Subject: [PATCH 07/14] test_suite_pk: fix guards in pk_psa_sign()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.function | 35 +++++++++++++++++++----------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index f4967644f3..bbdb42d2ac 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1830,9 +1830,11 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     unsigned char legacy_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
     unsigned char opaque_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
     size_t sig_len, legacy_pub_key_len, opaque_pub_key_len;
-    int ret;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+#if defined(MBEDTLS_RSA_C) || defined(MBEDTLS_PK_WRITE_C)
+    int ret;
+#endif /* MBEDTLS_RSA_C || MBEDTLS_PK_WRITE_C */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     mbedtls_ecp_group_id ecp_grp_id;
 #endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
@@ -1857,19 +1859,17 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
         TEST_ASSERT(mbedtls_rsa_gen_key(mbedtls_pk_rsa(pk), mbedtls_test_rnd_std_rand, NULL,
                                         bits, 3) == 0);
         TEST_EQUAL(mbedtls_rsa_set_padding(mbedtls_pk_rsa(pk), rsa_padding, MBEDTLS_MD_NONE), 0);
-    } else
+    }
+#else /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
+    (void) rsa_padding;
 #endif /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
         ecp_grp_id = mbedtls_ecc_group_from_psa(psa_type, bits);
         TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)) == 0);
         TEST_ASSERT(pk_genkey(&pk, ecp_grp_id) == 0);
-    } else
-#endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
-    {
-        (void) rsa_padding;
-        TEST_ASSUME(!"Opaque PK key not supported in this configuration");
     }
+#endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
 
     /* Export public key from the non-opaque PK context we just created. */
 #if defined(MBEDTLS_PK_WRITE_C)
@@ -1882,13 +1882,17 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
             legacy_pub_key + sizeof(legacy_pub_key) - legacy_pub_key_len,
             legacy_pub_key_len);
 #else /* MBEDTLS_PK_WRITE_C */
+#if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
         TEST_EQUAL(mbedtls_ecp_point_write_binary(&(mbedtls_pk_ec_ro(pk)->grp),
                                                   &(mbedtls_pk_ec_ro(pk)->Q),
                                                   MBEDTLS_ECP_PF_UNCOMPRESSED,
                                                   &legacy_pub_key_len, legacy_pub_key,
                                                   sizeof(legacy_pub_key)), 0);
-    } else {
+    }
+#endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
+#if defined(MBEDTLS_RSA_C)
+    if (PSA_KEY_TYPE_IS_RSA(psa_type)) {
         unsigned char *end = legacy_pub_key + sizeof(legacy_pub_key);
         ret = mbedtls_rsa_write_pubkey(mbedtls_pk_rsa(pk), legacy_pub_key, &end);
         legacy_pub_key_len = (size_t) ret;
@@ -1897,6 +1901,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
          * we shift that to the origin of the buffer instead. */
         memmove(legacy_pub_key, end, legacy_pub_key_len);
     }
+#endif /* MBEDTLS_RSA_C */
 #endif /* MBEDTLS_PK_WRITE_C */
 
     /* Turn the PK context into an opaque one. */
@@ -1932,7 +1937,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     memmove(opaque_pub_key,
             opaque_pub_key + sizeof(opaque_pub_key) - opaque_pub_key_len,
             opaque_pub_key_len);
-#else
+#else /* MBEDTLS_PK_WRITE_C */
     TEST_EQUAL(psa_export_public_key(key_id, opaque_pub_key, sizeof(opaque_pub_key),
                                      &opaque_pub_key_len), PSA_SUCCESS);
 #endif /* MBEDTLS_PK_WRITE_C */
@@ -1954,19 +1959,25 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     mbedtls_pk_init(&pk);
 #if defined(MBEDTLS_PK_WRITE_C) && defined(MBEDTLS_PK_PARSE_C)
     TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, legacy_pub_key, legacy_pub_key_len), 0);
-#else
+#else /* MBEDTLS_PK_WRITE_C && MBEDTLS_PK_PARSE_C */
+#if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
         TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
         TEST_EQUAL(mbedtls_ecp_group_load(&(mbedtls_pk_ec_rw(pk)->grp), ecp_grp_id), 0);
         TEST_EQUAL(mbedtls_ecp_point_read_binary(&(mbedtls_pk_ec_ro(pk)->grp),
                                                  &(mbedtls_pk_ec_rw(pk)->Q),
                                                  legacy_pub_key, legacy_pub_key_len), 0);
-    } else {
+    }
+#endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
+#if defined(MBEDTLS_RSA_C)
+    if (PSA_KEY_TYPE_IS_RSA(psa_type)) {
         TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)), 0);
         TEST_EQUAL(mbedtls_rsa_parse_pubkey(mbedtls_pk_rsa(pk), legacy_pub_key,
                                             legacy_pub_key_len), 0);
     }
-#endif
+#endif /* MBEDTLS_RSA_C */
+#endif /* MBEDTLS_PK_WRITE_C && MBEDTLS_PK_PARSE_C */
+
 #if defined(MBEDTLS_RSA_C)
     if (PSA_KEY_TYPE_IS_RSA(psa_type)) {
         TEST_EQUAL(mbedtls_rsa_set_padding(mbedtls_pk_rsa(pk), rsa_padding, MBEDTLS_MD_NONE), 0);

From d45836a1c325a724809c66e26bde56f84e81dc9d Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Wed, 20 Mar 2024 15:42:55 +0100
Subject: [PATCH 08/14] pk_wrap: fix algorithm selection in
 rsa_opaque_decrypt()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 library/pk_wrap.c | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 98b4f9a4e2..19196b559a 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -1469,16 +1469,29 @@ static int rsa_opaque_decrypt(mbedtls_pk_context *pk,
                               unsigned char *output, size_t *olen, size_t osize,
                               int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_algorithm_t alg;
+    psa_key_type_t type;
     psa_status_t status;
 
     /* PSA has its own RNG */
     (void) f_rng;
     (void) p_rng;
 
-    status = psa_asymmetric_decrypt(pk->priv_id, PSA_ALG_RSA_PKCS1V15_CRYPT,
-                                    input, ilen,
-                                    NULL, 0,
-                                    output, osize, olen);
+    status = psa_get_key_attributes(pk->priv_id, &attributes);
+    if (status != PSA_SUCCESS) {
+        return PSA_PK_TO_MBEDTLS_ERR(status);
+    }
+
+    type = psa_get_key_type(&attributes);
+    alg = psa_get_key_algorithm(&attributes);
+    psa_reset_key_attributes(&attributes);
+
+    if (!PSA_KEY_TYPE_IS_RSA(type)) {
+        return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+    }
+
+    status = psa_asymmetric_decrypt(pk->priv_id, alg, input, ilen, NULL, 0, output, osize, olen);
     if (status != PSA_SUCCESS) {
         return PSA_PK_RSA_TO_MBEDTLS_ERR(status);
     }
@@ -1517,9 +1530,7 @@ static int rsa_opaque_sign_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg
         return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
     }
 
-    /* make the signature */
-    status = psa_sign_hash(pk->priv_id, alg, hash, hash_len,
-                           sig, sig_size, sig_len);
+    status = psa_sign_hash(pk->priv_id, alg, hash, hash_len, sig, sig_size, sig_len);
     if (status != PSA_SUCCESS) {
         if (PSA_KEY_TYPE_IS_RSA(type)) {
             return PSA_PK_RSA_TO_MBEDTLS_ERR(status);

From 1b533ab205f4674dba31b563a357db464f96a39f Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Wed, 20 Mar 2024 15:43:33 +0100
Subject: [PATCH 09/14] test_suite_pk: test also RSA OAEP in
 pk_wrap_rsa_decrypt_test_vec()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.data     | 12 ++++++++----
 tests/suites/test_suite_pk.function |  6 ++++++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index b985ca050b..a929c82f4f 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -482,13 +482,17 @@ RSA decrypt test vector - PKCS1v2.1, but data is PKCS1v1.5 encrypted
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA1
 pk_rsa_decrypt_test_vec:"28818cb14236ad18f4527e7f1f7633e96cef021bc3234475d7f61e88702b6335b42a352ed3f3267ac7c3e9ba4af17e45096c63eefd8d9a7cb42dfc52fffb2f5b8afb305b46312c2eb50634123b4437a2287ac57b7509d59a583fb741989a49f32625e9267b4641a6607b7303d35c68489db53c8d387b620d0d46a852e72ea43c":1024:MBEDTLS_RSA_PKCS_V21:MBEDTLS_MD_SHA1:"eecfae81b1b9b3c908810b10a1b5600199eb9f44aef4fda493b81a9e3d84f632124ef0236e5d1e3b7e28fae7aa040a2d5b252176459d1f397541ba2a58fb6599":"c97fb1f027f453f6341233eaaad1d9353f6c42d08866b1d05a0f2035028b9d869840b41666b42e92ea0da3b43204b5cfce3352524d0416a5a441e700af461503":"bbf82f090682ce9c2338ac2b9da871f7368d07eed41043a440d6b6f07454f51fb8dfbaaf035c02ab61ea48ceeb6fcd4876ed520d60e1ec4619719d8a5b8b807fafb8e0a3dfc737723ee6b4b7d93a2584ee6a649d060953748834b2454598394ee0aab12d7b61a51f527a9a41f6c1687fe2537298ca2a8f5946f8e5fd091dbdcb":"11":"d436e99569fd32a7c8a05bbc90d32c49":MBEDTLS_ERR_RSA_INVALID_PADDING
 
-RSA Opaque decrypt test vector #1
+RSA Opaque PCKS1 v1.5 - decrypt test vector #1
 depends_on:MBEDTLS_PKCS1_V15
-pk_wrap_rsa_decrypt_test_vec:"a42eda41e56235e666e7faaa77100197f657288a1bf183e4820f0c37ce2c456b960278d6003e0bbcd4be4a969f8e8fd9231e1f492414f00ed09844994c86ec32db7cde3bec7f0c3dbf6ae55baeb2712fa609f5fc3207a824eb3dace31849cd6a6084318523912bccb84cf42e3c6d6d1685131d69bb545acec827d2b0dfdd5568b7dcc4f5a11d6916583fefa689d367f8c9e1d95dcd2240895a9470b0c1730f97cd6e8546860bd254801769f54be96e16362ddcbf34d56035028890199e0f48db38642cb66a4181e028a6443a404fea284ce02b4614b683367d40874e505611d23142d49f06feea831d52d347b13610b413c4efc43a6de9f0b08d2a951dc503b6":2048:"e79a373182bfaa722eb035f772ad2a9464bd842de59432c18bbab3a7dfeae318c9b915ee487861ab665a40bd6cda560152578e8579016c929df99fea05b4d64efca1d543850bc8164b40d71ed7f3fa4105df0fb9b9ad2a18ce182c8a4f4f975bea9aa0b9a1438a27a28e97ac8330ef37383414d1bd64607d6979ac050424fd17":"c6749cbb0db8c5a177672d4728a8b22392b2fc4d3b8361d5c0d5055a1b4e46d821f757c24eef2a51c561941b93b3ace7340074c058c9bb48e7e7414f42c41da4cccb5c2ba91deb30c586b7fb18af12a52995592ad139d3be429add6547e044becedaf31fa3b39421e24ee034fbf367d11f6b8f88ee483d163b431e1654ad3e89":"b38ac65c8141f7f5c96e14470e851936a67bf94cc6821a39ac12c05f7c0b06d9e6ddba2224703b02e25f31452f9c4a8417b62675fdc6df46b94813bc7b9769a892c482b830bfe0ad42e46668ace68903617faf6681f4babf1cc8e4b0420d3c7f61dc45434c6b54e2c3ee0fc07908509d79c9826e673bf8363255adb0add2401039a7bcd1b4ecf0fbe6ec8369d2da486eec59559dd1d54c9b24190965eafbdab203b35255765261cd0909acf93c3b8b8428cbb448de4715d1b813d0c94829c229543d391ce0adab5351f97a3810c1f73d7b1458b97daed4209c50e16d064d2d5bfda8c23893d755222793146d0a78c3d64f35549141486c3b0961a7b4c1a2034f":"3":"4E636AF98E40F3ADCFCCB698F4E80B9F":0
+pk_wrap_rsa_decrypt_test_vec:"a42eda41e56235e666e7faaa77100197f657288a1bf183e4820f0c37ce2c456b960278d6003e0bbcd4be4a969f8e8fd9231e1f492414f00ed09844994c86ec32db7cde3bec7f0c3dbf6ae55baeb2712fa609f5fc3207a824eb3dace31849cd6a6084318523912bccb84cf42e3c6d6d1685131d69bb545acec827d2b0dfdd5568b7dcc4f5a11d6916583fefa689d367f8c9e1d95dcd2240895a9470b0c1730f97cd6e8546860bd254801769f54be96e16362ddcbf34d56035028890199e0f48db38642cb66a4181e028a6443a404fea284ce02b4614b683367d40874e505611d23142d49f06feea831d52d347b13610b413c4efc43a6de9f0b08d2a951dc503b6":2048:"e79a373182bfaa722eb035f772ad2a9464bd842de59432c18bbab3a7dfeae318c9b915ee487861ab665a40bd6cda560152578e8579016c929df99fea05b4d64efca1d543850bc8164b40d71ed7f3fa4105df0fb9b9ad2a18ce182c8a4f4f975bea9aa0b9a1438a27a28e97ac8330ef37383414d1bd64607d6979ac050424fd17":"c6749cbb0db8c5a177672d4728a8b22392b2fc4d3b8361d5c0d5055a1b4e46d821f757c24eef2a51c561941b93b3ace7340074c058c9bb48e7e7414f42c41da4cccb5c2ba91deb30c586b7fb18af12a52995592ad139d3be429add6547e044becedaf31fa3b39421e24ee034fbf367d11f6b8f88ee483d163b431e1654ad3e89":"b38ac65c8141f7f5c96e14470e851936a67bf94cc6821a39ac12c05f7c0b06d9e6ddba2224703b02e25f31452f9c4a8417b62675fdc6df46b94813bc7b9769a892c482b830bfe0ad42e46668ace68903617faf6681f4babf1cc8e4b0420d3c7f61dc45434c6b54e2c3ee0fc07908509d79c9826e673bf8363255adb0add2401039a7bcd1b4ecf0fbe6ec8369d2da486eec59559dd1d54c9b24190965eafbdab203b35255765261cd0909acf93c3b8b8428cbb448de4715d1b813d0c94829c229543d391ce0adab5351f97a3810c1f73d7b1458b97daed4209c50e16d064d2d5bfda8c23893d755222793146d0a78c3d64f35549141486c3b0961a7b4c1a2034f":"3":MBEDTLS_RSA_PKCS_V15:"4E636AF98E40F3ADCFCCB698F4E80B9F":0
 
-RSA Opaque decrypt test vector #2
+RSA Opaque PCKS1 v2.1 - decrypt test vector #1
+depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA1
+pk_wrap_rsa_decrypt_test_vec:"1253e04dc0a5397bb44a7ab87e9bf2a039a33d1e996fc82a94ccd30074c95df763722017069e5268da5d1c0b4f872cf653c11df82314a67968dfeae28def04bb6d84b1c31d654a1970e5783bd6eb96a024c2ca2f4a90fe9f2ef5c9c140e5bb48da9536ad8700c84fc9130adea74e558d51a74ddf85d8b50de96838d6063e0955":1024:"eecfae81b1b9b3c908810b10a1b5600199eb9f44aef4fda493b81a9e3d84f632124ef0236e5d1e3b7e28fae7aa040a2d5b252176459d1f397541ba2a58fb6599":"c97fb1f027f453f6341233eaaad1d9353f6c42d08866b1d05a0f2035028b9d869840b41666b42e92ea0da3b43204b5cfce3352524d0416a5a441e700af461503":"bbf82f090682ce9c2338ac2b9da871f7368d07eed41043a440d6b6f07454f51fb8dfbaaf035c02ab61ea48ceeb6fcd4876ed520d60e1ec4619719d8a5b8b807fafb8e0a3dfc737723ee6b4b7d93a2584ee6a649d060953748834b2454598394ee0aab12d7b61a51f527a9a41f6c1687fe2537298ca2a8f5946f8e5fd091dbdcb":"11":MBEDTLS_RSA_PKCS_V21:"d436e99569fd32a7c8a05bbc90d32c49":0
+
+RSA Opaque PCKS1 v1.5 - decrypt test vector #2
 depends_on:MBEDTLS_PKCS1_V15
-pk_wrap_rsa_decrypt_test_vec:"a42eda41e56235e666e7faaa77100197f657288a1bf183e4820f0c37ce2c456b960278d6003e0bbcd4be4a969f8e8fd9231e1f492414f00ed09844994c86ec32db7cde3bec7f0c3dbf6ae55baeb2712fa609f5fc3207a824eb3dace31849cd6a6084318523912bccb84cf42e3c6d6d1685131d69bb545acec827d2b0dfdd5568b7dcc4f5a11d6916583fefa689d367f8c9e1d95dcd2240895a9470b0c1730f97cd6e8546860bd254801769f54be96e16362ddcbf34d56035028890199e0f48db38642cb66a4181e028a6443a404feb284ce02b4614b683367d40874e505611d23142d49f06feea831d52d347b13610b413c4efc43a6de9f0b08d2a951dc503b6":2048:"e79a373182bfaa722eb035f772ad2a9464bd842de59432c18bbab3a7dfeae318c9b915ee487861ab665a40bd6cda560152578e8579016c929df99fea05b4d64efca1d543850bc8164b40d71ed7f3fa4105df0fb9b9ad2a18ce182c8a4f4f975bea9aa0b9a1438a27a28e97ac8330ef37383414d1bd64607d6979ac050424fd17":"c6749cbb0db8c5a177672d4728a8b22392b2fc4d3b8361d5c0d5055a1b4e46d821f757c24eef2a51c561941b93b3ace7340074c058c9bb48e7e7414f42c41da4cccb5c2ba91deb30c586b7fb18af12a52995592ad139d3be429add6547e044becedaf31fa3b39421e24ee034fbf367d11f6b8f88ee483d163b431e1654ad3e89":"b38ac65c8141f7f5c96e14470e851936a67bf94cc6821a39ac12c05f7c0b06d9e6ddba2224703b02e25f31452f9c4a8417b62675fdc6df46b94813bc7b9769a892c482b830bfe0ad42e46668ace68903617faf6681f4babf1cc8e4b0420d3c7f61dc45434c6b54e2c3ee0fc07908509d79c9826e673bf8363255adb0add2401039a7bcd1b4ecf0fbe6ec8369d2da486eec59559dd1d54c9b24190965eafbdab203b35255765261cd0909acf93c3b8b8428cbb448de4715d1b813d0c94829c229543d391ce0adab5351f97a3810c1f73d7b1458b97daed4209c50e16d064d2d5bfda8c23893d755222793146d0a78c3d64f35549141486c3b0961a7b4c1a2034f":"3":"4E636AF98E40F3ADCFCCB698F4E80B9F":MBEDTLS_ERR_RSA_INVALID_PADDING
+pk_wrap_rsa_decrypt_test_vec:"a42eda41e56235e666e7faaa77100197f657288a1bf183e4820f0c37ce2c456b960278d6003e0bbcd4be4a969f8e8fd9231e1f492414f00ed09844994c86ec32db7cde3bec7f0c3dbf6ae55baeb2712fa609f5fc3207a824eb3dace31849cd6a6084318523912bccb84cf42e3c6d6d1685131d69bb545acec827d2b0dfdd5568b7dcc4f5a11d6916583fefa689d367f8c9e1d95dcd2240895a9470b0c1730f97cd6e8546860bd254801769f54be96e16362ddcbf34d56035028890199e0f48db38642cb66a4181e028a6443a404feb284ce02b4614b683367d40874e505611d23142d49f06feea831d52d347b13610b413c4efc43a6de9f0b08d2a951dc503b6":2048:"e79a373182bfaa722eb035f772ad2a9464bd842de59432c18bbab3a7dfeae318c9b915ee487861ab665a40bd6cda560152578e8579016c929df99fea05b4d64efca1d543850bc8164b40d71ed7f3fa4105df0fb9b9ad2a18ce182c8a4f4f975bea9aa0b9a1438a27a28e97ac8330ef37383414d1bd64607d6979ac050424fd17":"c6749cbb0db8c5a177672d4728a8b22392b2fc4d3b8361d5c0d5055a1b4e46d821f757c24eef2a51c561941b93b3ace7340074c058c9bb48e7e7414f42c41da4cccb5c2ba91deb30c586b7fb18af12a52995592ad139d3be429add6547e044becedaf31fa3b39421e24ee034fbf367d11f6b8f88ee483d163b431e1654ad3e89":"b38ac65c8141f7f5c96e14470e851936a67bf94cc6821a39ac12c05f7c0b06d9e6ddba2224703b02e25f31452f9c4a8417b62675fdc6df46b94813bc7b9769a892c482b830bfe0ad42e46668ace68903617faf6681f4babf1cc8e4b0420d3c7f61dc45434c6b54e2c3ee0fc07908509d79c9826e673bf8363255adb0add2401039a7bcd1b4ecf0fbe6ec8369d2da486eec59559dd1d54c9b24190965eafbdab203b35255765261cd0909acf93c3b8b8428cbb448de4715d1b813d0c94829c229543d391ce0adab5351f97a3810c1f73d7b1458b97daed4209c50e16d064d2d5bfda8c23893d755222793146d0a78c3d64f35549141486c3b0961a7b4c1a2034f":"3":MBEDTLS_RSA_PKCS_V15:"4E636AF98E40F3ADCFCCB698F4E80B9F":MBEDTLS_ERR_RSA_INVALID_PADDING
 
 EC nocrypt
 depends_on:MBEDTLS_PK_HAVE_ECC_KEYS
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index bbdb42d2ac..979eaf5cdc 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1576,6 +1576,7 @@ exit:
 void pk_wrap_rsa_decrypt_test_vec(data_t *cipher, int mod,
                                   char *input_P, char *input_Q,
                                   char *input_N, char *input_E,
+                                  int padding_mode,
                                   data_t *clear, int ret)
 {
     unsigned char output[256];
@@ -1610,6 +1611,11 @@ void pk_wrap_rsa_decrypt_test_vec(data_t *cipher, int mod,
     TEST_EQUAL(mbedtls_rsa_get_len(rsa), (mod + 7) / 8);
     TEST_EQUAL(mbedtls_rsa_complete(rsa), 0);
 
+    /* Set padding mode */
+    if (padding_mode == MBEDTLS_RSA_PKCS_V21) {
+        TEST_EQUAL(mbedtls_rsa_set_padding(rsa, padding_mode, MBEDTLS_MD_SHA1), 0);
+    }
+
     /* Turn PK context into an opaque one. */
     TEST_EQUAL(mbedtls_pk_get_psa_attributes(&pk, PSA_KEY_USAGE_DECRYPT, &key_attr), 0);
     TEST_EQUAL(mbedtls_pk_import_into_psa(&pk, &key_attr, &key_id), 0);

From 6fb2586dfd31b87befeafefda7de0ea6c364ed30 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Wed, 20 Mar 2024 16:55:14 +0100
Subject: [PATCH 10/14] test_suite_pk: fix guards in pk_psa_sign()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.function | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 979eaf5cdc..702c23e5ca 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1956,16 +1956,11 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     mbedtls_pk_free(&pk);
     TEST_ASSERT(PSA_SUCCESS == psa_destroy_key(key_id));
 
-    /* Create a new non-opaque PK context to verify the signature.
-     *
-     * Note: if we used "pk_write" previously, then we go for a "pk_parse" here;
-     * otherwise if we went for "ecp_point_write_binary" then we'll go
-     * for a "ecp_point_read_binary" here. This allows to drop dependencies
-     * on "PK_WRITE" and "PK_PARSE" if required */
+    /* Create a new non-opaque PK context to verify the signature. */
     mbedtls_pk_init(&pk);
-#if defined(MBEDTLS_PK_WRITE_C) && defined(MBEDTLS_PK_PARSE_C)
+#if defined(MBEDTLS_PK_PARSE_C)
     TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, legacy_pub_key, legacy_pub_key_len), 0);
-#else /* MBEDTLS_PK_WRITE_C && MBEDTLS_PK_PARSE_C */
+#else /* MBEDTLS_PK_PARSE_C */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
         TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
@@ -1982,7 +1977,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
                                             legacy_pub_key_len), 0);
     }
 #endif /* MBEDTLS_RSA_C */
-#endif /* MBEDTLS_PK_WRITE_C && MBEDTLS_PK_PARSE_C */
+#endif /* MBEDTLS_PK_PARSE_C */
 
 #if defined(MBEDTLS_RSA_C)
     if (PSA_KEY_TYPE_IS_RSA(psa_type)) {

From 027796c0cc2619a770d3024422fd53b5a9e271c0 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Wed, 20 Mar 2024 16:55:35 +0100
Subject: [PATCH 11/14] test_suite_pk: uniformly generate RSA and EC keys in
 pk_psa_sign()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.function | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 702c23e5ca..272e9f67a9 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1862,8 +1862,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     if (PSA_KEY_TYPE_IS_RSA(psa_type)) {
         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), mbedtls_test_rnd_std_rand, NULL,
-                                        bits, 3) == 0);
+        TEST_EQUAL(pk_genkey(&pk, bits), 0);
         TEST_EQUAL(mbedtls_rsa_set_padding(mbedtls_pk_rsa(pk), rsa_padding, MBEDTLS_MD_NONE), 0);
     }
 #else /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */

From 144c27b0f3b18cd10b9ee7fa4e573f8474c8cc89 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Wed, 20 Mar 2024 17:10:35 +0100
Subject: [PATCH 12/14] pkwrite: add new internal symbol for the max supported
 public key DER length

This is also used in pk_psa_sign() to properly size buffers holding
the public key.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 library/pkwrite.h                   |  9 +++++++++
 tests/suites/test_suite_pk.function | 12 +++++-------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/library/pkwrite.h b/library/pkwrite.h
index 544ab2f329..01dc3d2f0f 100644
--- a/library/pkwrite.h
+++ b/library/pkwrite.h
@@ -109,4 +109,13 @@
 #define MBEDTLS_PK_ECP_PRV_DER_MAX_BYTES   0
 
 #endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+
+/* Define the maximum available public key DER length based on the supported
+ * key types (EC and/or RSA). */
+#if (MBEDTLS_PK_RSA_PUB_DER_MAX_BYTES > MBEDTLS_PK_ECP_PUB_DER_MAX_BYTES)
+#define MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE    MBEDTLS_PK_RSA_PUB_DER_MAX_BYTES
+#else
+#define MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE    MBEDTLS_PK_ECP_PUB_DER_MAX_BYTES
+#endif
+
 #endif /* MBEDTLS_PK_WRITE_H */
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 272e9f67a9..d116050bf6 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -23,6 +23,9 @@
 
 #include <test/psa_exercise_key.h>
 
+/* Needed for the definition of MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE. */
+#include "pkwrite.h"
+
 /* Used for properly sizing the key buffer in pk_genkey_ec() */
 #include "psa_util_internal.h"
 
@@ -1828,13 +1831,8 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     mbedtls_pk_context pk;
     unsigned char hash[32];
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
-    /* These buffers will be used to contain the key in DER format. Unfortunately
-     * when only EC is supported on the PSA side (i.e. no RSA or DH)
-     * PSA_EXPORT_PUBLIC_KEY_MAX_SIZE falls to PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(),
-     * but this is not enough to contain the DER representation of the public key.
-     * Therefore we pick the RSA size to be safe. */
-    unsigned char legacy_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
-    unsigned char opaque_pub_key[RSA_WRITE_PUBKEY_MAX_SIZE];
+    unsigned char legacy_pub_key[MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE];
+    unsigned char opaque_pub_key[MBEDTLS_PK_WRITE_PUBKEY_MAX_SIZE];
     size_t sig_len, legacy_pub_key_len, opaque_pub_key_len;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

From ea01efa589cde8e984f37a9aa1898825532b3ec6 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Wed, 20 Mar 2024 17:19:08 +0100
Subject: [PATCH 13/14] add changelog

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 ChangeLog.d/8938.txt | 3 +++
 1 file changed, 3 insertions(+)
 create mode 100644 ChangeLog.d/8938.txt

diff --git a/ChangeLog.d/8938.txt b/ChangeLog.d/8938.txt
new file mode 100644
index 0000000000..68a1c084be
--- /dev/null
+++ b/ChangeLog.d/8938.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Fix RSA opaque keys always using PKCS1 v1.5 algorithms instead of the
+     primary algorithm of the wrapped PSA key.

From 2833050bb660765d1685f803de77f58b9cf9c8d2 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Thu, 21 Mar 2024 05:24:12 +0100
Subject: [PATCH 14/14] test_suite_pk: fix guards in pk_psa_sign()

If the public key is exported with mbedtls_pk_write_pubkey_der()
it should be re-imported with mbedtls_pk_parse_public_key().
Alternative options (when PK_WRITE is not defined), i.e.
mbedtls_ecp_point_write_binary() and mbedtls_rsa_write_pubkey(),
export the key in a different format which cannot be parsed by
pk_parse module so mbedtls_ecp_point_read_binary() and
mbedtls_rsa_parse_pubkey() should be used respectively in this
case.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.function | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index d116050bf6..388879d1a1 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1875,7 +1875,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
 #endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
 
     /* Export public key from the non-opaque PK context we just created. */
-#if defined(MBEDTLS_PK_WRITE_C)
+#if defined(MBEDTLS_PK_PARSE_C) && defined(MBEDTLS_PK_WRITE_C)
     ret = mbedtls_pk_write_pubkey_der(&pk, legacy_pub_key, sizeof(legacy_pub_key));
     TEST_ASSERT(ret >= 0);
     legacy_pub_key_len = (size_t) ret;
@@ -1884,7 +1884,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     memmove(legacy_pub_key,
             legacy_pub_key + sizeof(legacy_pub_key) - legacy_pub_key_len,
             legacy_pub_key_len);
-#else /* MBEDTLS_PK_WRITE_C */
+#else /* MBEDTLS_PK_PARSE_C && MBEDTLS_PK_WRITE_C */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
         TEST_EQUAL(mbedtls_ecp_point_write_binary(&(mbedtls_pk_ec_ro(pk)->grp),
@@ -1905,7 +1905,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
         memmove(legacy_pub_key, end, legacy_pub_key_len);
     }
 #endif /* MBEDTLS_RSA_C */
-#endif /* MBEDTLS_PK_WRITE_C */
+#endif /* MBEDTLS_PK_PARSE_C && MBEDTLS_PK_WRITE_C */
 
     /* Turn the PK context into an opaque one. */
     TEST_EQUAL(mbedtls_pk_get_psa_attributes(&pk, PSA_KEY_USAGE_SIGN_HASH, &attributes), 0);
@@ -1932,7 +1932,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     }
 
     /* Export public key from the opaque PK context. */
-#if defined(MBEDTLS_PK_WRITE_C)
+#if defined(MBEDTLS_PK_PARSE_C) && defined(MBEDTLS_PK_WRITE_C)
     ret = mbedtls_pk_write_pubkey_der(&pk, opaque_pub_key, sizeof(opaque_pub_key));
     TEST_ASSERT(ret >= 0);
     opaque_pub_key_len = (size_t) ret;
@@ -1940,10 +1940,10 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
     memmove(opaque_pub_key,
             opaque_pub_key + sizeof(opaque_pub_key) - opaque_pub_key_len,
             opaque_pub_key_len);
-#else /* MBEDTLS_PK_WRITE_C */
+#else /* MBEDTLS_PK_PARSE_C && MBEDTLS_PK_WRITE_C */
     TEST_EQUAL(psa_export_public_key(key_id, opaque_pub_key, sizeof(opaque_pub_key),
                                      &opaque_pub_key_len), PSA_SUCCESS);
-#endif /* MBEDTLS_PK_WRITE_C */
+#endif /* MBEDTLS_PK_PARSE_C && MBEDTLS_PK_WRITE_C */
 
     /* Check that the public keys of opaque and non-opaque PK contexts match. */
     TEST_EQUAL(opaque_pub_key_len, legacy_pub_key_len);
@@ -1955,9 +1955,9 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
 
     /* Create a new non-opaque PK context to verify the signature. */
     mbedtls_pk_init(&pk);
-#if defined(MBEDTLS_PK_PARSE_C)
+#if defined(MBEDTLS_PK_PARSE_C) && defined(MBEDTLS_PK_WRITE_C)
     TEST_EQUAL(mbedtls_pk_parse_public_key(&pk, legacy_pub_key, legacy_pub_key_len), 0);
-#else /* MBEDTLS_PK_PARSE_C */
+#else /* MBEDTLS_PK_PARSE_C && MBEDTLS_PK_WRITE_C */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(psa_type)) {
         TEST_EQUAL(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)), 0);
@@ -1974,7 +1974,7 @@ void pk_psa_sign(int psa_type, int bits, int rsa_padding)
                                             legacy_pub_key_len), 0);
     }
 #endif /* MBEDTLS_RSA_C */
-#endif /* MBEDTLS_PK_PARSE_C */
+#endif /* MBEDTLS_PK_PARSE_C && MBEDTLS_PK_WRITE_C */
 
 #if defined(MBEDTLS_RSA_C)
     if (PSA_KEY_TYPE_IS_RSA(psa_type)) {