Add PSA interuptable key agreement APIs

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 5f44cc3..840bb6d 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -7764,6 +7764,128 @@
     return status;
 }
 
+uint32_t psa_key_agreement_iop_get_num_ops(
+    psa_key_agreement_iop_t *operation)
+{
+#if defined(MBEDTLS_ECP_RESTARTABLE) && \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+    return psa_driver_wrapper_key_agreement_get_num_ops(operation);
+#else
+    (void) operation;
+    return 0;
+#endif
+}
+
+psa_status_t psa_key_agreement_iop_setup(
+    psa_key_agreement_iop_t *operation,
+    mbedtls_svc_key_id_t private_key,
+    const uint8_t *peer_key,
+    size_t peer_key_length,
+    psa_algorithm_t alg,
+    const psa_key_attributes_t *attributes)
+{
+#if defined(MBEDTLS_ECP_RESTARTABLE) && \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+    psa_status_t status = PSA_SUCCESS;
+    uint8_t *private_key_buffer = NULL;
+    size_t key_size = 0;
+    size_t key_len = 0;
+    psa_key_attributes_t private_key_attributes;
+
+    status = psa_get_key_attributes(private_key, &private_key_attributes);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    if (!PSA_KEY_TYPE_IS_ECC_KEY_PAIR(private_key_attributes.type) ||
+        !PSA_ALG_IS_ECDH(alg)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    key_size = PSA_BITS_TO_BYTES(psa_get_key_bits(&private_key_attributes));
+    if (key_size == 0) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+    private_key_buffer = mbedtls_calloc(key_size, 1);
+    if (private_key_buffer == NULL) {
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    }
+
+    status = psa_export_key(private_key, private_key_buffer, key_size, &key_len);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    operation->ctx.mbedtls_ctx.attributes = attributes;
+
+    status = psa_driver_wrapper_key_agreement_setup(operation, private_key_buffer,
+                                                    key_len, peer_key,
+                                                    peer_key_length,
+                                                    &private_key_attributes);
+
+exit:
+    mbedtls_free(private_key_buffer);
+    return status;
+    #else
+    (void) operation;
+    (void) private_key;
+    (void) peer_key;
+    (void) peer_key_length;
+    (void) alg;
+    (void) attributes;
+    return PSA_ERROR_NOT_SUPPORTED;
+    #endif
+}
+
+psa_status_t psa_key_agreement_iop_complete(
+    psa_key_agreement_iop_t *operation,
+    mbedtls_svc_key_id_t *key)
+{
+#if defined(MBEDTLS_ECP_RESTARTABLE) && \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+
+    psa_status_t status;
+    uint8_t intermediate_key[PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE];
+    size_t key_len = 0;
+
+    status = psa_driver_wrapper_key_agreement_complete(operation, intermediate_key,
+                                                       PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE,
+                                                       &key_len);
+
+    operation->num_ops = psa_driver_wrapper_key_agreement_get_num_ops(operation);
+
+    if (status == PSA_SUCCESS) {
+        status = psa_import_key(operation->ctx.mbedtls_ctx.attributes, intermediate_key,
+                                key_len, key);
+    }
+
+    mbedtls_platform_zeroize(intermediate_key, PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE);
+    return status;
+#else
+    (void) operation;
+    (void) key;
+    return PSA_ERROR_NOT_SUPPORTED;
+#endif
+}
+
+psa_status_t psa_key_agreement_iop_abort(
+    psa_key_agreement_iop_t *operation)
+{
+#if defined(MBEDTLS_ECP_RESTARTABLE) && \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+    psa_status_t status;
+
+    status = psa_driver_wrapper_key_agreement_abort(operation);
+
+    return status;
+#else
+    (void) operation;
+    return PSA_ERROR_NOT_SUPPORTED;
+#endif
+
+
+}
+
 /****************************************************************/
 /* Random generation */
 /****************************************************************/
diff --git a/tf-psa-crypto/core/psa_crypto_core.h b/tf-psa-crypto/core/psa_crypto_core.h
index 1753554..33950e0 100644
--- a/tf-psa-crypto/core/psa_crypto_core.h
+++ b/tf-psa-crypto/core/psa_crypto_core.h
@@ -701,6 +701,110 @@
     size_t *shared_secret_length);
 
 /**
+ * \brief Get the total number of ops that a key agreement operation has taken
+ *        Since it's start.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       key_agreement_get_num_ops entry point. This function behaves as an
+ *       key_agreement_get_num_ops entry point as defined in the PSA driver
+ *       interface specification for transparent drivers.
+ *
+ * \param[in]   operation           The \c mbedtls_psa_key_agreement_interruptible_operation_t to use.
+ *                                  This must be initialized first.
+ *
+ * \return                      Total number of operations.
+ */
+uint32_t mbedtls_psa_key_agreement_get_num_ops(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation);
+
+/**
+ * \brief  Setup a new interruptible key agreement operation.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       key_agreement_setup entry point. This function behaves as a
+ *       key_agreement_setup entry point as defined in the PSA driver interface
+ *       specification for transparent drivers.
+ *
+ *  \param[in] operation                 The \c psa_key_agreement_iop_t to use.
+ *                                       This must be initialized first.
+ *  \param[in] private_key_buffer        The buffer containing the private key
+ *                                       context.
+ *  \param[in] private_key_buffer_len    Size of the \p private_key_buffer buffer in
+ *                                       bytes.
+ *  \param[in] peer_key                  The buffer containing the key context
+ *                                       of the peer's public key.
+ *  \param[in]  peer_key_length          Size of the \p peer_key buffer in
+ *                                       bytes.
+ *  \param[in]  attributes               The attributes of the private key to use for the
+ *                                       operation.
+ *  \retval #PSA_SUCCESS
+ *         The operation started successfully - call \c psa_key_agreement_complete()
+ *         with the same context to complete the operation
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         An unsupported, incorrectly formatted or incorrect type of key was
+ *         used.
+ * \retval #PSA_ERROR_NOT_SUPPORTED Either no internal interruptible operations
+ *         are currently supported, or the key type is currently unsupported.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ *         There was insufficient memory to load the key representation.
+ */
+psa_status_t mbedtls_psa_key_agreement_setup(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation,
+    const uint8_t *private_key_buffer,
+    size_t private_key_buffer_len,
+    const uint8_t *peer_key,
+    size_t peer_key_length,
+    const psa_key_attributes_t *attributes);
+
+/**
+ * \brief Continue and eventually complete a key agreement operation.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       key_agreement_complete entry point. This function behaves as a
+ *       key_agreement_complete entry point as defined in the PSA driver
+ *       interface specification for transparent drivers.
+ *
+ * \param[in] operation                  The \c mbedtls_psa_key_agreement_interruptible_operation_t to use.
+ *                                       This must be initialized first.
+ * \param[out] shared_secret             The buffer to which the shared secret
+ *                                       is to be written.
+ * \param[in]  shared_secret_size        Size of the \p shared_secret buffer in
+ *                                       bytes.
+ * \param[out] shared_secret_length      On success, the number of bytes that make
+ *                                       up the returned shared secret.
+ * \retval #PSA_SUCCESS
+ *         The shared secret was calculated successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT \emptydescription
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         Internal interruptible operations are currently supported.
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ *         \p shared_secret_size is too small
+ */
+psa_status_t mbedtls_psa_key_agreement_complete(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation,
+    uint8_t *shared_secret,
+    size_t shared_secret_size,
+    size_t *shared_secret_length);
+
+/**
+ * \brief Abort a key agreement operation.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       key_agreement_abort entry point. This function behaves as a
+ *       key_agreement_abort entry point as defined in the PSA driver
+ *       interface specification for transparent drivers.
+ *
+ * \param[in] operation                  The \c mbedtls_psa_key_agreement_interruptible_operation_t to abort.
+ *                                       This must be initialized first.
+ *
+ * \retval #PSA_SUCCESS
+ *         The operation was aborted successfully.
+ */
+psa_status_t mbedtls_psa_key_agreement_abort(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation);
+
+
+/**
  * \brief Set the maximum number of ops allowed to be executed by an
  *        interruptible function in a single call.
  *
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..77c8ac2 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -12,6 +12,7 @@
 
 #include <psa/crypto.h>
 #include "psa_crypto_core.h"
+#include "psa_crypto_driver_wrappers.h"
 #include "psa_crypto_ecp.h"
 #include "psa_crypto_random_impl.h"
 #include "mbedtls/psa_util.h"
@@ -630,47 +631,133 @@
 /* Interruptible ECC Key Agreement */
 /****************************************************************/
 
