Refactor call hierarchy for ECDH so that it goes through the driver wrapper in a similar fashion to ECDSA.
Add component_test_psa_config_accel_ecdh to all.sh to test key agreement driver wrapper with libtestdriver1.

Signed-off-by: Aditya Deshpande <aditya.deshpande@arm.com>
diff --git a/tests/include/test/drivers/crypto_config_test_driver_extension.h b/tests/include/test/drivers/crypto_config_test_driver_extension.h
index 0bbca4a..fbfe8da 100644
--- a/tests/include/test/drivers/crypto_config_test_driver_extension.h
+++ b/tests/include/test/drivers/crypto_config_test_driver_extension.h
@@ -54,6 +54,14 @@
 #endif
 #endif
 
+#if defined(PSA_WANT_ALG_ECDH)
+#if defined(MBEDTLS_PSA_ACCEL_ALG_ECDH)
+#undef MBEDTLS_PSA_ACCEL_ALG_ECDH
+#else
+#define MBEDTLS_PSA_ACCEL_ALG_ECDH 1
+#endif
+#endif
+
 #if defined(PSA_WANT_ALG_MD5)
 #if defined(MBEDTLS_PSA_ACCEL_ALG_MD5)
 #undef MBEDTLS_PSA_ACCEL_ALG_MD5
@@ -202,7 +210,6 @@
 #define MBEDTLS_PSA_ACCEL_ALG_CCM 1
 #define MBEDTLS_PSA_ACCEL_ALG_CMAC 1
 #define MBEDTLS_PSA_ACCEL_ALG_ECB_NO_PADDING 1
-#define MBEDTLS_PSA_ACCEL_ALG_ECDH 1
 #define MBEDTLS_PSA_ACCEL_ALG_GCM 1
 #define MBEDTLS_PSA_ACCEL_ALG_HKDF 1
 #define MBEDTLS_PSA_ACCEL_ALG_HKDF_EXTRACT 1
@@ -215,6 +222,7 @@
 #define MBEDTLS_PSA_ACCEL_ALG_TLS12_PSK_TO_MS 1
 
 #if defined(MBEDTLS_PSA_ACCEL_ALG_ECDSA)
+#if defined(MBEDTLS_PSA_ACCEL_ALG_ECDH)
 #define MBEDTLS_PSA_ACCEL_ECC_BRAINPOOL_P_R1_256 1
 #define MBEDTLS_PSA_ACCEL_ECC_BRAINPOOL_P_R1_384 1
 #define MBEDTLS_PSA_ACCEL_ECC_BRAINPOOL_P_R1_512 1
@@ -229,6 +237,7 @@
 #define MBEDTLS_PSA_ACCEL_ECC_SECP_R1_384 1
 #define MBEDTLS_PSA_ACCEL_ECC_SECP_R1_521 1
 #endif
+#endif
 
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_DERIVE 1
 #define MBEDTLS_PSA_ACCEL_KEY_TYPE_HMAC 1
diff --git a/tests/scripts/all.sh b/tests/scripts/all.sh
index a4c6c86..7f1723b 100755
--- a/tests/scripts/all.sh
+++ b/tests/scripts/all.sh
@@ -1885,6 +1885,46 @@
     make test
 }
 
