Merge pull request #9808 from waleed-elmelegy-arm/add-iop-export-pub-key-complete

Add PSA interruptible export public-key complete API
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 60707c6..4db4fbc 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -1676,6 +1676,8 @@
 
     status = mbedtls_psa_ecp_export_public_key_iop_abort(&operation->ctx);
 
+    memset(&operation->ctx, 0, sizeof(operation->ctx));
+
     operation->id = 0;
 
     return status;
@@ -1694,9 +1696,8 @@
 #if defined(MBEDTLS_ECP_RESTARTABLE)
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
-    size_t key_size = 0;
-    psa_key_attributes_t private_key_attributes;
-    psa_key_type_t private_key_type;
+    psa_key_attributes_t key_attributes;
+    psa_key_type_t key_type;
     psa_key_slot_t *slot = NULL;
 
     if (operation->id != 0 || operation->error_occurred) {
@@ -1713,29 +1714,22 @@
         goto exit;
     }
 
-    private_key_attributes = slot->attr;
+    key_attributes = slot->attr;
 
-    private_key_type = psa_get_key_type(&private_key_attributes);
+    key_type = psa_get_key_type(&key_attributes);
 
-    if (!PSA_KEY_TYPE_IS_KEY_PAIR(private_key_type)) {
+    if (!PSA_KEY_TYPE_IS_ASYMMETRIC(key_type)) {
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }
 
-    if (!PSA_KEY_TYPE_IS_ECC_KEY_PAIR(private_key_type)) {
-        status = PSA_ERROR_NOT_SUPPORTED;
-        goto exit;
-    }
-
-    key_size = PSA_EXPORT_KEY_OUTPUT_SIZE(private_key_type,
-                                          psa_get_key_bits(&private_key_attributes));
-    if (key_size == 0) {
+    if (!PSA_KEY_TYPE_IS_ECC(key_type)) {
         status = PSA_ERROR_NOT_SUPPORTED;
         goto exit;
     }
 
     status = mbedtls_psa_ecp_export_public_key_iop_setup(&operation->ctx, slot->key.data,
-                                                         slot->key.bytes, &private_key_attributes);
+                                                         slot->key.bytes, &key_attributes);
 
 exit:
     unlock_status = psa_unregister_read_under_mutex(slot);
@@ -1757,12 +1751,33 @@
                                                 size_t data_size,
                                                 size_t *data_length)
 {
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if (operation->id == 0 || operation->error_occurred) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    status = mbedtls_psa_ecp_export_public_key_iop_complete(&operation->ctx, data, data_size,
+                                                            data_length);
+
+    if (status != PSA_OPERATION_INCOMPLETE) {
+        psa_export_public_key_iop_abort_internal(operation);
+
+        if (status != PSA_SUCCESS) {
+            operation->error_occurred = 1;
+        }
+    }
+
+    return status;
+#else
     (void) operation;
     (void) data;
     (void) data_size;
     (void) data_length;
 
-    return PSA_ERROR_NOT_SUPPORTED;
+    return PSA_ERROR_BAD_STATE;
+#endif
 }
 
 psa_status_t psa_export_public_key_iop_abort(psa_export_public_key_iop_t *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 f90274e..3ca28fa 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -665,18 +665,18 @@
 }
 
 psa_status_t mbedtls_psa_ecp_export_public_key_iop_setup(
-    mbedtls_psa_export_public_key_iop_operation_t *operation,
-    uint8_t *private_key,
-    size_t private_key_len,
-    const psa_key_attributes_t *private_key_attributes)
+    mbedtls_psa_export_public_key_iop_t *operation,
+    uint8_t *key,
+    size_t key_len,
+    const psa_key_attributes_t *key_attributes)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
     status = mbedtls_psa_ecp_load_representation(
-        psa_get_key_type(private_key_attributes),
-        psa_get_key_bits(private_key_attributes),
-        private_key,
-        private_key_len,
+        psa_get_key_type(key_attributes),
+        psa_get_key_bits(key_attributes),
+        key,
+        key_len,
         &operation->key);
     if (status != PSA_SUCCESS) {
         goto exit;
@@ -689,8 +689,35 @@
     return status;
 }
 
