pk: add checks for the returned ECC family

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pk.c b/library/pk.c
index 826c29a..9c4aa16 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -224,6 +224,9 @@
 
     pk->ec_family = mbedtls_ecc_group_to_psa(ecp_keypair->grp.id,
                                              &pk->ec_bits);
+    if (pk->ec_family == 0) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
 
     return 0;
 }
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index e21ec2b..3a3d399 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -1150,6 +1150,10 @@
 #endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
     const size_t curve_bytes = PSA_BITS_TO_BYTES(curve_bits);
 
+    if (curve == 0) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
     psa_set_key_type(&key_attr, PSA_KEY_TYPE_ECC_KEY_PAIR(curve));
     psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_EXPORT);
 
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index d397374..7227f92 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -38,6 +38,10 @@
     size_t key_len;
     int ret;
 
+    if (curve == 0) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
     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);