Remove interuptible key agreement driver interface

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
index 9d10f61..d3b7d6f 100644
--- a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
+++ b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
@@ -2701,121 +2701,6 @@
     }
 }
 
-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)
-
-    /* Add cases for drivers here */
-
-#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)
-
-    /* Add cases for drivers here */
-
-#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)
-
-    /* Add cases for drivers here */
-
-#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)
-
-    /* Add cases for drivers here */
-
-#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 )
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 9accc7f..9d9080f 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -7771,7 +7771,13 @@
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    status = psa_driver_wrapper_key_agreement_abort(operation);
+    if (operation->id == 0) {
+        return PSA_SUCCESS;
+    }
+
+    status = mbedtls_psa_key_agreement_iop_abort(&operation->mbedtls_ctx);
+
+    operation->id = 0;
 
     return status;
 }
@@ -7830,12 +7836,15 @@
 
     operation->num_ops = 0;
 
-    status = psa_driver_wrapper_key_agreement_setup(operation, slot->key.data,
-                                                    slot->key.bytes, peer_key,
-                                                    peer_key_length,
-                                                    &slot->attr);
+    /* To be removed later when driver dispatch is added. */
+    operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
 
-    operation->num_ops = psa_driver_wrapper_key_agreement_get_num_ops(operation);
+    status = mbedtls_psa_key_agreement_iop_setup(&operation->mbedtls_ctx,
+                                                 &slot->attr, slot->key.data,
+                                                 slot->key.bytes, peer_key,
+                                                 peer_key_length);
+
+    operation->num_ops = mbedtls_psa_key_agreement_iop_get_num_ops(&operation->mbedtls_ctx);
 
 exit:
     unlock_status = psa_unregister_read_under_mutex(slot);
@@ -7871,11 +7880,11 @@
     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);
+    status = mbedtls_psa_key_agreement_iop_complete(&operation->mbedtls_ctx, intermediate_key,
+                                                    sizeof(intermediate_key),
+                                                    &key_len);
 
-    operation->num_ops = psa_driver_wrapper_key_agreement_get_num_ops(operation);
+    operation->num_ops = mbedtls_psa_key_agreement_iop_get_num_ops(&operation->mbedtls_ctx);
 
     if (status == PSA_SUCCESS) {
         status = psa_import_key(&operation->attributes, intermediate_key,
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 086c0c2..5a484fc 100644
--- a/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
@@ -147,10 +147,5 @@
 #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 f5241b8..7dbeadc 100644
--- a/tf-psa-crypto/include/psa/crypto_struct.h
+++ b/tf-psa-crypto/include/psa/crypto_struct.h
@@ -508,7 +508,7 @@
      * 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);
+    mbedtls_psa_key_agreement_interruptible_operation_t MBEDTLS_PRIVATE(mbedtls_ctx);
     uint32_t MBEDTLS_PRIVATE(num_ops);
     psa_key_attributes_t MBEDTLS_PRIVATE(attributes);
     unsigned int MBEDTLS_PRIVATE(error_occurred) : 1;