-uint32_t psa_key_agreement_iop_get_num_ops(
-    psa_key_agreement_iop_t *operation)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+
+uint32_t mbedtls_psa_key_agreement_get_num_ops(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation)
 {
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    return operation->num_ops;
+#else
     (void) operation;
     return 0;
+#endif
 }
 
-psa_status_t psa_key_agreement_iop_setup(
-    psa_key_agreement_iop_t *operation,
-    psa_key_id_t private_key,
+psa_status_t mbedtls_psa_key_agreement_setup(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation,
+    const uint8_t *private_key_buffer,
+    size_t private_key_buffer_len,
     const uint8_t *peer_key,
     size_t peer_key_length,
-    psa_algorithm_t alg,
     const psa_key_attributes_t *attributes)
 {
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    mbedtls_ecp_keypair *ecp = NULL;
+    mbedtls_ecp_keypair *their_key = NULL;
+    size_t bits = 0;
+
+    status = mbedtls_psa_ecp_load_representation(
+        attributes->type,
+        attributes->bits,
+        private_key_buffer,
+        private_key_buffer_len,
+        &ecp);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    psa_ecc_family_t curve = mbedtls_ecc_group_to_psa(ecp->grp.id, &bits);
+
+    mbedtls_ecdh_init(&operation->ctx);
+
+    status = mbedtls_psa_ecp_load_representation(
+        PSA_KEY_TYPE_ECC_PUBLIC_KEY(curve),
+        bits,
+        peer_key,
+        peer_key_length,
+        &their_key);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    status = mbedtls_to_psa_error(
+        mbedtls_ecdh_get_params(&operation->ctx, their_key, MBEDTLS_ECDH_THEIRS));
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    status = mbedtls_to_psa_error(
+        mbedtls_ecdh_get_params(&operation->ctx, ecp, MBEDTLS_ECDH_OURS));
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    mbedtls_ecdh_enable_restart(&operation->ctx);
+    operation->num_ops = 0;
+
+exit:
+    mbedtls_ecp_keypair_free(their_key);
+    mbedtls_ecp_keypair_free(ecp);
+    mbedtls_free(ecp);
+    mbedtls_free(their_key);
+    return status;
+#else
     (void) operation;
-    (void) private_key;
+    (void) private_key_buffer;
+    (void) private_key_buffer;
+    (void) private_key_buffer_len;
     (void) peer_key;
     (void) peer_key_length;
-    (void) alg;
     (void) attributes;
-
-    return PSA_SUCCESS;
+    return PSA_ERROR_NOT_SUPPORTED;
+#endif
 }
 
-psa_status_t psa_key_agreement_iop_complete(
-    psa_key_agreement_iop_t *operation,
-    psa_key_id_t *key)
+psa_status_t mbedtls_psa_key_agreement_complete(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation,
+    uint8_t *shared_secret,
+    size_t shared_secret_size,
+    size_t *shared_secret_length)
 {
-    (void) operation;
-    (void) key;
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    return PSA_SUCCESS;
+    mbedtls_psa_interruptible_set_max_ops(psa_interruptible_get_max_ops());
+
+    status = mbedtls_to_psa_error(mbedtls_ecdh_calc_secret(&operation->ctx, shared_secret_length,
+                                                           shared_secret,
+                                                           shared_secret_size,
+                                                           mbedtls_psa_get_random,
+                                                           MBEDTLS_PSA_RANDOM_STATE));
+
+    operation->num_ops += operation->ctx.rs.ops_done;
+
+    return status;
+#else
+    (void) operation;
+    (void) shared_secret;
+    (void) shared_secret_size;
+    (void) shared_secret_length;
+    return PSA_ERROR_NOT_SUPPORTED;
+#endif
 }
 
-psa_status_t psa_key_agreement_iop_abort(
-    psa_key_agreement_iop_t *operation)
+psa_status_t mbedtls_psa_key_agreement_abort(
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation)
 {
-    (void) operation;
-
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    operation->num_ops = 0;
+    operation->attributes = NULL;
+    mbedtls_ecdh_free(&operation->ctx);
     return PSA_SUCCESS;
+#else
+    (void) operation;
+    return PSA_ERROR_NOT_SUPPORTED;
+#endif
 }
 
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_ECDH */
+
 #endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/tf-psa-crypto/include/psa/crypto.h b/tf-psa-crypto/include/psa/crypto.h
index 58b6887..62f1dca 100644
--- a/tf-psa-crypto/include/psa/crypto.h
+++ b/tf-psa-crypto/include/psa/crypto.h
@@ -5123,7 +5123,7 @@
 
 psa_status_t psa_key_agreement_iop_setup(
     psa_key_agreement_iop_t *operation,
-    psa_key_id_t private_key,
+    mbedtls_svc_key_id_t private_key,
     const uint8_t *peer_key,
     size_t peer_key_length,
     psa_algorithm_t alg,
@@ -5208,7 +5208,7 @@
  */
 psa_status_t psa_key_agreement_iop_complete(
     psa_key_agreement_iop_t *operation,
-    psa_key_id_t *key);
+    mbedtls_svc_key_id_t *key);
 
 /**
  * \brief                       Abort a key agreement operation.
diff --git a/tf-psa-crypto/include/psa/crypto_builtin_composites.h b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
index c9c0c6b..c5dde0d 100644
--- a/tf-psa-crypto/include/psa/crypto_builtin_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
@@ -33,6 +33,10 @@
 #endif
 #include "mbedtls/chachapoly.h"
 
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+#include "mbedtls/ecdh.h"
+#endif
+
 /*
  * MAC multi-part operation definitions.
  */
@@ -227,4 +231,22 @@
 #define MBEDTLS_PSA_GENERATE_KEY_IOP_INIT { 0 }
 #endif
 
+/* Context structure for the Mbed TLS interruptible key agreement implementation. */
+typedef struct {
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) && defined(MBEDTLS_ECP_RESTARTABLE)
+    mbedtls_ecdh_context MBEDTLS_PRIVATE(ctx);
+    const psa_key_attributes_t *MBEDTLS_PRIVATE(attributes);
+    uint32_t MBEDTLS_PRIVATE(num_ops);
+#else
+    /* Make the struct non-empty if algs not supported. */
+    unsigned MBEDTLS_PRIVATE(dummy);
+#endif
+} mbedtls_psa_key_agreement_interruptible_operation_t;
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) && defined(MBEDTLS_ECP_RESTARTABLE)
+#define MBEDTLS_PSA_KEY_AGREEMENT_INTERRUPTIBLE_OPERATION_INIT { { 0 }, { 0 }, 0 }
+#else
+#define MBEDTLS_PSA_KEY_AGREEMENT_INTERRUPTIBLE_OPERATION_INIT { 0 }
+#endif
+
 #endif /* PSA_CRYPTO_BUILTIN_COMPOSITES_H */