+psa_status_t mbedtls_psa_ecp_export_public_key_iop_complete(
+    mbedtls_psa_export_public_key_iop_t *operation,
+    uint8_t *pub_key,
+    size_t pub_key_size,
+    size_t *pub_key_len)
+{
+    int ret = 0;
+
+    if (mbedtls_ecp_is_zero(&operation->key->Q)) {
+        mbedtls_psa_interruptible_set_max_ops(psa_interruptible_get_max_ops());
+
+        ret = mbedtls_ecp_mul_restartable(&operation->key->grp, &operation->key->Q,
+                                          &operation->key->d, &operation->key->grp.G,
+                                          mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE,
+                                          &operation->restart_ctx);
+        operation->num_ops += operation->restart_ctx.ops_done;
+    }
+
+    if (ret == 0) {
+        ret = mbedtls_ecp_write_public_key(operation->key,
+                                           MBEDTLS_ECP_PF_UNCOMPRESSED, pub_key_len,
+                                           pub_key, pub_key_size);
+    }
+
+    return mbedtls_to_psa_error(ret);
+}
+
 psa_status_t mbedtls_psa_ecp_export_public_key_iop_abort(
-    mbedtls_psa_export_public_key_iop_operation_t *operation)
+    mbedtls_psa_export_public_key_iop_t *operation)
 {
     mbedtls_ecp_keypair_free(operation->key);
     mbedtls_free(operation->key);
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 261b873..c220e82 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h
@@ -123,7 +123,7 @@
 /**
  * \brief Setup a new interruptible export public-key operation.
  *
- *  \param[in] operation                 The \c mbedtls_psa_export_public_key_iop_operation_t to use.
+ *  \param[in] operation                 The \c mbedtls_psa_export_public_key_iop_t to use.
  *                                       This must be initialized first.
  *  \param[in] private_key               pointer to private key.
  *  \param[in] private_key_len           size of \p private_key in bytes.
@@ -142,21 +142,46 @@
  *
  */
 psa_status_t mbedtls_psa_ecp_export_public_key_iop_setup(
-    mbedtls_psa_export_public_key_iop_operation_t *operation,
+    mbedtls_psa_export_public_key_iop_t *operation,
     uint8_t *private_key,
     size_t private_key_len,
     const psa_key_attributes_t *private_key_attributes);
 
+
+/**
+ * \brief Continue and eventually complete an export public-key operation.
+ *
+ * \param[in] operation                  The \c mbedtls_psa_export_public_key_iop_t to use.
+ *                                       This must be initialized first and
+ *                                       had \c mbedtls_psa_ecp_export_public_key_iop_setup()
+ *                                       called successfully.
+ * \param[out] pub_key                   Buffer where the public key data is to be written.
+ * \param[in]  pub_key_size              Size of the \p pub_key buffer in bytes.
+ * \param[out] pub_key_len               On success, the number of bytes that make up the public key data.
+ *
+ * \retval #PSA_SUCCESS
+ *         The key was exported successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT \emptydescription
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL \emptydescription
+ * \retval #PSA_ERROR_NOT_SUPPORTED \emptydescription
+ *
+ */
+psa_status_t mbedtls_psa_ecp_export_public_key_iop_complete(
+    mbedtls_psa_export_public_key_iop_t *operation,
+    uint8_t *pub_key,
+    size_t pub_key_size,
+    size_t *pub_key_len);
+
 /**
  * \brief Abort an interruptible export public-key operation.
  *
- * \param[in] operation               The \c mbedtls_psa_export_public_key_iop_operation_t to abort.
+ * \param[in] operation               The \c mbedtls_psa_export_public_key_iop_t to abort.
  *
  * \retval #PSA_SUCCESS
  *         The operation was aborted successfully.
  */
 psa_status_t mbedtls_psa_ecp_export_public_key_iop_abort(
-    mbedtls_psa_export_public_key_iop_operation_t *operation);
+    mbedtls_psa_export_public_key_iop_t *operation);
 
 /**
  * \brief Generate an ECP key.
diff --git a/tf-psa-crypto/include/psa/crypto_builtin_composites.h b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
index 9bd58f9..ba5375f 100644
--- a/tf-psa-crypto/include/psa/crypto_builtin_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
@@ -258,7 +258,7 @@
     /* Make the struct non-empty if algs not supported. */
     unsigned MBEDTLS_PRIVATE(dummy);
 #endif
