Add PSA interruptible key generation setup & abort APIs

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h b/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h
index 8367cd3..22d5d84 100644
--- a/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h
+++ b/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h
@@ -238,6 +238,8 @@
 }
 mbedtls_mpi;
 
+#define MBEDTLS_MPI_INIT { 0, 0, 0 }
+
 /**
  * \brief           Initialize an MPI context.
  *
diff --git a/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h b/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
index 7b0a80d..b340614 100644
--- a/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
+++ b/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
@@ -162,6 +162,8 @@
 }
 mbedtls_ecp_point;
 
+#define MBEDTLS_ECP_POINT_INIT { MBEDTLS_MPI_INIT, MBEDTLS_MPI_INIT, MBEDTLS_MPI_INIT }
+
 /**
  * \brief           The ECP group structure.
  *
@@ -250,6 +252,9 @@
 }
 mbedtls_ecp_group;
 
+#define MBEDTLS_ECP_GROUP_INIT { MBEDTLS_ECP_DP_NONE, MBEDTLS_MPI_INIT, MBEDTLS_MPI_INIT, \
+                                 MBEDTLS_MPI_INIT, MBEDTLS_ECP_POINT_INIT, MBEDTLS_MPI_INIT, \
+                                 0, 0, 0, NULL, NULL, NULL, NULL, NULL, 0 }
 /**
  * \name SECTION: Module settings
  *
@@ -419,6 +424,9 @@
 }
 mbedtls_ecp_keypair;
 
+#define MBEDTLS_ECP_KEYPAIR_INIT { MBEDTLS_ECP_GROUP_INIT, MBEDTLS_MPI_INIT, \
+                                   MBEDTLS_ECP_POINT_INIT }
+
 /**
  * The uncompressed point format for Short Weierstrass curves
  * (MBEDTLS_ECP_DP_SECP_XXX and MBEDTLS_ECP_DP_BP_XXX).
diff --git a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
index 749e11b..33a9797 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -596,41 +596,38 @@
 /* Interruptible ECC Key Generation */
 /****************************************************************/
 
-uint32_t psa_generate_key_iop_get_num_ops(
-    psa_generate_key_iop_t *operation)
-{
-    (void) operation;
-    return 0;
-}
+#if defined(MBEDTLS_ECP_RESTARTABLE) && defined(MBEDTLS_ECP_C)
 
-psa_status_t psa_generate_key_iop_setup(
-    psa_generate_key_iop_t *operation,
+psa_status_t mbedtls_psa_generate_key_iop_setup(
+    mbedtls_psa_generate_key_iop_t *operation,
     const psa_key_attributes_t *attributes)
 {
-    (void) operation;
-    (void) attributes;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    return PSA_ERROR_NOT_SUPPORTED;
+    mbedtls_ecp_keypair_init(&operation->ecp);
+
+    psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY(
+        psa_get_key_type(attributes));
+    mbedtls_ecp_group_id grp_id =
+        mbedtls_ecc_group_from_psa(curve, psa_get_key_bits(attributes));
+    if (grp_id == MBEDTLS_ECP_DP_NONE) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    status = mbedtls_ecp_group_load(&operation->ecp.grp, grp_id);
+
+    return mbedtls_to_psa_error(status);
 }
 
-psa_status_t psa_generate_key_iop_complete(
-    psa_generate_key_iop_t *operation,
-    psa_key_id_t *key)
+psa_status_t mbedtls_psa_generate_key_iop_abort(
+    mbedtls_psa_generate_key_iop_t *operation)
 {
-    (void) operation;
-    (void) key;
-
-    return PSA_ERROR_NOT_SUPPORTED;
+    mbedtls_ecp_keypair_free(&operation->ecp);
+    operation->num_ops = 0;
+    return PSA_SUCCESS;
 }
 
-psa_status_t psa_generate_key_iop_abort(
-    psa_generate_key_iop_t *operation)
-{
-    (void) operation;
-
-    return PSA_ERROR_NOT_SUPPORTED;
-}
-
+#endif
 /****************************************************************/
 /* Interruptible ECC Key Agreement */
 /****************************************************************/