Repository: Mbed-TLS/mbedtls (C)

Issue:
Implement interruptible ECC Export Public Key (Setup / Abort)
Stage 2 of implementing interruptible ECC export public key.

This should involve implementing `psa_export_public_key_iop_setup()` and `psa_export_public_key_iop_abort()` and include any testing that we can do of the above two functions to ensure they are functionally correct.

Remember that abort should be called internally in case of any error.

Function prototypes for reference: 

```
psa_status_t psa_export_public_key_iop_abort(psa_export_public_key_iop_t * operation);
psa_status_t psa_export_public_key_iop_setup(psa_export_public_key_iop_t * operation,
                                             psa_key_id_t key);

```



Fix (patch):
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 4fb6ab5097e1..beb17d5276b5 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -1665,6 +1665,23 @@ psa_status_t psa_export_public_key(mbedtls_svc_key_id_t key,
 /* Interruptible ECC Export Public-key */
 /****************************************************************/
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+static psa_status_t psa_export_public_key_iop_abort_internal(psa_export_public_key_iop_t *operation)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if (operation->id == 0) {
+        return PSA_SUCCESS;
+    }
+
+    status = mbedtls_psa_ecp_export_public_key_iop_abort(&operation->ctx);
+
+    operation->id = 0;
+
+    return status;
+}
+#endif
+
 uint32_t psa_export_public_key_iop_get_num_ops(psa_export_public_key_iop_t *operation)
 {
     (void) operation;
@@ -1672,12 +1689,67 @@ 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;
+
+    private_key_type = psa_get_key_type(&private_key_attributes);
+
+    if (!PSA_KEY_TYPE_IS_KEY_PAIR(private_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) {
+        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);
+
+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,
@@ -1695,9 +1767,19 @@ psa_status_t psa_export_public_key_iop_complete(psa_export_public_key_iop_t *ope
 
 psa_status_t psa_export_public_key_iop_abort(psa_export_public_key_iop_t *operation)
 {
-    (void) operation;
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    return PSA_ERROR_NOT_SUPPORTED;
+    status = psa_export_public_key_iop_abort_internal(operation);
+
+    operation->num_ops = 0;
+    operation->error_occurred = 0;
+
+    return status;
+#else
+    (void) operation;
+    return PSA_SUCCESS;
+#endif
 }
 
 /** Validate that a key policy is internally well-formed.
@@ -8443,7 +8525,6 @@ psa_status_t psa_generate_key_iop_abort(
 #endif
 }
 
-
 /****************************************************************/
 /* Module setup */
 /****************************************************************/
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 b2764b0649a3..b58587f7280c 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -654,6 +654,41 @@ psa_status_t mbedtls_psa_ecp_generate_key_iop_abort(
     return PSA_SUCCESS;
 }
 
+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)
+{
+    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,
+        &operation->key);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    mbedtls_ecp_restart_init(&operation->restart_ctx);
+    operation->num_ops = 0;
+
+exit:
+    return status;
+}
+
+psa_status_t mbedtls_psa_ecp_export_public_key_iop_abort(
+    mbedtls_psa_export_public_key_iop_operation_t *operation)
+{
+    mbedtls_ecp_keypair_free(operation->key);
+    mbedtls_free(operation->key);
+    mbedtls_ecp_restart_free(&operation->restart_ctx);
+    operation->num_ops = 0;
+    return PSA_SUCCESS;
+}
+
 #endif
 /****************************************************************/
 /* Interruptible ECC Key Agreement */
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 f3ff32328bb6..506516de7624 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h
@@ -120,6 +120,44 @@ psa_status_t mbedtls_psa_ecp_export_public_key(
     const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length);
 
+/**
+ * \brief Setup a new interruptible export public-key operation.
+ *
+ *  \param[in] operation                 The \c mbedtls_psa_export_public_key_iop_operation_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.
+ *  \param[in] private_key_attributes    Key attributes of the private key.
+ *
+ *  \retval #PSA_SUCCESS
+ *         The operation started successfully - call \c mbedtls_psa_ecp_export_public_key_iop_complete()
+ *         with the same operation to complete the operation.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *           Either no internal interruptible operations are
+ *           currently supported, or the key attributes are not unsupported.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ *         There was insufficient memory to load the key representation.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED \emptydescription
+ *
+ */
+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);
+
+/**
+ * \brief Abort an interruptible export public-key operation.
+ *
+ * \param[in] operation               The \c mbedtls_psa_export_public_key_iop_operation_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);
+
 /**
  * \brief Generate an ECP key.
  *
diff --git a/tf-psa-crypto/include/psa/crypto.h b/tf-psa-crypto/include/psa/crypto.h
index c1ef041f7769..6ea945d3ec75 100644
--- a/tf-psa-crypto/include/psa/crypto.h
+++ b/tf-psa-crypto/include/psa/crypto.h
@@ -5677,7 +5677,7 @@ uint32_t psa_export_public_key_iop_get_num_ops(psa_export_public_key_iop_t *oper
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY \emptydescription
  */
 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);
 
 /**
  * \brief                       Continue and eventually complete the action of
diff --git a/tf-psa-crypto/include/psa/crypto_builtin_composites.h b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
index 14e759b13d86..9bd58f9d9e8c 100644
--- a/tf-psa-crypto/include/psa/crypto_builtin_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
@@ -248,4 +248,22 @@ typedef struct {
 #define MBEDTLS_PSA_KEY_AGREEMENT_IOP_INIT { 0 }
 #endif
 
+/* Context structure for the Mbed TLS interruptible export public-key implementation. */
+typedef struct {
+#if defined(MBEDTLS_ECP_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    mbedtls_ecp_keypair *MBEDTLS_PRIVATE(key);
+    mbedtls_ecp_restart_ctx MBEDTLS_PRIVATE(restart_ctx);
+    uint32_t MBEDTLS_PRIVATE(num_ops);
+#else
+    /* Make the struct non-empty if algs not supported. */
+    unsigned MBEDTLS_PRIVATE(dummy);
+#endif
+} mbedtls_psa_export_public_key_iop_operation_t;
+
+#if defined(MBEDTLS_ECP_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+#define MBEDTLS_PSA_EXPORT_PUBLIC_KEY_IOP_INIT { NULL, MBEDTLS_ECP_RESTART_INIT, 0 }
+#else
+#define MBEDTLS_PSA_EXPORT_PUBLIC_KEY_IOP_INIT { 0 }
+#endif
+
 #endif /* PSA_CRYPTO_BUILTIN_COMPOSITES_H */