-} mbedtls_psa_export_public_key_iop_operation_t;
+} mbedtls_psa_export_public_key_iop_t;
 
 #if defined(MBEDTLS_ECP_C) && defined(MBEDTLS_ECP_RESTARTABLE)
 #define MBEDTLS_PSA_EXPORT_PUBLIC_KEY_IOP_INIT { NULL, MBEDTLS_ECP_RESTART_INIT, 0 }
diff --git a/tf-psa-crypto/include/psa/crypto_struct.h b/tf-psa-crypto/include/psa/crypto_struct.h
index d0300da..ffaf6c2 100644
--- a/tf-psa-crypto/include/psa/crypto_struct.h
+++ b/tf-psa-crypto/include/psa/crypto_struct.h
@@ -584,7 +584,7 @@
      * any driver (i.e. none of the driver contexts are active).
      */
     unsigned int MBEDTLS_PRIVATE(id);
-    mbedtls_psa_export_public_key_iop_operation_t MBEDTLS_PRIVATE(ctx);
+    mbedtls_psa_export_public_key_iop_t MBEDTLS_PRIVATE(ctx);
     unsigned int MBEDTLS_PRIVATE(error_occurred) : 1;
     uint32_t MBEDTLS_PRIVATE(num_ops);
 #endif
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
index ade9cba..fc8ebb5 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
@@ -156,10 +156,6 @@
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_SIGN:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_IMPORT:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_EXPORT:PSA_CRYPTO_DRIVER_TEST
 import_export:"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_KEY_USAGE_EXPORT:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION( PSA_KEY_PERSISTENCE_VOLATILE, TEST_DRIVER_LOCATION ):1024:-1:PSA_ERROR_BUFFER_TOO_SMALL:1
 
-PSA import/export RSA keypair: trailing garbage rejected, opaque
-depends_on:PSA_WANT_ALG_RSA_PKCS1V15_SIGN:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_IMPORT:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_EXPORT:PSA_CRYPTO_DRIVER_TEST
-import_with_data:"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b2400":PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_ERROR_INVALID_ARGUMENT
-
 PSA import RSA keypair: truncated
 depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_IMPORT
 import_with_data:"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b":PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_ERROR_INVALID_ARGUMENT
