Add PSA interuptable key agreement APIs

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>

scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja

@@ -2701,6 +2701,116 @@ static inline psa_status_t psa_driver_wrapper_key_agreement(
     }
 }
 
+static inline uint32_t psa_driver_wrapper_key_agreement_get_num_ops(
+    psa_key_agreement_iop_t *operation)
+{
+    switch( operation->id )
+    {
+        /* If uninitialised, return 0, as no work can have been done. */
+        case 0:
+            return 0;
+
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            return(mbedtls_psa_key_agreement_get_num_ops(&operation->ctx.mbedtls_ctx));
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+    }
+
+    return 0;
+}
+
+static inline psa_status_t psa_driver_wrapper_key_agreement_setup(
+    psa_key_agreement_iop_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)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
+                                                    psa_get_key_lifetime(attributes) );
+
+    switch( location )
+    {
+        case PSA_KEY_LOCATION_LOCAL_STORAGE:
+            /* Key is stored in the slot in export representation, so
+             * cycle through all known transparent accelerators */
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+
+
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+
+            /* Fell through, meaning no accelerator supports this operation */
+            operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
+            status = mbedtls_psa_key_agreement_setup(&operation->ctx.mbedtls_ctx, private_key_buffer,
+                                                    private_key_buffer_len, peer_key,
+                                                    peer_key_length,
+                                                    attributes);
+            break;
+
+            /* Add cases for opaque driver here */
+
+            default:
+                /* Key is declared with a lifetime not known to us */
+                status = PSA_ERROR_INVALID_ARGUMENT;
+                break;
+    }
+    return( status );
+
+}
+
+static inline psa_status_t psa_driver_wrapper_key_agreement_complete(
+    psa_key_agreement_iop_t *operation,
+    uint8_t *shared_secret,
+    size_t shared_secret_size,
+    size_t *shared_secret_length)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    switch( operation->id )
+    {
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            status = mbedtls_psa_key_agreement_complete(&operation->ctx.mbedtls_ctx, shared_secret,
+                                                        shared_secret_size,
+                                                        shared_secret_length);
+            break;
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+
+
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            break;
+    }
+    return( status );
+}
+
+static inline psa_status_t psa_driver_wrapper_key_agreement_abort(
+    psa_key_agreement_iop_t *operation)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    switch( operation->id )
+    {
+        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
+            status = mbedtls_psa_key_agreement_abort(&operation->ctx.mbedtls_ctx);
+            break;
+
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+
+
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            break;
+    }
+    return( status );
+}
+
+
 static inline psa_status_t psa_driver_wrapper_pake_setup(
     psa_pake_operation_t *operation,
     const psa_crypto_driver_pake_inputs_t *inputs )

tests/include/test/psa_test_wrappers.h

