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>
diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index a0dacf0..102aee2 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -648,44 +648,44 @@
 
 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 afc1e34..e1a8e1c 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1817,7 +1817,7 @@
 /* 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 @@
     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 @@
     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 @@
 
     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 @@
                                       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 @@
     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 @@
                                   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);