@@ -422,7 +418,7 @@
 
 PSA import/export-public EC brainpool512r1: good, opaque
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_IMPORT:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_EXPORT:PSA_WANT_ECC_BRAINPOOL_P_R1_512:PSA_CRYPTO_DRIVER_TEST
-import_export_public_key:"372c9778f69f726cbca3f4a268f16b4d617d10280d79a6a029cd51879fe1012934dfe5395455337df6906dc7d6d2eea4dbb2065c0228f73b3ed716480e7d71d2":PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):PSA_ALG_ECDSA_ANY:0:0:PSA_SUCCESS:"0438b7ec92b61c5c6c7fbc28a4ec759d48fcd4e2e374defd5c4968a54dbef7510e517886fbfc38ea39aa529359d70a7156c35d3cbac7ce776bdb251dd64bce71234424ee7049eed072f0dbc4d79996e175d557e263763ae97095c081e73e7db2e38adc3d4c9a0487b1ede876dc1fca61c902e9a1d8722b8612928f18a24845591a"
+import_export_public_key:"372c9778f69f726cbca3f4a268f16b4d617d10280d79a6a029cd51879fe1012934dfe5395455337df6906dc7d6d2eea4dbb2065c0228f73b3ed716480e7d71d2":PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):PSA_ALG_ECDSA_ANY:PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION( PSA_KEY_PERSISTENCE_VOLATILE, TEST_DRIVER_LOCATION ):0:PSA_SUCCESS:"0438b7ec92b61c5c6c7fbc28a4ec759d48fcd4e2e374defd5c4968a54dbef7510e517886fbfc38ea39aa529359d70a7156c35d3cbac7ce776bdb251dd64bce71234424ee7049eed072f0dbc4d79996e175d557e263763ae97095c081e73e7db2e38adc3d4c9a0487b1ede876dc1fca61c902e9a1d8722b8612928f18a24845591a"
 
 PSA import/export EC curve25519 key pair: good (already properly masked), opaque
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_IMPORT:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_EXPORT:PSA_WANT_ECC_MONTGOMERY_255:PSA_CRYPTO_DRIVER_TEST
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 00b935b..6be0f60 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
@@ -222,6 +222,9 @@
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = mbedtls_svc_key_id_make(1, 0x6964);
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_export_public_key_iop_t export_key_iop = PSA_EXPORT_PUBLIC_KEY_IOP_INIT;
+#endif
     uint8_t buffer[1];
     size_t length;
     int ok = 0;
@@ -248,6 +251,11 @@
                                      buffer, sizeof(buffer), &length),
                PSA_ERROR_INVALID_HANDLE);
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    TEST_EQUAL(psa_export_public_key_iop_setup(&export_key_iop, key),
+               PSA_ERROR_INVALID_HANDLE);
+#endif
+
     ok = 1;
 
 exit:
@@ -1886,6 +1894,8 @@
     size_t export_size = expected_public_key->len + export_size_delta;
     size_t exported_length = INVALID_EXPORT_LENGTH;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_export_public_key_iop_t export_key_operation = PSA_EXPORT_PUBLIC_KEY_IOP_INIT;
+
 
     PSA_ASSERT(psa_crypto_init());
 
@@ -1917,6 +1927,47 @@
         TEST_MEMORY_COMPARE(expected_public_key->x, expected_public_key->len,
                             exported, exported_length);
     }
