Fix handling of ECC public keys under MBEDTLS_PK_USE_PSA_EC_DATA

The test code to construct test keys and the implementation had matching
errors: both assumed that there was a PSA public key object. Fix this.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index f218558..13b960a 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -253,6 +253,8 @@
      *   inside the ecp_keypair structure
      * - the following fields are used for all public key operations: signature
      *   verify, key pair check and key write.
+     * - For a key pair, priv_id contains the private key. For a public key,
+     *   priv_id is null.
      * Of course, when MBEDTLS_PK_USE_PSA_EC_DATA is not enabled, the legacy
      * ecp_keypair structure is used for storing the public key and performing
      * all the operations.
diff --git a/library/pk.c b/library/pk.c
index 3b9c537..d0869b8 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -464,16 +464,12 @@
             int sign_ok = (pk_type != MBEDTLS_PK_ECKEY_DH);
             int derive_ok = (pk_type != MBEDTLS_PK_ECDSA);
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-            psa_key_attributes_t old_attributes = PSA_KEY_ATTRIBUTES_INIT;
-            psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-            status = psa_get_key_attributes(pk->priv_id, &old_attributes);
-            if (status != PSA_SUCCESS) {
-                return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+            psa_ecc_family_t family = pk->ec_family;
+            size_t bits = pk->ec_bits;
+            int has_private = 0;
+            if (pk->priv_id != MBEDTLS_SVC_KEY_ID_INIT) {
+                has_private = 1;
             }
-            psa_key_type_t old_type = psa_get_key_type(&old_attributes);
-            int has_private = PSA_KEY_TYPE_IS_KEY_PAIR(old_type);
-            size_t bits = psa_get_key_bits(&old_attributes);
-            psa_ecc_family_t family = PSA_KEY_TYPE_ECC_GET_FAMILY(old_type);
 #else
             const mbedtls_ecp_keypair *ec = mbedtls_pk_ec_ro(*pk);
             int has_private = (ec->d.n != 0);
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 3d7a179..efbe6b0 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -283,9 +283,7 @@
                                         PSA_KEY_USAGE_VERIFY_HASH);
                 psa_set_key_algorithm(&pub_attributes, PSA_ALG_ECDSA_ANY);
                 PSA_ASSERT(psa_destroy_key(pk->priv_id));
-                PSA_ASSERT(psa_import_key(&pub_attributes,
-                                          pk->pub_raw, pk->pub_raw_len,
-                                          &pk->priv_id));
+                pk->priv_id = MBEDTLS_SVC_KEY_ID_INIT;
 #else
                 mbedtls_ecp_keypair *ec = mbedtls_pk_ec_rw(*pk);
                 mbedtls_mpi_free(&ec->d);