diff --git a/library/psa_crypto_ffdh.c b/library/psa_crypto_ffdh.c
index aba62ee..de6cd75 100644
--- a/library/psa_crypto_ffdh.c
+++ b/library/psa_crypto_ffdh.c
@@ -64,33 +64,39 @@
     static const unsigned char dhm_G_8192[] =
         MBEDTLS_DHM_RFC7919_FFDHE8192_G_BIN;
 
-    if (key_size == sizeof(dhm_P_2048)) {
-        dhm_P = dhm_P_2048;
-        dhm_G = dhm_G_2048;
-        dhm_size_P = sizeof(dhm_P_2048);
-        dhm_size_G = sizeof(dhm_G_2048);
-    } else if (key_size == sizeof(dhm_P_3072)) {
-        dhm_P = dhm_P_3072;
-        dhm_G = dhm_G_3072;
-        dhm_size_P = sizeof(dhm_P_3072);
-        dhm_size_G = sizeof(dhm_G_3072);
-    } else if (key_size == sizeof(dhm_P_4096)) {
-        dhm_P = dhm_P_4096;
-        dhm_G = dhm_G_4096;
-        dhm_size_P = sizeof(dhm_P_4096);
-        dhm_size_G = sizeof(dhm_G_4096);
-    } else if (key_size == sizeof(dhm_P_6144)) {
-        dhm_P = dhm_P_6144;
-        dhm_G = dhm_G_6144;
-        dhm_size_P = sizeof(dhm_P_6144);
-        dhm_size_G = sizeof(dhm_G_6144);
-    } else if (key_size == sizeof(dhm_P_8192)) {
-        dhm_P = dhm_P_8192;
-        dhm_G = dhm_G_8192;
-        dhm_size_P = sizeof(dhm_P_8192);
-        dhm_size_G = sizeof(dhm_G_8192);
-    } else {
-        return PSA_ERROR_INVALID_ARGUMENT;
+    switch(key_size) {
+        case sizeof(dhm_P_2048):
+            dhm_P = dhm_P_2048;
+            dhm_G = dhm_G_2048;
+            dhm_size_P = sizeof(dhm_P_2048);
+            dhm_size_G = sizeof(dhm_G_2048);
+            break;
+        case sizeof(dhm_P_3072):
+            dhm_P = dhm_P_3072;
+            dhm_G = dhm_G_3072;
+            dhm_size_P = sizeof(dhm_P_3072);
+            dhm_size_G = sizeof(dhm_G_3072);
+            break;
+        case sizeof(dhm_P_4096):
+            dhm_P = dhm_P_4096;
+            dhm_G = dhm_G_4096;
+            dhm_size_P = sizeof(dhm_P_4096);
+            dhm_size_G = sizeof(dhm_G_4096);
+            break;
+        case sizeof(dhm_P_6144):
+            dhm_P = dhm_P_6144;
+            dhm_G = dhm_G_6144;
+            dhm_size_P = sizeof(dhm_P_6144);
+            dhm_size_G = sizeof(dhm_G_6144);
+            break;
+        case sizeof(dhm_P_8192):
+            dhm_P = dhm_P_8192;
+            dhm_G = dhm_G_8192;
+            dhm_size_P = sizeof(dhm_P_8192);
+            dhm_size_G = sizeof(dhm_G_8192);
+            break;
+        default:
+            return PSA_ERROR_INVALID_ARGUMENT;
     }
 
     if (P != NULL) {