+
+    /* Adjust expected_status for interruptible export public-key.
+     * Interruptible export public-key is only supported for ECC keys and even
+     * for those only when MBEDTLS_ECP_RESTARTABLE is on.
+     */
+    if ((PSA_KEY_TYPE_IS_KEY_PAIR(type) || PSA_KEY_TYPE_IS_PUBLIC_KEY(type)) &&
+        !PSA_KEY_TYPE_IS_ECC(type)) {
+        expected_export_status = PSA_ERROR_NOT_SUPPORTED;
+    }
+
+#if !defined(MBEDTLS_ECP_RESTARTABLE)
+    expected_export_status = PSA_ERROR_NOT_SUPPORTED;
+#endif
+
+    if (PSA_KEY_LIFETIME_GET_LOCATION(psa_get_key_lifetime(&attributes)) !=
+        PSA_KEY_LOCATION_LOCAL_STORAGE) {
+        expected_export_status = PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    status = psa_export_public_key_iop_setup(&export_key_operation, key);
+    TEST_EQUAL(status, expected_export_status);
+
+    if (status != PSA_SUCCESS) {
+        expected_export_status = PSA_ERROR_BAD_STATE;
+    }
+
+    memset(exported, 0, export_size);
+
+    do {
+        status = psa_export_public_key_iop_complete(&export_key_operation,
+                                                    exported,
+                                                    export_size,
+                                                    &exported_length);
+    } while (status == PSA_OPERATION_INCOMPLETE);
+    TEST_EQUAL(status, expected_export_status);
+
+    if (status == PSA_SUCCESS) {
+        TEST_MEMORY_COMPARE(expected_public_key->x, expected_public_key->len,
+                            exported, exported_length);
+    }
+
 exit:
     /*
      * Key attributes may have been returned by psa_get_key_attributes()
@@ -10390,6 +10441,11 @@
     psa_status_t expected_status = expected_status_arg;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_export_public_key_iop_t export_key_operation = PSA_EXPORT_PUBLIC_KEY_IOP_INIT;
+    uint8_t output[PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)] = { 0 };
+    size_t output_len = 0;
+    uint8_t refrence_output[PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)] =
+    { 0 };
+    size_t refrence_output_len = 0;
     psa_status_t status;
 
     PSA_ASSERT(psa_crypto_init());
@@ -10406,6 +10462,15 @@
     status = psa_generate_key(&attributes, &iop_key);
     TEST_EQUAL(status, PSA_SUCCESS);
 
+    /* Test calling complete() without calling setup() will fail. */
+    status = psa_export_public_key_iop_complete(&export_key_operation,
+                                                output,
+                                                sizeof(output),
+                                                &output_len);
+    TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+
+    PSA_ASSERT(psa_export_public_key_iop_abort(&export_key_operation));
+
     status = psa_export_public_key_iop_setup(&export_key_operation, iop_key);
     TEST_EQUAL(status, expected_status);
 
@@ -10415,12 +10480,56 @@
     TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
 #endif
 
-    TEST_EQUAL(psa_export_public_key_iop_abort(&export_key_operation), PSA_SUCCESS);
+    PSA_ASSERT(psa_export_public_key_iop_abort(&export_key_operation));
 
     /* Test that after calling abort operation is reset to it's fresh state */
     status = psa_export_public_key_iop_setup(&export_key_operation, iop_key);
     TEST_EQUAL(status, expected_status);
 
+    if (expected_status != PSA_SUCCESS) {
+        expected_status = PSA_ERROR_BAD_STATE;
+    }
+
+    do {
+        status = psa_export_public_key_iop_complete(&export_key_operation,
+                                                    output,
+                                                    sizeof(output),
+                                                    &output_len);
+    } while (status == PSA_OPERATION_INCOMPLETE);
+    TEST_EQUAL(status, expected_status);
+
+    /* Test calling complete() 2 times consecutively will fail. */
+    status = psa_export_public_key_iop_complete(&export_key_operation,
+                                                output,
+                                                sizeof(output),
+                                                &output_len);
+    TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+
+    if (expected_status == PSA_SUCCESS) {
+        status = psa_export_public_key(iop_key,
+                                       refrence_output,
+                                       sizeof(refrence_output),
+                                       &refrence_output_len);
+        TEST_EQUAL(status, PSA_SUCCESS);
+
+        TEST_MEMORY_COMPARE(refrence_output, refrence_output_len, output, output_len);
+
+        /* Test psa_export_public_key_iop_complete() returns right error code when
+           output buffer is not enough. */
+        PSA_ASSERT(psa_export_public_key_iop_abort(&export_key_operation));
+
+        status = psa_export_public_key_iop_setup(&export_key_operation, iop_key);
+        TEST_EQUAL(status, PSA_SUCCESS);
+
+        do {
+            status = psa_export_public_key_iop_complete(&export_key_operation,
+                                                        output,
+                                                        refrence_output_len-1,
+                                                        &output_len);
+        } while (status == PSA_OPERATION_INCOMPLETE);
+        TEST_EQUAL(status, PSA_ERROR_BUFFER_TOO_SMALL);
+    }
+
 exit:
     psa_export_public_key_iop_abort(&export_key_operation);
     psa_destroy_key(iop_key);