diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index bc306cb..5fe969c 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3620,10 +3620,14 @@
     ret = mbedtls_pk_parse_public_key( &pk, peer_key, peer_key_length );
     if( ret != 0 )
         goto exit;
-    if( mbedtls_pk_get_type( &pk ) != MBEDTLS_PK_ECKEY )
+    switch( mbedtls_pk_get_type( &pk ) )
     {
-        ret = MBEDTLS_ERR_ECP_INVALID_KEY;
-        goto exit;
+        case MBEDTLS_PK_ECKEY:
+        case MBEDTLS_PK_ECKEY_DH:
+            break;
+        default:
+            ret = MBEDTLS_ERR_ECP_INVALID_KEY;
+            goto exit;
     }
     their_key = mbedtls_pk_ec( pk );
     ret = mbedtls_ecdh_get_params( &ecdh, their_key, MBEDTLS_ECDH_THEIRS );