diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3c3e80d..d9ba543 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -134,6 +134,15 @@
     return global_data.drivers_initialized;
 }
 
+static int psa_is_dh_key_size_valid(size_t bits) {
+    if (bits != 2048 && bits != 3072 && bits != 4096 &&
+        bits != 6144 && bits != 8192) {
+        return 0;
+    }
+
+    return 1;
+}
+
 psa_status_t mbedtls_to_psa_error(int ret)
 {
     /* Mbed TLS error codes can combine a high-level error code and a
@@ -632,9 +641,7 @@
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_FFDH_KEY_PAIR) || \
         defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_FFDH_PUBLIC_KEY)
         if (PSA_KEY_TYPE_IS_DH(type)) {
-            if (data_length != 256 && data_length != 384 &&
-                data_length != 512 && data_length != 768 &&
-                data_length != 1024) {
+            if (psa_is_dh_key_size_valid(PSA_BYTES_TO_BITS(data_length)) == 0) {
                 return PSA_ERROR_INVALID_ARGUMENT;
             }
 
@@ -6980,8 +6987,7 @@
 
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_FFDH_KEY_PAIR)
     if (PSA_KEY_TYPE_IS_DH(type) && PSA_KEY_TYPE_IS_KEY_PAIR(type)) {
-        if (bits != 2048 && bits != 3072 && bits != 4096 &&
-            bits != 6144 && bits != 8192) {
+        if (psa_is_dh_key_size_valid(bits) == 0) {
             return PSA_ERROR_NOT_SUPPORTED;
         }
     } else