tests: fixes for using the new public key raw format

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pk.c b/library/pk.c
index 7a047df..47c19b2 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -196,6 +196,39 @@
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+int mbedtls_pk_update_public_key_from_keypair(mbedtls_pk_context *pk,
+                                              mbedtls_ecp_keypair *ecp_keypair)
+{
+    int ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+
+    if (pk == NULL) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+    /* The raw public key storing mechanism is only supported for EC keys so
+     * we fail silently for other ones. */
+    if ((pk->pk_info->type != MBEDTLS_PK_ECKEY) &&
+        (pk->pk_info->type != MBEDTLS_PK_ECKEY_DH) &&
+        (pk->pk_info->type != MBEDTLS_PK_ECDSA)) {
+        return 0;
+    }
+
+    ret = mbedtls_ecp_point_write_binary(&ecp_keypair->grp, &ecp_keypair->Q,
+                                         MBEDTLS_ECP_PF_UNCOMPRESSED,
+                                         &pk->pub_raw_len,
+                                         pk->pub_raw,
+                                         MBEDTLS_PK_MAX_EC_PUBKEY_RAW_LEN);
+    if (ret != 0) {
+        return ret;
+    }
+
+    pk->ec_family = mbedtls_ecc_group_to_psa(ecp_keypair->grp.id,
+                                             &pk->ec_bits);
+
+    return 0;
+}
+#endif /* MBEDTLS_PK_USE_PSA_EC_PUB_KEY */
+
 #if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
 /*
  * Initialize an RSA-alt context
diff --git a/library/pk_internal.h b/library/pk_internal.h
index 09f3d85..dbb7bc1 100644
--- a/library/pk_internal.h
+++ b/library/pk_internal.h
@@ -23,10 +23,16 @@
 #ifndef MBEDTLS_PK_INTERNAL_H
 #define MBEDTLS_PK_INTERNAL_H
 
+#include "mbedtls/pk.h"
+
 #if defined(MBEDTLS_ECP_LIGHT)
 #include "mbedtls/ecp.h"
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#include "psa/crypto.h"
+#endif
+
 #if defined(MBEDTLS_ECP_LIGHT)
 /**
  * Public function mbedtls_pk_ec() can be used to get direct access to the
@@ -85,4 +91,19 @@
 #endif /* MBEDTLS_ECP_DP_CURVE25519_ENABLED || MBEDTLS_ECP_DP_CURVE448_ENABLED */
 #endif /* MBEDTLS_ECP_LIGHT */
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+/**
+ * \brief   Copy the public key content in raw format from "ctx->pk_ctx"
+ *          (which is an ecp_keypair) into the internal "ctx->pub_raw" buffer.
+ *
+ * \note    This is a temporary function that can be removed as soon as the pk
+ *          module is free from ECP_C
+ *
+ * \param pk   It is the pk_context which is going to be updated. It acts both
+ *             as input and output.
+ */
+int mbedtls_pk_update_public_key_from_keypair(mbedtls_pk_context *pk,
+                                              mbedtls_ecp_keypair *ecp_keypair);
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+
 #endif /* MBEDTLS_PK_INTERNAL_H */
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index f36c6be..d397374 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1,5 +1,6 @@
 /* BEGIN_HEADER */
 #include "mbedtls/pk.h"
+#include "pk_internal.h"
 
 /* For error codes */
 #include "mbedtls/asn1.h"
@@ -24,16 +25,15 @@
 #define RSA_KEY_SIZE 512
 #define RSA_KEY_LEN   64
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO) && defined(MBEDTLS_ECP_LIGHT)
-static int pk_genkey_ec(mbedtls_ecp_group *grp,
-                        mbedtls_mpi *d, mbedtls_ecp_point *Q)
+#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;
+    mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
     psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     size_t curve_bits;
-    psa_ecc_family_t curve = mbedtls_ecc_group_to_psa(grp->id,
-                                                      &curve_bits);
+    psa_ecc_family_t curve = mbedtls_ecc_group_to_psa(grp_id, &curve_bits);
     unsigned char key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
     size_t key_len;
     int ret;
@@ -53,26 +53,33 @@
         goto exit;
     }
 
-    ret = mbedtls_mpi_read_binary(d, key_buf, key_len);
+    ret = mbedtls_mpi_read_binary(&eck->d, key_buf, key_len);
     if (ret != 0) {
         goto exit;
     }
 
-    status = psa_export_public_key(key_id, key_buf, sizeof(key_buf),
-                                   &key_len);
+    status = psa_export_public_key(key_id, pk->pub_raw, sizeof(pk->pub_raw),
+                                   &pk->pub_raw_len);
     if (status != PSA_SUCCESS) {
         ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
         goto exit;
     }
 
