pk_wrap: use specific lengths for EC's private key and key-pair

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/include/mbedtls/psa_util.h b/include/mbedtls/psa_util.h
index dc74ac6..f6070dc 100644
--- a/include/mbedtls/psa_util.h
+++ b/include/mbedtls/psa_util.h
@@ -257,6 +257,9 @@
 #define MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH \
     PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)
 
+#define MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH \
+    PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)
+
 /* Expose whatever RNG the PSA subsystem uses to applications using the
  * mbedtls_xxx API. The declarations and definitions here need to be
  * consistent with the implementation in library/psa_crypto_random_impl.h.
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 525f6bc..6fba6e9 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -692,8 +692,7 @@
     /* This buffer contains first the public key (consisting of two public
      * points plus a header byte), then the signature (consisting of two
      * public points). Size it for the public key which is one byte larger. */
-    unsigned char buf[PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(
-            PSA_VENDOR_ECC_MAX_CURVE_BITS )];
+    unsigned char buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
     unsigned char *p;
     psa_algorithm_t psa_sig_md = PSA_ALG_ECDSA_ANY;
     size_t curve_bits;
@@ -876,8 +875,7 @@
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_status_t status;
-    unsigned char buf[PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(
-            PSA_VENDOR_ECC_MAX_CURVE_BITS )];
+    unsigned char buf[MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH];
 #if defined(MBEDTLS_ECDSA_DETERMINISTIC)
     psa_algorithm_t psa_sig_md =
         PSA_ALG_DETERMINISTIC_ECDSA( mbedtls_hash_info_psa_from_md( md_alg ) );