diff --git a/library/psa_crypto_ffdh.c b/library/psa_crypto_ffdh.c
index d8c153a..6b07a7d 100644
--- a/library/psa_crypto_ffdh.c
+++ b/library/psa_crypto_ffdh.c
@@ -131,6 +131,10 @@
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
+    if (!PSA_KEY_TYPE_IS_DH_KEY_PAIR(psa_get_key_type(attributes))) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     mbedtls_mpi_init(&P); mbedtls_mpi_init(&G);
     mbedtls_mpi_init(&X); mbedtls_mpi_init(&GY);
     mbedtls_mpi_init(&K);