Add PSA interruptible key generation get num ops API

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index e18647f..bb29283 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -8121,8 +8121,7 @@
 uint32_t psa_generate_key_iop_get_num_ops(
     psa_generate_key_iop_t *operation)
 {
-    (void) operation;
-    return 0;
+    return operation->num_ops;
 }
 
 psa_status_t psa_generate_key_iop_setup(
@@ -8197,6 +8196,8 @@
         goto exit;
     }
 
+    operation->num_ops = mbedtls_psa_generate_key_iop_get_num_ops(&operation->ctx);
+
     status = psa_import_key(&operation->attributes,
                             key_data + (sizeof(key_data) - key_len),
                             key_len,
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 b43923e..c6036e7 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -596,6 +596,12 @@
 
 #if defined(MBEDTLS_ECP_RESTARTABLE)
 
+uint32_t mbedtls_psa_generate_key_iop_get_num_ops(
+    mbedtls_psa_generate_key_iop_t *operation)
+{
+    return operation->num_ops;
+}
+
 psa_status_t mbedtls_psa_ecp_generate_key_iop_setup(
     mbedtls_psa_generate_key_iop_t *operation,
     const psa_key_attributes_t *attributes)
diff --git a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h
index f3ff323..8f51fc9 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h
@@ -144,6 +144,17 @@
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length);
 
 /**
+ * \brief Get the total number of ops that a key generation operation has taken
+ *        Since it's start.
+ *
+ * \param[in] operation                 The \c mbedtls_psa_generate_key_iop_t to use.
+ *                                      This must be initialized first.
+ * \return Total number of operations.
+ */
+uint32_t mbedtls_psa_generate_key_iop_get_num_ops(
+    mbedtls_psa_generate_key_iop_t *operation);
+
+/**
  * \brief Setup a new interruptible key generation operation.
  *
  *  \param[in] operation                 The \c mbedtls_psa_generate_key_iop_t to use.