ssl_tls: fix guards for accelerated ECDH

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/include/mbedtls/check_config.h b/include/mbedtls/check_config.h
index ab6b5dd..08874e3 100644
--- a/include/mbedtls/check_config.h
+++ b/include/mbedtls/check_config.h
@@ -304,15 +304,26 @@
 #endif
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
+/* Helper for ECDH dependencies, will be undefined at the end of the file */
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(PSA_HAVE_FULL_ECDH)
+#define MBEDTLS_PK_HAVE_ECDH
+#endif
+#else /* MBEDTLS_USE_PSA_CRYPTO */
+#if defined(MBEDTLS_ECDH_C)
+#define MBEDTLS_PK_HAVE_ECDH
+#endif
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED) &&                 \
-    ( !defined(MBEDTLS_ECDH_C) ||                                       \
+    ( !defined(MBEDTLS_PK_HAVE_ECDH) ||                                       \
       !defined(MBEDTLS_PK_HAVE_ECDSA) ||                                \
       !defined(MBEDTLS_X509_CRT_PARSE_C) )
 #error "MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED defined, but not all prerequisites"
 #endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) &&                 \
-    ( !defined(MBEDTLS_ECDH_C) || !defined(MBEDTLS_RSA_C) ||          \
+    ( !defined(MBEDTLS_PK_HAVE_ECDH) || !defined(MBEDTLS_RSA_C) ||          \
       !defined(MBEDTLS_X509_CRT_PARSE_C) )
 #error "MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED defined, but not all prerequisites"
 #endif
@@ -1100,6 +1111,7 @@
 /* Undefine helper symbols */
 #undef MBEDTLS_PK_HAVE_ECDSA
 #undef MBEDTLS_PK_HAVE_JPAKE
+#undef MBEDTLS_PK_HAVE_ECDH
 
 /*
  * Avoid warning from -pedantic. This is a convenient place for this
diff --git a/include/mbedtls/config_psa.h b/include/mbedtls/config_psa.h
index cedce1e..d0c24d8 100644
--- a/include/mbedtls/config_psa.h
+++ b/include/mbedtls/config_psa.h
@@ -853,6 +853,11 @@
 #define PSA_HAVE_FULL_JPAKE 1
 #endif
 
+#if defined(PSA_WANT_ALG_ECDH) && defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR) && \
+    defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
+#define PSA_HAVE_FULL_ECDH 1
+#endif
+
 /* These features are always enabled. */
 #define PSA_WANT_KEY_TYPE_DERIVE 1
 #define PSA_WANT_KEY_TYPE_PASSWORD 1
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index 3de7a8f..f029265 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -193,6 +193,12 @@
 #endif /* PSA_WANT_ALG_ECDSA */
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
+/* Symbol for ECDH capabilities, no matter how it is provided */
+#if (defined(MBEDTLS_USE_PSA_CRYPTO) && defined(PSA_HAVE_FULL_ECDH)) || \
+    (!defined(MBEDTLS_USE_PSA_CRYPTO) && defined(MBEDTLS_ECDH_C))
+#define MBEDTLS_PK_CAN_ECDH
+#endif
+
 #if defined(MBEDTLS_PK_CAN_ECDSA_VERIFY) || defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
 #define MBEDTLS_PK_CAN_ECDSA_SOME
 #endif
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index a6129da..c91a12f 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -4220,7 +4220,7 @@
     mbedtls_ssl_buffering_free(ssl);
 #endif /* MBEDTLS_SSL_PROTO_DTLS */
 
-#if defined(MBEDTLS_ECDH_C) && \
+#if defined(MBEDTLS_PK_CAN_ECDH) && \
     (defined(MBEDTLS_USE_PSA_CRYPTO) || defined(MBEDTLS_SSL_PROTO_TLS1_3))
     if (handshake->ecdh_psa_privkey_is_external == 0) {
         psa_destroy_key(handshake->ecdh_psa_privkey);