+component_test_psa_crypto_config_accel_ecdh () {
+    msg "test: MBEDTLS_PSA_CRYPTO_CONFIG with accelerated ECDH"
+
+    # Disable ALG_STREAM_CIPHER and ALG_ECB_NO_PADDING to avoid having
+    # partial support for cipher operations in the driver test library.
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_STREAM_CIPHER
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_ECB_NO_PADDING
+
+    # SHA384 needed for some ECDSA signature tests.
+    scripts/config.py -f tests/include/test/drivers/config_test_driver.h set MBEDTLS_SHA384_C
+    scripts/config.py -f tests/include/test/drivers/config_test_driver.h set MBEDTLS_SHA512_C
+
+    loc_accel_list="ALG_ECDH KEY_TYPE_ECC_KEY_PAIR KEY_TYPE_ECC_PUBLIC_KEY"
+    loc_accel_flags=$( echo "$loc_accel_list" | sed 's/[^ ]* */-DLIBTESTDRIVER1_MBEDTLS_PSA_ACCEL_&/g' )
+    make -C tests libtestdriver1.a CFLAGS=" -g3 $ASAN_CFLAGS $loc_accel_flags" LDFLAGS="$ASAN_CFLAGS"
+
+    # Restore test driver base configuration
+    scripts/config.py -f tests/include/test/drivers/config_test_driver.h unset MBEDTLS_SHA384_C
+    scripts/config.py -f tests/include/test/drivers/config_test_driver.h unset MBEDTLS_SHA512_C
+
+    scripts/config.py set MBEDTLS_PSA_CRYPTO_DRIVERS
+    scripts/config.py set MBEDTLS_PSA_CRYPTO_CONFIG
+    scripts/config.py unset MBEDTLS_USE_PSA_CRYPTO
+    scripts/config.py unset MBEDTLS_SSL_PROTO_TLS1_3
+    scripts/config.py unset MBEDTLS_ECDH_C
+    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED
+    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED
+    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
+    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED
+    scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED
+
+    loc_accel_flags="$loc_accel_flags $( echo "$loc_accel_list" | sed 's/[^ ]* */-DMBEDTLS_PSA_ACCEL_&/g' )"
+    make CFLAGS="$ASAN_CFLAGS -O -Werror -I../tests/include -I../tests -I../../tests -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_TEST_LIBTESTDRIVER1 -g3 $loc_accel_flags" LDFLAGS="-ltestdriver1 $ASAN_CFLAGS"
+
+    not grep mbedtls_ecdh_ library/ecdh.o
+
+    msg "test: MBEDTLS_PSA_CRYPTO_CONFIG with accelerated ECDH"
+    make test
+}
+
 component_test_psa_crypto_config_accel_rsa_signature () {
     msg "test: MBEDTLS_PSA_CRYPTO_CONFIG with accelerated RSA signature"
 
diff --git a/tests/src/drivers/test_driver_key_agreement.c b/tests/src/drivers/test_driver_key_agreement.c
index 884899f..ccea61d 100644
--- a/tests/src/drivers/test_driver_key_agreement.c
+++ b/tests/src/drivers/test_driver_key_agreement.c
@@ -19,15 +19,22 @@
 
 #include <test/helpers.h>
 
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
+
 #include "psa/crypto.h"
 #include "psa_crypto_core.h"
+#include "psa_crypto_ecp.h"
 
 #include "test/drivers/key_agreement.h"
 #include "test/drivers/test_driver.h"
 
 #include <string.h>
+#include <stdio.h>
 
-#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1)
+#include "libtestdriver1/include/psa/crypto.h"
+#include "libtestdriver1/library/psa_crypto_ecp.h"
+#endif
 
 mbedtls_test_driver_key_agreement_hooks_t
     mbedtls_test_driver_key_agreement_hooks = MBEDTLS_TEST_DRIVER_KEY_AGREEMENT_INIT;
@@ -58,16 +65,30 @@
         return( PSA_SUCCESS );
     }
 
-    return( psa_key_agreement_raw_builtin(
+    if( PSA_ALG_IS_ECDH(alg) )
+    {
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+    (LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+        return( libtestdriver1_mbedtls_psa_key_agreement_ecdh(
+                    (const libtestdriver1_psa_key_attributes_t *) attributes,
+                    key_buffer, key_buffer_size,
+                    alg, peer_key, peer_key_length,
+                    shared_secret, shared_secret_size,
+                    shared_secret_length ) );
+#elif defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+        return( mbedtls_psa_key_agreement_ecdh(
                 attributes,
-                key_buffer,
-                key_buffer_size,
-                alg,
-                peer_key,
-                peer_key_length,
-                shared_secret,
-                shared_secret_size,
+                key_buffer, key_buffer_size,
+                alg, peer_key, peer_key_length,
+                shared_secret, shared_secret_size,
                 shared_secret_length ) );
+#endif   
+    }
+    else
+    {
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+
 }
 
 #endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && PSA_CRYPTO_DRIVER_TEST */