@@ -475,13 +475,13 @@ psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_abort(
 
 psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_complete(
     psa_key_agreement_iop_t *arg0_operation,
-    psa_key_id_t *arg1_key);
+    mbedtls_svc_key_id_t *arg1_key);
 #define psa_key_agreement_iop_complete(arg0_operation, arg1_key) \
     mbedtls_test_wrap_psa_key_agreement_iop_complete(arg0_operation, arg1_key)
 
 psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_setup(
     psa_key_agreement_iop_t *arg0_operation,
-    psa_key_id_t arg1_private_key,
+    mbedtls_svc_key_id_t arg1_private_key,
     const uint8_t *arg2_peer_key,
     size_t arg3_peer_key_length,
     psa_algorithm_t arg4_alg,

tests/psa-client-server/psasim/src/psa_sim_serialise.c

@@ -1624,6 +1624,42 @@ int psasim_deserialise_mbedtls_svc_key_id_t(uint8_t **pos,
     return 1;
 }
 
+size_t psasim_serialise_psa_key_agreement_iop_t_needs(
+    psa_key_agreement_iop_t value)
+{
+    return sizeof(value);
+}
+
+int psasim_serialise_psa_key_agreement_iop_t(uint8_t **pos,
+                                             size_t *remaining,
+                                             psa_key_agreement_iop_t value)
+{
+    if (*remaining < sizeof(value)) {
+        return 0;
+    }
+
+    memcpy(*pos, &value, sizeof(value));
+    *pos += sizeof(value);
+
+    return 1;
+}
+
+int psasim_deserialise_psa_key_agreement_iop_t(uint8_t **pos,
+                                               size_t *remaining,
+                                               psa_key_agreement_iop_t *value)
+{
+    if (*remaining < sizeof(*value)) {
+        return 0;
+    }
+
+    memcpy(value, *pos, sizeof(*value));
+
+    *pos += sizeof(*value);
+    *remaining -= sizeof(*value);
+
+    return 1;
+}
+
 void psa_sim_serialize_reset(void)
 {
     memset(hash_operation_handles, 0,

tests/psa-client-server/psasim/src/psa_sim_serialise.h

@@ -1301,3 +1301,46 @@ int psasim_serialise_mbedtls_svc_key_id_t(uint8_t **pos,
 int psasim_deserialise_mbedtls_svc_key_id_t(uint8_t **pos,
                                             size_t *remaining,
                                             mbedtls_svc_key_id_t *value);
+
+/** Return how much buffer space is needed by \c psasim_serialise_psa_key_agreement_iop_t()
+ *  to serialise a `psa_key_agreement_iop_t`.
+ *
+ * \param value              The value that will be serialised into the buffer
+ *                           (needed in case some serialisations are value-
+ *                           dependent).
+ *
+ * \return                   The number of bytes needed in the buffer by
+ *                           \c psasim_serialise_psa_key_agreement_iop_t() to serialise
+ *                           the given value.
+ */
+size_t psasim_serialise_psa_key_agreement_iop_t_needs(
+    psa_key_agreement_iop_t value);
+
+/** Serialise a `psa_key_agreement_iop_t` into a buffer.
+ *
+ * \param pos[in,out]        Pointer to a `uint8_t *` holding current position
+ *                           in the buffer.
+ * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
+ *                           remaining in the buffer.
+ * \param value              The value to serialise into the buffer.
+ *
+ * \return                   \c 1 on success ("okay"), \c 0 on error.
+ */
+int psasim_serialise_psa_key_agreement_iop_t(uint8_t **pos,
+                                             size_t *remaining,
+                                             psa_key_agreement_iop_t value);
+
+/** Deserialise a `psa_key_agreement_iop_t` from a buffer.
+ *
+ * \param pos[in,out]        Pointer to a `uint8_t *` holding current position
+ *                           in the buffer.
+ * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
+ *                           remaining in the buffer.
+ * \param value              Pointer to a `psa_key_agreement_iop_t` to receive the value
+ *                           deserialised from the buffer.
+ *
+ * \return                   \c 1 on success ("okay"), \c 0 on error.
+ */
+int psasim_deserialise_psa_key_agreement_iop_t(uint8_t **pos,
+                                               size_t *remaining,
+                                               psa_key_agreement_iop_t *value);

tests/psa-client-server/psasim/src/psa_sim_serialise.pl

@@ -48,7 +48,8 @@ my @types = qw(unsigned-int int size_t
                psa_key_derivation_operation_t
                psa_sign_hash_interruptible_operation_t
                psa_verify_hash_interruptible_operation_t
-               mbedtls_svc_key_id_t);
+               mbedtls_svc_key_id_t
+               psa_key_agreement_iop_t);
 
 grep(s/-/ /g, @types);
 

tests/src/psa_test_wrappers.c

@@ -835,7 +835,7 @@ psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_abort(
 /* Wrapper for psa_key_agreement_iop_complete */
 psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_complete(
     psa_key_agreement_iop_t *arg0_operation,
-    psa_key_id_t *arg1_key)
+    mbedtls_svc_key_id_t *arg1_key)
 {
     psa_status_t status = (psa_key_agreement_iop_complete)(arg0_operation, arg1_key);
     return status;
@@ -844,7 +844,7 @@ psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_complete(
 /* Wrapper for psa_key_agreement_iop_setup */
 psa_status_t mbedtls_test_wrap_psa_key_agreement_iop_setup(
     psa_key_agreement_iop_t *arg0_operation,
-    psa_key_id_t arg1_private_key,
+    mbedtls_svc_key_id_t arg1_private_key,
     const uint8_t *arg2_peer_key,
     size_t arg3_peer_key_length,
     psa_algorithm_t arg4_alg,

tf-psa-crypto/core/psa_crypto.c

@@ -7764,6 +7764,128 @@ psa_status_t psa_key_agreement(mbedtls_svc_key_id_t private_key,
     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 */
 /****************************************************************/

tf-psa-crypto/core/psa_crypto_core.h

@@ -700,6 +700,110 @@ psa_status_t psa_key_agreement_raw_builtin(
     size_t shared_secret_size,
     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.

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 @@ psa_status_t mbedtls_psa_generate_key_iop_abort(
 /* Interruptible ECC Key Agreement */
 /****************************************************************/
 
-uint32_t psa_key_agreement_iop_get_num_ops(
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
-    psa_key_agreement_iop_t *operation)
+
+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_status_t mbedtls_psa_key_agreement_setup(
-    psa_key_agreement_iop_t *operation,
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation,
-    psa_key_id_t private_key,
+    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_ERROR_NOT_SUPPORTED;
-    return PSA_SUCCESS;
+#endif
 }
 
-psa_status_t psa_key_agreement_iop_complete(
+psa_status_t mbedtls_psa_key_agreement_complete(
-    psa_key_agreement_iop_t *operation,
+    mbedtls_psa_key_agreement_interruptible_operation_t *operation,
-    psa_key_id_t *key)
+    uint8_t *shared_secret,
+    size_t shared_secret_size,
+    size_t *shared_secret_length)
 {
-    (void) operation;
+#if defined(MBEDTLS_ECP_RESTARTABLE)
-    (void) key;
+    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_status_t mbedtls_psa_key_agreement_abort(
-    psa_key_agreement_iop_t *operation)
+    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 */

tf-psa-crypto/include/psa/crypto.h

@@ -5123,7 +5123,7 @@ uint32_t psa_key_agreement_iop_get_num_ops(psa_key_agreement_iop_t *operation);
 
 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_setup(
  */
 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.

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 @@ typedef struct {
 #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 */

tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h

@@ -147,5 +147,10 @@ typedef union {
 #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. */

tf-psa-crypto/include/psa/crypto_struct.h

@@ -508,14 +508,15 @@ struct psa_key_agreement_iop_s {
      * 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