diff --git a/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h b/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
index 5a484fc..086c0c2 100644
--- a/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
@@ -147,5 +147,10 @@
 #endif
 } psa_driver_pake_context_t;
 
+typedef union {
+    unsigned dummy; /* Make sure this union is always non-empty */
+    mbedtls_psa_key_agreement_interruptible_operation_t mbedtls_ctx;
+} psa_driver_key_agreement_interruptible_context_t;
+
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_COMPOSITES_H */
 /* End of automatically generated file. */
diff --git a/tf-psa-crypto/include/psa/crypto_struct.h b/tf-psa-crypto/include/psa/crypto_struct.h
index 76ef5c4..b300557 100644
--- a/tf-psa-crypto/include/psa/crypto_struct.h
+++ b/tf-psa-crypto/include/psa/crypto_struct.h
@@ -508,14 +508,15 @@
      * any driver (i.e. none of the driver contexts are active).
      */
     unsigned int MBEDTLS_PRIVATE(id);
-
+    psa_driver_key_agreement_interruptible_context_t MBEDTLS_PRIVATE(ctx);
+    uint32_t MBEDTLS_PRIVATE(num_ops);
 #endif
 };
 
 #if defined(MBEDTLS_PSA_CRYPTO_CLIENT) && !defined(MBEDTLS_PSA_CRYPTO_C)
 #define PSA_KEY_AGREEMENT_IOP_INIT { 0 }
 #else
-#define PSA_KEY_AGREEMENT_IOP_INIT { 0 }
+#define PSA_KEY_AGREEMENT_IOP_INIT { 0, { 0 }, 0 }
 #endif
 
 static inline struct psa_key_agreement_iop_s