diff --git a/tf-psa-crypto/include/psa/crypto_struct.h b/tf-psa-crypto/include/psa/crypto_struct.h
index 09f4c181752b..d0300da00473 100644
--- a/tf-psa-crypto/include/psa/crypto_struct.h
+++ b/tf-psa-crypto/include/psa/crypto_struct.h
@@ -584,15 +584,20 @@ struct psa_export_public_key_iop_s {
      * 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);
+    unsigned int MBEDTLS_PRIVATE(error_occurred) : 1;
+    uint32_t MBEDTLS_PRIVATE(num_ops);
 #endif
 };
+
 #if defined(MBEDTLS_PSA_CRYPTO_CLIENT) && !defined(MBEDTLS_PSA_CRYPTO_C)
 #define PSA_EXPORT_PUBLIC_KEY_IOP_INIT { 0 }
 #else
-#define PSA_EXPORT_PUBLIC_KEY_IOP_INIT { 0 }
+#define PSA_EXPORT_PUBLIC_KEY_IOP_INIT { 0, MBEDTLS_PSA_EXPORT_PUBLIC_KEY_IOP_INIT, 0, 0 }
 #endif
 
-static inline struct psa_export_public_key_iop_s psa_export_public_key_iop_init(void)
+static inline struct psa_export_public_key_iop_s
+psa_export_public_key_iop_init(void)
 {
     const struct psa_export_public_key_iop_s v = PSA_EXPORT_PUBLIC_KEY_IOP_INIT;