Add PSA interruptible key generation complete API

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tests/include/test/psa_test_wrappers.h b/tests/include/test/psa_test_wrappers.h
index 7ab2bea..ef115f5 100644
--- a/tests/include/test/psa_test_wrappers.h
+++ b/tests/include/test/psa_test_wrappers.h
@@ -370,7 +370,7 @@
 
 psa_status_t mbedtls_test_wrap_psa_generate_key_iop_complete(
     psa_generate_key_iop_t *arg0_operation,
-    psa_key_id_t *arg1_key);
+    mbedtls_svc_key_id_t *arg1_key);
 #define psa_generate_key_iop_complete(arg0_operation, arg1_key) \
     mbedtls_test_wrap_psa_generate_key_iop_complete(arg0_operation, arg1_key)
 
diff --git a/tests/src/psa_test_wrappers.c b/tests/src/psa_test_wrappers.c
index 6e42a8b..b5419b9 100644
--- a/tests/src/psa_test_wrappers.c
+++ b/tests/src/psa_test_wrappers.c
@@ -633,7 +633,7 @@
 /* Wrapper for psa_generate_key_iop_complete */
 psa_status_t mbedtls_test_wrap_psa_generate_key_iop_complete(
     psa_generate_key_iop_t *arg0_operation,
-    psa_key_id_t *arg1_key)
+    mbedtls_svc_key_id_t *arg1_key)
 {
     psa_status_t status = (psa_generate_key_iop_complete)(arg0_operation, arg1_key);
     return status;
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 5f44cc3..d41a232 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -8110,6 +8110,8 @@
 
     status = mbedtls_psa_generate_key_iop_abort(&operation->ctx);
 
+    psa_reset_key_attributes(&operation->attributes);
+
     operation->id = 0;
 
     return status;
@@ -8178,12 +8180,41 @@
 
 psa_status_t psa_generate_key_iop_complete(
     psa_generate_key_iop_t *operation,
-    psa_key_id_t *key)
+    mbedtls_svc_key_id_t *key)
 {
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status;
+    uint8_t key_data[MBEDTLS_ECP_MAX_BYTES] = { 0 };
+    size_t key_len = 0;
+
+    if (operation->id == 0 || operation->error_occurred) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    status = mbedtls_psa_generate_key_complete(&operation->ctx, key_data,
+                                               MBEDTLS_ECP_MAX_BYTES, &key_len);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    status = psa_import_key(&operation->attributes, key_data, key_len, key);
+
+exit:
+    if (status != PSA_OPERATION_INCOMPLETE) {
+        if (status != PSA_SUCCESS) {
+            operation->error_occurred = 1;
+        }
+        psa_generate_key_iop_abort_internal(operation);
+    }
+
+    mbedtls_platform_zeroize(key_data, MBEDTLS_ECP_MAX_BYTES);
+    return status;
+#else
     (void) operation;
     (void) key;
 
-    return PSA_ERROR_NOT_SUPPORTED;
+    return PSA_ERROR_BAD_STATE;
+#endif
 }
 
 psa_status_t psa_generate_key_iop_abort(
diff --git a/tf-psa-crypto/core/psa_crypto_core.h b/tf-psa-crypto/core/psa_crypto_core.h
index 1753554..cc9a78d 100644
--- a/tf-psa-crypto/core/psa_crypto_core.h
+++ b/tf-psa-crypto/core/psa_crypto_core.h
@@ -456,6 +456,35 @@
     mbedtls_psa_generate_key_iop_t *operation,
     const psa_key_attributes_t *attributes);
 
