From 6989284cfaf7c8f5e3c223d9259f80e72425c389 Mon Sep 17 00:00:00 2001
From: Valerio Setti <valerio.setti@nordicsemi.no>
Date: Fri, 5 May 2023 15:55:25 +0200
Subject: [PATCH] pk: update pk_derive_public_key to use the new public key
 solution

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
---
 library/pkparse.c | 28 +++++++++++++---------------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/library/pkparse.c b/library/pkparse.c
index f21c45b16b54..1be49711fdcd 100644
--- a/library/pkparse.c
+++ b/library/pkparse.c
@@ -525,21 +525,17 @@ static int pk_use_ecparams(const mbedtls_asn1_buf *params, mbedtls_pk_context *p
 /*
  * Helper function for deriving a public key from its private counterpart.
  */
-static int pk_derive_public_key(mbedtls_ecp_keypair *eck,
+static int pk_derive_public_key(mbedtls_pk_context *pk,
                                 const unsigned char *d, size_t d_len,
                                 int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
     int ret;
+    mbedtls_ecp_keypair *eck = mbedtls_pk_ec(*pk);
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     psa_status_t status, destruction_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(eck->grp.id, &curve_bits);
-    /* This buffer is used to store the private key at first and then the
-     * public one (but not at the same time). Therefore we size it for the
-     * latter since it's bigger. */
-    unsigned char key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
-    size_t key_len;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
 
     (void) f_rng;
@@ -554,9 +550,8 @@ static int pk_derive_public_key(mbedtls_ecp_keypair *eck,
         return ret;
     }
 
-    mbedtls_platform_zeroize(key_buf, sizeof(key_buf));
-
-    status = psa_export_public_key(key_id, key_buf, sizeof(key_buf), &key_len);
+    status = psa_export_public_key(key_id, pk->pk_raw, sizeof(pk->pk_raw),
+                                   &pk->pk_raw_len);
     ret = psa_pk_status_to_mbedtls(status);
     destruction_status = psa_destroy_key(key_id);
     if (ret != 0) {
@@ -565,10 +560,12 @@ static int pk_derive_public_key(mbedtls_ecp_keypair *eck,
         return psa_pk_status_to_mbedtls(destruction_status);
     }
 
-    ret = mbedtls_ecp_point_read_binary(&eck->grp, &eck->Q, key_buf, key_len);
+    ret = mbedtls_ecp_point_read_binary(&eck->grp, &eck->Q, pk->pk_raw,
+                                        pk->pk_raw_len);
 #else /* MBEDTLS_USE_PSA_CRYPTO */
     (void) d;
     (void) d_len;
+    (void) pk;
 
     ret = mbedtls_ecp_mul(&eck->grp, &eck->Q, &eck->d, &eck->grp.G, f_rng, p_rng);
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
@@ -607,10 +604,11 @@ static int pk_use_ecparams_rfc8410(const mbedtls_asn1_buf *params,
  *
  * CurvePrivateKey ::= OCTET STRING
  */
-static int pk_parse_key_rfc8410_der(mbedtls_ecp_keypair *eck,
+static int pk_parse_key_rfc8410_der(mbedtls_pk_context *pk,
                                     unsigned char *key, size_t keylen, const unsigned char *end,
                                     int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
+    mbedtls_ecp_keypair *eck = mbedtls_pk_ec(*pk);
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     size_t len;
 
@@ -630,7 +628,7 @@ static int pk_parse_key_rfc8410_der(mbedtls_ecp_keypair *eck,
     // pk_parse_key_pkcs8_unencrypted_der() only supports version 1 PKCS8 keys,
     // which never contain a public key. As such, derive the public key
     // unconditionally.
-    if ((ret = pk_derive_public_key(eck, key, len, f_rng, p_rng)) != 0) {
+    if ((ret = pk_derive_public_key(pk, key, len, f_rng, p_rng)) != 0) {
         mbedtls_ecp_keypair_free(eck);
         return ret;
     }
@@ -1235,7 +1233,7 @@ static int pk_parse_key_sec1_der(mbedtls_pk_context *pk,
     }
 
     if (!pubkey_done) {
-        if ((ret = pk_derive_public_key(eck, d, d_len, f_rng, p_rng)) != 0) {
+        if ((ret = pk_derive_public_key(pk, d, d_len, f_rng, p_rng)) != 0) {
             mbedtls_ecp_keypair_free(eck);
             return ret;
         }
@@ -1347,9 +1345,9 @@ static int pk_parse_key_pkcs8_unencrypted_der(
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (mbedtls_pk_is_rfc8410_curve(ec_grp_id)) {
             if ((ret =
-                     pk_use_ecparams_rfc8410(&params, ec_grp_id, &mbedtls_pk_ec(*pk)->grp)) != 0 ||
+                     pk_use_ecparams_rfc8410(&params, ec_grp_id, pk)) != 0 ||
                 (ret =
-                     pk_parse_key_rfc8410_der(mbedtls_pk_ec(*pk), p, len, end, f_rng,
+                     pk_parse_key_rfc8410_der(pk, p, len, end, f_rng,
                                               p_rng)) != 0) {
                 mbedtls_pk_free(pk);
                 return ret;