test_suite_pk: modify pk_genkey() in order to use predefined keys

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 442a362..c39b8f7 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -26,9 +26,6 @@
 /* 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"
-
 #define RSA_KEY_SIZE   MBEDTLS_RSA_GEN_KEY_MIN_BITS
 #define RSA_KEY_LEN   (MBEDTLS_RSA_GEN_KEY_MIN_BITS/8)
 
@@ -185,120 +182,55 @@
 #define MBEDTLS_MD_ALG_FOR_TEST         MBEDTLS_MD_SHA512
 #endif
 
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-static int pk_genkey_ec(mbedtls_pk_context *pk, mbedtls_ecp_group_id grp_id)
-{
-    psa_status_t status;
-    psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
-    size_t curve_bits;
-    psa_ecc_family_t curve = mbedtls_ecc_group_to_psa(grp_id, &curve_bits);
-    int ret;
+const char *curve_names_lut[] = {
+    [MBEDTLS_ECP_DP_SECP192R1] = "secp192r1",
+    [MBEDTLS_ECP_DP_SECP256R1] = "secp256r1",
+    [MBEDTLS_ECP_DP_SECP384R1] = "secp384r1",
+    [MBEDTLS_ECP_DP_SECP521R1] = "secp521r1",
+    [MBEDTLS_ECP_DP_BP256R1] = "brainpoolP256r1",
+    [MBEDTLS_ECP_DP_BP384R1] = "brainpoolP384r1",
+    [MBEDTLS_ECP_DP_BP512R1] = "brainpoolP512r1",
+    [MBEDTLS_ECP_DP_CURVE25519] = "x25519",
+    [MBEDTLS_ECP_DP_SECP192K1] = "secp192k1",
+    [MBEDTLS_ECP_DP_SECP256K1] = "secp256k1",
+    [MBEDTLS_ECP_DP_CURVE448] = "x448",
+};
 
-    if (curve == 0) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
-
-    psa_set_key_type(&key_attr, PSA_KEY_TYPE_ECC_KEY_PAIR(curve));
-    psa_set_key_bits(&key_attr, curve_bits);
-    psa_key_usage_t usage = PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_COPY;
-    psa_algorithm_t sign_alg = 0;
-    psa_algorithm_t derive_alg = 0;
-    if (mbedtls_pk_get_type(pk) != MBEDTLS_PK_ECDSA) {
-        usage |= PSA_KEY_USAGE_DERIVE;
-        derive_alg = PSA_ALG_ECDH;
-    }
-    if (mbedtls_pk_get_type(pk) != MBEDTLS_PK_ECKEY_DH &&
-        curve != PSA_ECC_FAMILY_MONTGOMERY) {
-        usage |= PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_SIGN_MESSAGE;
-#if defined(MBEDTLS_ECDSA_DETERMINISTIC)
-        sign_alg = PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH);
-#else
-        sign_alg = PSA_ALG_ECDSA(PSA_ALG_ANY_HASH);
-#endif
-    }
-    if (derive_alg != 0) {
-        psa_set_key_algorithm(&key_attr, derive_alg);
-        if (sign_alg != 0) {
-            psa_set_key_enrollment_algorithm(&key_attr, sign_alg);
-        }
-    } else {
-        psa_set_key_algorithm(&key_attr, sign_alg);
-    }
-    psa_set_key_usage_flags(&key_attr, usage);
-
-    status = psa_generate_key(&key_attr, &pk->priv_id);
-    if (status != PSA_SUCCESS) {
-        return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-    }
-
-    status = psa_export_public_key(pk->priv_id, pk->pub_raw, sizeof(pk->pub_raw),
-                                   &pk->pub_raw_len);
-    if (status != PSA_SUCCESS) {
-        ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-        goto exit;
-    }
-
-    pk->ec_family = curve;
-    pk->ec_bits = curve_bits;
-
-    return 0;
-
-exit:
-    status = psa_destroy_key(pk->priv_id);
-    return (ret != 0) ? ret : psa_pk_status_to_mbedtls(status);
-}
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
-
-/** Generate a key of the desired type.
+/** Fill the provided PK context with a proper key.
+ *
+ * Instead of generating a new key every time, use predefined ones to speed up
+ * testing.
+ * This function assumes that the PK context has already been setup
+ * (mbedtls_pk_setup() has been called on the PK context ) so that it
+ * can determine the key type to be loaded from the PK context itself.
  *
  * \param pk               The PK object to fill. It must have been initialized
  *                         with mbedtls_pk_setup().
  * \param curve_or_keybits - For RSA keys, the key size in bits.
  *                         - For EC keys, the curve (\c MBEDTLS_ECP_DP_xxx).
  *
- * \return                 The status from the underlying type-specific key
- *                         generation function.
- * \return                 -1 if the key type is not recognized.
+ * \return                 0 on success or a negative value otherwise.
  */
 static int pk_genkey(mbedtls_pk_context *pk, int curve_or_keybits)
 {
-    (void) pk;
-    (void) curve_or_keybits;
+    char file_name[128] = { 0 };
+    int ret;
+    /* Dump pk_info since this is overridden by mbedtls_pk_parse_keyfile(). */
+    const mbedtls_pk_info_t *original_pk_info = pk->pk_info;
 
-#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
     if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_RSA) {
-        return mbedtls_rsa_gen_key(mbedtls_pk_rsa(*pk),
-                                   mbedtls_test_rnd_std_rand, NULL,
-                                   curve_or_keybits, 3);
+        sprintf(file_name, "data_files/rsa_%d.der", curve_or_keybits);
+    } else if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY ||
+               mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY_DH ||
+               mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECDSA) {
+        sprintf(file_name, "data_files/ec_%s.der", curve_names_lut[curve_or_keybits]);
     }
-#endif
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
-    if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY ||
-        mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY_DH ||
-        mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECDSA) {
-        int ret;
 
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-        ret = pk_genkey_ec(pk, curve_or_keybits);
-        if (ret != 0) {
-            return ret;
-        }
+    ret = mbedtls_pk_parse_keyfile(pk, file_name, NULL, mbedtls_test_rnd_std_rand, NULL);
+    /* Restore pk_info. */
+    pk->pk_info = original_pk_info;
 
-        return 0;
-#else
-        ret = mbedtls_ecp_group_load(&mbedtls_pk_ec_rw(*pk)->grp, curve_or_keybits);
-        if (ret != 0) {
-            return ret;
-        }
-        return mbedtls_ecp_gen_keypair(&mbedtls_pk_ec_rw(*pk)->grp,
-                                       &mbedtls_pk_ec_rw(*pk)->d,
-                                       &mbedtls_pk_ec_rw(*pk)->Q,
-                                       mbedtls_test_rnd_std_rand, NULL);
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
-
-    }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
-    return -1;
+    return ret;
 }
 
 #if defined(MBEDTLS_PSA_CRYPTO_C)