+
+/**
+ * \brief Continue and eventually complete a key generation operation.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       generate_key_complete entry point. This function behaves as a
+ *       generate_key_complete entry point as defined in the PSA driver
+ *       interface specification for transparent drivers.
+ *
+ * \param[in] operation                  The \c mbedtls_psa_generate_key_iop_t to use.
+ *                                       This must be initialized first and
+ *                                       had \c mbedtls_psa_generate_key_iop_setup()
+ *                                       called successfully.
+ * \param[out] key_output                The buffer to which the generated key
+ *                                       is to be written.
+ * \param[out] key_len                   On success, the number of bytes that make
+ *                                       up the returned key output.
+ * \retval #PSA_SUCCESS
+ *         The key was generated successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT \emptydescription
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY \emptydescription
+ *
+ */
+psa_status_t mbedtls_psa_generate_key_complete(
+    mbedtls_psa_generate_key_iop_t *operation,
+    uint8_t *key_output,
+    size_t key_output_size,
+    size_t *key_len);
+
 /**
  * \brief Abort a key generation operation.
  *
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 acb2482..498072a 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -617,6 +617,32 @@
     return mbedtls_to_psa_error(status);
 }
 
+psa_status_t mbedtls_psa_generate_key_complete(
+    mbedtls_psa_generate_key_iop_t *operation,
+    uint8_t *key_output,
+    size_t key_output_size,
+    size_t *key_len)
+{
+    *key_len = 0;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    status = mbedtls_ecp_gen_privkey(&operation->ecp.grp, &operation->ecp.d,
+                                     mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE);
+
+    if (status) {
+        return mbedtls_to_psa_error(status);
+    }
+
+    operation->num_ops = 1;
+
+    *key_len = mbedtls_mpi_size(&operation->ecp.d);
+    if (*key_len > key_output_size) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+    memcpy(key_output, operation->ecp.d.p, *key_len);
+
+    return mbedtls_to_psa_error(status);
+}
+
 psa_status_t mbedtls_psa_generate_key_iop_abort(
     mbedtls_psa_generate_key_iop_t *operation)
 {
diff --git a/tf-psa-crypto/include/psa/crypto.h b/tf-psa-crypto/include/psa/crypto.h
index 58b6887..cb3b579 100644
--- a/tf-psa-crypto/include/psa/crypto.h
+++ b/tf-psa-crypto/include/psa/crypto.h
@@ -5501,7 +5501,7 @@
  */
 psa_status_t psa_generate_key_iop_complete(
     psa_generate_key_iop_t *operation,
-    psa_key_id_t *key);
+    mbedtls_svc_key_id_t *key);
 
 /**
  * \brief                       Abort a key generation operation.
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
index eaafd90..b6e30c4 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
@@ -10089,6 +10089,7 @@
                   int is_large_key)
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    mbedtls_svc_key_id_t iop_key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t type = type_arg;
     psa_key_usage_t usage = usage_arg;
     size_t bits = bits_arg;
@@ -10096,6 +10097,7 @@
     psa_status_t expected_status = expected_status_arg;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_attributes_t got_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_attributes_t iop_attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_generate_key_iop_t operation = PSA_GENERATE_KEY_IOP_INIT;
 
     PSA_ASSERT(psa_crypto_init());
@@ -10135,6 +10137,12 @@
     expected_status = PSA_ERROR_NOT_SUPPORTED;
 #endif
 
+    /* Test calling complete() without calling setup() will fail. */
+    status = psa_generate_key_iop_complete(&operation, &iop_key);
+    TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+
+    psa_generate_key_iop_abort(&operation);
+
     status = psa_generate_key_iop_setup(&operation, &attributes);
     TEST_EQUAL(status, expected_status);
 
@@ -10150,6 +10158,22 @@
     status = psa_generate_key_iop_setup(&operation, &attributes);
     TEST_EQUAL(status, expected_status);
 
+    do {
+        status = psa_generate_key_iop_complete(&operation, &iop_key);
+    } while (status == PSA_OPERATION_INCOMPLETE);
+
+    TEST_EQUAL(status, PSA_SUCCESS);
+
+    PSA_ASSERT(psa_get_key_attributes(iop_key, &iop_attributes));
+    TEST_EQUAL(psa_get_key_type(&iop_attributes), type);
+    TEST_EQUAL(psa_get_key_bits(&iop_attributes), bits);
+
+    TEST_EQUAL(mbedtls_test_psa_exercise_key(iop_key, usage, alg, 0), 1);
+
+    /* Test calling complete() 2 times consecutively will fail. */
+    status = psa_generate_key_iop_complete(&operation, &iop_key);
+    TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+
 exit:
     psa_generate_key_iop_abort(&operation);
     /*
@@ -10157,8 +10181,10 @@
      * thus reset them as required.
      */
     psa_reset_key_attributes(&got_attributes);
+    psa_reset_key_attributes(&iop_attributes);
 
     psa_destroy_key(key);
+    psa_destroy_key(iop_key);
     PSA_DONE();
 }
 /* END_CASE */