From 7237d5ff5b269f9859d1b97b317a4d2b535b8f9a Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Thu, 18 May 2023 19:00:22 +0200
Subject: [PATCH] test: minor enhancement for using the new private key format

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 tests/suites/test_suite_pk.function      | 47 +++++++++---------------
 tests/suites/test_suite_pkparse.function |  7 +++-
 2 files changed, 22 insertions(+), 32 deletions(-)

diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 7227f9278..a5b50dec4 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -29,13 +29,9 @@
 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);
-    unsigned char key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
-    size_t key_len;
     int ret;
 
     if (curve == 0) {
@@ -44,25 +40,21 @@ static int pk_genkey_ec(mbedtls_pk_context *pk, mbedtls_ecp_group_id grp_id)
 
     psa_set_key_type(&key_attr, PSA_KEY_TYPE_ECC_KEY_PAIR(curve));
     psa_set_key_bits(&key_attr, curve_bits);
-    psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_EXPORT);
+    psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_EXPORT |
+                            PSA_KEY_USAGE_SIGN_HASH |
+                            PSA_KEY_USAGE_SIGN_MESSAGE);
+#if defined(MBEDTLS_ECDSA_DETERMINISTIC)
+    psa_set_key_algorithm(&key_attr, PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH));
+#else
+    psa_set_key_algorithm(&key_attr, PSA_ALG_ECDSA(PSA_ALG_ANY_HASH));
+#endif
 
-    status = psa_generate_key(&key_attr, &key_id);
+    status = psa_generate_key(&key_attr, &pk->priv_id);
     if (status != PSA_SUCCESS) {
         return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
     }
 
-    status = psa_export_key(key_id, key_buf, sizeof(key_buf), &key_len);
-    if (status != PSA_SUCCESS) {
-        ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-        goto exit;
-    }
-
-    ret = mbedtls_mpi_read_binary(&eck->d, key_buf, key_len);
-    if (ret != 0) {
-        goto exit;
-    }
-
-    status = psa_export_public_key(key_id, pk->pub_raw, sizeof(pk->pub_raw),
+    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;
@@ -72,15 +64,10 @@ static int pk_genkey_ec(mbedtls_pk_context *pk, mbedtls_ecp_group_id grp_id)
     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:
-    status = psa_destroy_key(key_id);
+    status = psa_destroy_key(pk->priv_id);
     return (ret != 0) ? ret : psa_pk_status_to_mbedtls(status);
 }
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
@@ -114,10 +101,16 @@ static int pk_genkey(mbedtls_pk_context *pk, int parameter)
         mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECDSA) {
         int ret;
 
+#if defined(MBEDTLS_ECP_C)
         ret = mbedtls_ecp_group_load(&mbedtls_pk_ec_rw(*pk)->grp, parameter);
         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_ECP_C */
 
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
         mbedtls_ecp_group grp;
@@ -136,12 +129,6 @@ static int pk_genkey(mbedtls_pk_context *pk, int parameter)
 
         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 */
diff --git a/tests/suites/test_suite_pkparse.function b/tests/suites/test_suite_pkparse.function
index a49b6d319..6fa78c149 100644
--- a/tests/suites/test_suite_pkparse.function
+++ b/tests/suites/test_suite_pkparse.function
@@ -117,10 +117,13 @@ void pk_parse_keyfile_ec(char *key_file, char *password, int result)
     TEST_ASSERT(res == result);
 
     if (res == 0) {
-        const mbedtls_ecp_keypair *eckey;
         TEST_ASSERT(mbedtls_pk_can_do(&ctx, MBEDTLS_PK_ECKEY));
-        eckey = mbedtls_pk_ec_ro(ctx);
+#if defined(MBEDTLS_ECP_C)
+        const mbedtls_ecp_keypair *eckey = mbedtls_pk_ec_ro(ctx);
         TEST_ASSERT(mbedtls_ecp_check_privkey(&eckey->grp, &eckey->d) == 0);
+#else
+        /* PSA keys are already checked on import so nothing to do here. */
+#endif
     }
 
 exit: