diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 9fd908dce2..310dab8a22 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -1689,12 +1689,62 @@ uint32_t psa_export_public_key_iop_get_num_ops(psa_export_public_key_iop_t *oper
 }
 
 psa_status_t psa_export_public_key_iop_setup(psa_export_public_key_iop_t *operation,
-                                             psa_key_id_t key)
+                                             mbedtls_svc_key_id_t key)
 {
+#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_slot_t *slot = NULL;
+
+    if (operation->id != 0 || operation->error_occurred) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    /* We only support the builtin/Mbed TLS driver for now. */
+    operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
+
+    status = psa_get_and_lock_transparent_key_slot_with_policy(key, &slot,
+                                                               0,
+                                                               0);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    private_key_attributes = slot->attr;
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    private_key_type = psa_get_key_type(&private_key_attributes);
+    if (!PSA_KEY_TYPE_IS_ECC_KEY_PAIR(private_key_type)) {
+        goto exit;
+    }
+
+    key_size = PSA_EXPORT_KEY_OUTPUT_SIZE(private_key_type,
+                                          psa_get_key_bits(&private_key_attributes));
+    if (key_size == 0) {
+        goto exit;
+    }
+
+    status = mbedtls_psa_ecp_export_public_key_iop_setup(&operation->ctx, slot->key.data,
+                                                         slot->key.bytes, &private_key_attributes);
+
+exit:
+    unlock_status = psa_unregister_read_under_mutex(slot);
+    if (status != PSA_SUCCESS) {
+        psa_export_public_key_iop_abort_internal(operation);
+        operation->error_occurred = 1;
+        return status;
+    }
+    return unlock_status;
+#else
     (void) operation;
     (void) key;
-
     return PSA_ERROR_NOT_SUPPORTED;
+#endif
 }
 
 psa_status_t psa_export_public_key_iop_complete(psa_export_public_key_iop_t *operation,