-    ret = mbedtls_ecp_point_read_binary(grp, Q, key_buf, key_len);
+    pk->ec_family = curve;
+    pk->ec_bits = curve_bits;
+
+    status = psa_destroy_key(key_id);
+    if (status != PSA_SUCCESS) {
+        return psa_pk_status_to_mbedtls(status);
+    }
+
+    return 0;
 
 exit:
-    psa_destroy_key(key_id);
-
-    return ret;
+    status = psa_destroy_key(key_id);
+    return (ret != 0) ? ret : psa_pk_status_to_mbedtls(status);
 }
-#endif /* MBEDTLS_USE_PSA_CRYPTO && MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
 /** Generate a key of the desired type.
  *
@@ -102,22 +109,36 @@
         mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY_DH ||
         mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECDSA) {
         int ret;
-        if ((ret = mbedtls_ecp_group_load(&mbedtls_pk_ec_rw(*pk)->grp,
-                                          parameter)) != 0) {
+
+        ret = mbedtls_ecp_group_load(&mbedtls_pk_ec_rw(*pk)->grp, parameter);
+        if (ret != 0) {
             return ret;
         }
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-        return pk_genkey_ec(&mbedtls_pk_ec_rw(*pk)->grp,
-                            &mbedtls_pk_ec_rw(*pk)->d,
-                            &mbedtls_pk_ec_rw(*pk)->Q);
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+        mbedtls_ecp_group grp;
+        /* Duplicating the mbedtls_ecp_group_load call to make this part
+         * more future future proof for when ECP_C will not be defined. */
+        mbedtls_ecp_group_init(&grp);
+        ret = mbedtls_ecp_group_load(&grp, parameter);
+        if (ret != 0) {
+            return ret;
+        }
+        ret = pk_genkey_ec(pk, grp.id);
+        if (ret != 0) {
+            return ret;
+        }
+        mbedtls_ecp_group_free(&grp);
+
+        return 0;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 #if defined(MBEDTLS_ECP_C)
         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_ECP_C */
+
     }
 #endif /* MBEDTLS_ECP_LIGHT */
     return -1;
@@ -702,7 +723,6 @@
                     data_t *sig, int ret)
 {
     mbedtls_pk_context pk;
-    mbedtls_ecp_keypair *eckey;
 
     mbedtls_pk_init(&pk);
     USE_PSA_INIT();
@@ -710,11 +730,23 @@
     TEST_ASSERT(mbedtls_pk_setup(&pk, mbedtls_pk_info_from_type(type)) == 0);
 
     TEST_ASSERT(mbedtls_pk_can_do(&pk, MBEDTLS_PK_ECDSA));
-    eckey = mbedtls_pk_ec_rw(pk);
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    mbedtls_ecp_keypair ecp;
+    mbedtls_ecp_keypair_init(&ecp);
+
+    TEST_ASSERT(mbedtls_ecp_group_load(&ecp.grp, id) == 0);
+    TEST_ASSERT(mbedtls_ecp_point_read_binary(&ecp.grp, &ecp.Q,
+                                              key->x, key->len) == 0);
+    TEST_ASSERT(mbedtls_pk_update_public_key_from_keypair(&pk, &ecp) == 0);
+
+    mbedtls_ecp_keypair_free(&ecp);
+#else
+    mbedtls_ecp_keypair *eckey = (mbedtls_ecp_keypair *) mbedtls_pk_ec(pk);
 
     TEST_ASSERT(mbedtls_ecp_group_load(&eckey->grp, id) == 0);
     TEST_ASSERT(mbedtls_ecp_point_read_binary(&eckey->grp, &eckey->Q,
                                               key->x, key->len) == 0);
+#endif
 
     // MBEDTLS_MD_NONE is used since it will be ignored.
     TEST_ASSERT(mbedtls_pk_verify(&pk, MBEDTLS_MD_NONE,
diff --git a/tests/suites/test_suite_pkparse.function b/tests/suites/test_suite_pkparse.function
index e0e3300..a49b6d3 100644
--- a/tests/suites/test_suite_pkparse.function
+++ b/tests/suites/test_suite_pkparse.function
@@ -84,10 +84,16 @@
     TEST_ASSERT(res == result);
 
     if (res == 0) {
-        const mbedtls_ecp_keypair *eckey;
         TEST_ASSERT(mbedtls_pk_can_do(&ctx, MBEDTLS_PK_ECKEY));
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+        /* No need to check whether the parsed public point is on the curve or
+         * not because this is already done by the internal "pk_get_ecpubkey()"
+         * function */
+#else
+        const mbedtls_ecp_keypair *eckey;
         eckey = mbedtls_pk_ec_ro(ctx);
         TEST_ASSERT(mbedtls_ecp_check_pubkey(&eckey->grp, &eckey->Q) == 0);
+#endif
     }
 
 exit: