test_suite_pk: various minor fixes

- removed redundant info from data file (i.e. informations that
  can be extrapolated somehow)
- removed unecessary parameters in functions
- added some extra check on the generated PK contexts
- etc...

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index e08751e..69546b2 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -428,49 +428,67 @@
 #endif
 
 #if defined(MBEDTLS_PSA_CRYPTO_C)
-mbedtls_svc_key_id_t pk_psa_pub_key_from_priv(mbedtls_svc_key_id_t priv_id,
-                                              psa_key_type_t type, psa_key_usage_t usage,
-                                              psa_algorithm_t alg, size_t bits)
+static mbedtls_svc_key_id_t psa_pub_key_from_priv(mbedtls_svc_key_id_t priv_id)
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_type_t type;
+    psa_algorithm_t alg;
+    psa_key_usage_t usage;
     unsigned char pub_key_buf[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
     size_t pub_key_len;
     mbedtls_svc_key_id_t pub_key = MBEDTLS_SVC_KEY_ID_INIT;
 
+    /* Get attributes from the private key. */
+    PSA_ASSERT(psa_get_key_attributes(priv_id, &attributes));
+    type = psa_get_key_type(&attributes);
+    usage = psa_get_key_usage_flags(&attributes);
+    alg = psa_get_key_algorithm(&attributes);
+    psa_reset_key_attributes(&attributes);
+
+    /* Export the public key and then import it in a new slot. */
     PSA_ASSERT(psa_export_public_key(priv_id, pub_key_buf, sizeof(pub_key_buf), &pub_key_len));
 
+    /* Notes:
+     * - psa_import_key() automatically determines the key's bit length
+     *   from the provided key data. That's why psa_set_key_bits() is not used
+     *   below.
+     * - public keys are always exportable by default even if PSA_KEY_USAGE_EXPORT
+     *   is not set.
+     */
+    type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(type);
+    usage &= ~PSA_KEY_USAGE_EXPORT;
+    psa_set_key_type(&attributes, type);
     psa_set_key_usage_flags(&attributes, usage);
     psa_set_key_algorithm(&attributes, alg);
-    psa_set_key_type(&attributes, type);
-    psa_set_key_bits(&attributes, bits);
 
     PSA_ASSERT(psa_import_key(&attributes, pub_key_buf, pub_key_len, &pub_key));
 
 exit:
+    psa_reset_key_attributes(&attributes);
     return pub_key;
 }
 
 psa_status_t pk_psa_import_key(unsigned char *key_data, size_t key_len,
                                psa_key_type_t type, psa_key_usage_t usage,
-                               psa_algorithm_t alg, size_t bits,
-                               mbedtls_svc_key_id_t *key)
+                               psa_algorithm_t alg, mbedtls_svc_key_id_t *key)
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status;
 
     *key = MBEDTLS_SVC_KEY_ID_INIT;
 
+    /* Note: psa_import_key() automatically determines the key's bit length
+     * from the provided key data. That's why psa_set_key_bits() is not used below. */
     psa_set_key_usage_flags(&attributes, usage);
     psa_set_key_algorithm(&attributes, alg);
     psa_set_key_type(&attributes, type);
-    psa_set_key_bits(&attributes, bits);
     status = psa_import_key(&attributes, key_data, key_len, key);
 
     return status;
 }
 
-psa_status_t pk_psa_genkey_generic(psa_key_type_t type, psa_key_usage_t usage,
-                                   psa_algorithm_t alg, size_t bits,
+psa_status_t pk_psa_genkey_generic(psa_key_type_t type, size_t bits,
+                                   psa_key_usage_t usage, psa_algorithm_t alg,
                                    mbedtls_svc_key_id_t *key)
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -496,9 +514,9 @@
 {
     mbedtls_svc_key_id_t key;
 
-    pk_psa_genkey_generic(PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1),
+    pk_psa_genkey_generic(PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1), 256,
                           PSA_KEY_USAGE_SIGN_HASH, PSA_ALG_ECDSA(PSA_ALG_SHA_256),
-                          256, &key);
+                          &key);
 
     return key;
 }
@@ -511,8 +529,8 @@
 {
     mbedtls_svc_key_id_t key;
 
-    pk_psa_genkey_generic(PSA_KEY_TYPE_RSA_KEY_PAIR, PSA_KEY_USAGE_SIGN_HASH,
-                          PSA_ALG_RSA_PKCS1V15_SIGN_RAW, 1024, &key);
+    pk_psa_genkey_generic(PSA_KEY_TYPE_RSA_KEY_PAIR, 1024, PSA_KEY_USAGE_SIGN_HASH,
+                          PSA_ALG_RSA_PKCS1V15_SIGN_RAW, &key);
 
     return key;
 }
@@ -2271,8 +2289,8 @@
 
 #if defined(PSA_WANT_KEY_TYPE_DH_KEY_PAIR_GENERATE)
     /* Generate a key type that is not handled by the PK module. */
-    PSA_ASSERT(pk_psa_genkey_generic(PSA_KEY_TYPE_DH_KEY_PAIR(PSA_DH_FAMILY_RFC7919),
-                                     PSA_KEY_USAGE_EXPORT, PSA_ALG_NONE, 2048, &key_id));
+    PSA_ASSERT(pk_psa_genkey_generic(PSA_KEY_TYPE_DH_KEY_PAIR(PSA_DH_FAMILY_RFC7919), 2048,
+                                     PSA_KEY_USAGE_EXPORT, PSA_ALG_NONE, &key_id));
     TEST_EQUAL(mbedtls_pk_copy_from_psa(key_id, &pk_ctx), MBEDTLS_ERR_PK_BAD_INPUT_DATA);
     psa_destroy_key(key_id);
 #endif /* PSA_WANT_KEY_TYPE_DH_KEY_PAIR_GENERATE */
@@ -2280,29 +2298,13 @@
 #if defined(MBEDTLS_PK_HAVE_ECC_KEYS) && defined(PSA_WANT_ECC_SECP_R1_256) && \
     defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE)
     /* Generate an EC key which cannot be exported. */
-    PSA_ASSERT(pk_psa_genkey_generic(PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1),
-                                     0, PSA_ALG_NONE, 256, &key_id));
-    TEST_EQUAL(mbedtls_pk_copy_from_psa(key_id, &pk_ctx), MBEDTLS_ERR_PK_BAD_INPUT_DATA);
-    psa_destroy_key(key_id);
-
-    /* Use valid exportable EC key with wrong alorithm. */
-    PSA_ASSERT(pk_psa_genkey_generic(PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1),
-                                     PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH,
-                                     PSA_ALG_ECDH, 256, &key_id));
+    PSA_ASSERT(pk_psa_genkey_generic(PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1), 256,
+                                     0, PSA_ALG_NONE, &key_id));
     TEST_EQUAL(mbedtls_pk_copy_from_psa(key_id, &pk_ctx), MBEDTLS_ERR_PK_BAD_INPUT_DATA);
     psa_destroy_key(key_id);
 #endif /* MBEDTLS_PK_HAVE_ECC_KEYS && PSA_WANT_ECC_SECP_R1_256 &&
           PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE */
 
-#if defined(MBEDTLS_RSA_C) && defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE)
-    /* Use valid exportable RSA key with wrong alorithm. */
-    PSA_ASSERT(pk_psa_genkey_generic(PSA_KEY_TYPE_RSA_KEY_PAIR,
-                                     PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH,
-                                     PSA_ALG_CMAC, 2048, &key_id));
-    TEST_EQUAL(mbedtls_pk_copy_from_psa(key_id, &pk_ctx), MBEDTLS_ERR_PK_BAD_INPUT_DATA);
-    psa_destroy_key(key_id);
-#endif /* MBEDTLS_RSA_C && PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE */
-
 exit:
     mbedtls_pk_free(&pk_ctx);
     psa_destroy_key(key_id);
@@ -2312,14 +2314,12 @@
 
 /* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C*/
 void pk_copy_from_psa_success(data_t *priv_key_data, int key_type_arg,
-                              int key_bits_arg, int key_usage_arg,
                               int key_alg_arg)
 {
     psa_key_type_t key_type = key_type_arg;
-    psa_key_type_t pub_key_type;
-    size_t key_bits = key_bits_arg;
-    psa_key_usage_t key_usage = key_usage_arg;
     psa_algorithm_t key_alg = key_alg_arg;
+    psa_key_usage_t key_usage = PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH |
+                                PSA_KEY_USAGE_EXPORT;
     mbedtls_pk_context pk_priv, pk_pub;
     mbedtls_svc_key_id_t priv_key_id = MBEDTLS_SVC_KEY_ID_INIT;
     mbedtls_svc_key_id_t pub_key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -2329,6 +2329,10 @@
     unsigned char out_buf2[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
     size_t out_buf_len, out_buf2_len;
 
+    mbedtls_pk_init(&pk_priv);
+    mbedtls_pk_init(&pk_pub);
+    PSA_INIT();
+
     /* Get the MD type to be used for the tests below from the provided key policy. */
     mbedtls_md_type_t md_for_test = MBEDTLS_MD_SHA256; /* Default */
     int is_psa_hash_alg_specified = 0;
@@ -2342,54 +2346,49 @@
     TEST_CALLOC(in_buf, in_buf_len);
     memset(in_buf, 0x1, in_buf_len);
 
-    mbedtls_pk_init(&pk_priv);
-    mbedtls_pk_init(&pk_pub);
-    PSA_INIT();
-
-    /* Generate a private key in PSA and create a PK context from it. */
-    PSA_ASSERT(pk_psa_import_key(priv_key_data->x, priv_key_data->len,
-                                 key_type, key_usage, key_alg, key_bits, &priv_key_id));
-    TEST_EQUAL(mbedtls_pk_copy_from_psa(priv_key_id, &pk_priv), 0);
-
-    /* Starting from the private key above, create another PSA slot for the public
-     * one and create a new PK context from it. */
-    if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(key_type)) {
-        pub_key_type = PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_KEY_TYPE_ECC_GET_FAMILY(key_type));
-    } else if (key_type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
-        pub_key_type = PSA_KEY_TYPE_RSA_PUBLIC_KEY;
-    } else {
-        TEST_FAIL("Key type can only be EC or RSA key pair");
+    if (key_type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
+        key_usage |= PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT;
     }
 
+    /* Create a private key in PSA and create a PK context from it. */
+    PSA_ASSERT(pk_psa_import_key(priv_key_data->x, priv_key_data->len,
+                                 key_type, key_usage, key_alg, &priv_key_id));
+    TEST_EQUAL(mbedtls_pk_copy_from_psa(priv_key_id, &pk_priv), 0);
+
     /* Generate a 2nd PK contex using only the public key derived from its private
      * counterpart generated above. */
-    pub_key_id = pk_psa_pub_key_from_priv(priv_key_id, pub_key_type, key_usage, key_alg, key_bits);
+    pub_key_id = psa_pub_key_from_priv(priv_key_id);
     TEST_EQUAL(mbedtls_pk_copy_from_psa(pub_key_id, &pk_pub), 0);
 
-    /* Test #1: check that the 2 generated PK contexts form a valid private/public key pair. */
+    /* Test #1: check that the generated PK contexts are of the correct type. */
+    if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(key_type)) {
+        TEST_EQUAL(mbedtls_pk_get_type(&pk_priv), MBEDTLS_PK_ECKEY);
+        TEST_EQUAL(mbedtls_pk_get_type(&pk_pub), MBEDTLS_PK_ECKEY);
+    } else {
+        TEST_EQUAL(mbedtls_pk_get_type(&pk_priv), MBEDTLS_PK_RSA);
+        TEST_EQUAL(mbedtls_pk_get_type(&pk_pub), MBEDTLS_PK_RSA);
+    }
+
+    /* Test #2: check that the 2 generated PK contexts form a valid private/public key pair. */
     TEST_EQUAL(mbedtls_pk_check_pair(&pk_pub, &pk_priv, mbedtls_test_rnd_std_rand, NULL), 0);
 
-    /* Test #2: sign/verify with the following pattern:
+    /* Test #3: sign/verify with the following pattern:
      * - Sign using the PK context generated from the private key.
      * - Verify from the same PK context used for signature.
      * - Verify with the PK context generated using public key.
      * - Verify using the public PSA key directly.
      *
-     * Note: mbedtls_pk_verify_ext() is tested only for RSA keys with PKCS1 v2.1
-     *      padding. For other key types this functions works the same as
-     *      mbedtls_pk_verify().
-     *
      * Note: PSS requires the hash to be specified on sign operation (i.e. not
      *       null or any), so in case PSA_ALG_ANY_HASH is provided as input, we
      *       use mbedtls_pk_sign_ext() instead of mbedtls_pk_sign().
      */
-    mbedtls_pk_rsassa_pss_options pss_opt = {
-        .mgf1_hash_id = md_for_test,
-        .expected_salt_len = MBEDTLS_RSA_SALT_LEN_ANY,
-    };
-
     if ((PSA_ALG_IS_RSA_OAEP(key_alg) || PSA_ALG_IS_RSA_PSS(key_alg)) &&
         (!is_psa_hash_alg_specified)) {
+        /* Ensure pk_sign() fails without crashing. */
+        TEST_EQUAL(mbedtls_pk_sign(&pk_priv, md_for_test, in_buf, in_buf_len,
+                                   out_buf, sizeof(out_buf), &out_buf_len,
+                                   mbedtls_test_rnd_std_rand, NULL),
+                   MBEDTLS_ERR_RSA_BAD_INPUT_DATA);
         TEST_EQUAL(mbedtls_pk_sign_ext(MBEDTLS_PK_RSASSA_PSS, &pk_priv, md_for_test,
                                        in_buf, in_buf_len,
                                        out_buf, sizeof(out_buf), &out_buf_len,
@@ -2405,52 +2404,40 @@
     TEST_EQUAL(mbedtls_pk_verify(&pk_pub, md_for_test, in_buf, in_buf_len,
                                  out_buf, out_buf_len), 0);
 
-    if (PSA_ALG_IS_RSA_OAEP(key_alg) || PSA_ALG_IS_RSA_PSS(key_alg)) {
-        TEST_EQUAL(mbedtls_pk_verify_ext(MBEDTLS_PK_RSASSA_PSS, &pss_opt,
-                                         &pk_priv, md_for_test, in_buf, in_buf_len,
-                                         out_buf, out_buf_len), 0);
-        TEST_EQUAL(mbedtls_pk_verify_ext(MBEDTLS_PK_RSASSA_PSS, &pss_opt,
-                                         &pk_pub, md_for_test, in_buf, in_buf_len,
-                                         out_buf, out_buf_len), 0);
-    }
-
-#if defined(MBEDTLS_PSA_UTIL_HAVE_ECDSA)
     if (PSA_ALG_IS_SIGN(key_alg) && is_psa_hash_alg_specified) {
+#if defined(MBEDTLS_PSA_UTIL_HAVE_ECDSA)
         /* ECDSA signature requires PK->PSA format conversion. */
         if (PSA_ALG_IS_ECDSA(key_alg)) {
             TEST_EQUAL(mbedtls_ecdsa_der_to_raw(mbedtls_pk_get_bitlen(&pk_pub),
                                                 out_buf, out_buf_len, out_buf,
                                                 sizeof(out_buf), &out_buf_len), 0);
         }
+#endif /* MBEDTLS_PSA_UTIL_HAVE_ECDSA */
         PSA_ASSERT(psa_verify_hash(pub_key_id, key_alg, in_buf, in_buf_len,
                                    out_buf, out_buf_len));
     }
 
-    /* Test #3: check sign/verify interoperability also in the opposite direction:
+    /* Test #4: check sign/verify interoperability also in the opposite direction:
      * sign with PSA and verify with PK. Key's policy must include a valid hash
      * algorithm (not any).
      */
     if (PSA_ALG_IS_SIGN(key_alg) && is_psa_hash_alg_specified) {
         PSA_ASSERT(psa_sign_hash(priv_key_id, key_alg, in_buf, in_buf_len,
                                  out_buf, sizeof(out_buf), &out_buf_len));
+#if defined(MBEDTLS_PSA_UTIL_HAVE_ECDSA)
         /*  ECDSA signature requires PSA->PK format conversion */
         if (PSA_ALG_IS_ECDSA(key_alg)) {
             TEST_EQUAL(mbedtls_ecdsa_raw_to_der(mbedtls_pk_get_bitlen(&pk_pub),
                                                 out_buf, out_buf_len, out_buf,
                                                 sizeof(out_buf), &out_buf_len), 0);
         }
+#endif /* MBEDTLS_PSA_UTIL_HAVE_ECDSA */
         TEST_EQUAL(mbedtls_pk_verify(&pk_pub, md_for_test, in_buf, in_buf_len,
                                      out_buf, out_buf_len), 0);
     }
-#endif /* MBEDTLS_PSA_UTIL_HAVE_ECDSA */
 
-    /* Test #4: in case of RSA key pair try also encryption/decryption.
-     * Hash algorithm does not matter when padding mode is PKCS1 v1.5, whereas
-     * it must be not null or any in case of v2.1 (OAEP). */
-    int test_encryption = (PSA_ALG_IS_RSA_PKCS1V15_SIGN(key_alg) ||
-                           (key_alg == PSA_ALG_RSA_PKCS1V15_CRYPT) ||
-                           is_psa_hash_alg_specified);
-    if ((key_type == PSA_KEY_TYPE_RSA_KEY_PAIR) && test_encryption) {
+    /* Test #5: in case of RSA key pair try also encryption/decryption. */
+    if (PSA_ALG_IS_ASYMMETRIC_ENCRYPTION(key_alg)) {
         /* Encrypt with the public key only PK context. */
         TEST_EQUAL(mbedtls_pk_encrypt(&pk_pub, in_buf, in_buf_len,
                                       out_buf, &out_buf_len, sizeof(out_buf),