diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index 42cabf5..030135b 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -1513,7 +1513,7 @@
     return 0;
 }
 
-static psa_key_type_t mbedtls_psa_parse_tls_ffdh_group(
+static psa_status_t  mbedtls_ssl_get_psa_ffdh_info_from_tls_id(
     uint16_t tls_ecc_grp_reg_id, size_t *bits, psa_key_type_t *key_type)
 {
     switch (tls_ecc_grp_reg_id) {
@@ -1556,28 +1556,21 @@
     mbedtls_ssl_handshake_params *handshake = ssl->handshake;
     size_t bits = 0;
     psa_key_type_t key_type = 0;
+    psa_algorithm_t alg = 0;
     size_t buf_size = (size_t) (end - buf);
 
-
     MBEDTLS_SSL_DEBUG_MSG(1, ("Perform PSA-based ECDH/FFDH computation."));
 
     /* Convert EC's TLS ID to PSA key type. */
 #if defined(PSA_WANT_ALG_ECDH)
-    psa_ecc_family_t ec_psa_family = 0;
     if (mbedtls_ssl_get_psa_curve_info_from_tls_id(
-            named_group, &ec_psa_family, &bits) == PSA_SUCCESS) {
-        key_type = PSA_KEY_TYPE_ECC_KEY_PAIR(ec_psa_family);
+            named_group, &key_type, &bits) == PSA_SUCCESS) {
+        alg = PSA_ALG_ECDH;
     }
 #endif
 #if defined(PSA_WANT_ALG_FFDH)
-    if (mbedtls_psa_parse_tls_ffdh_group(named_group, &bits, &key_type) == PSA_SUCCESS) {
-        if (PSA_KEY_TYPE_IS_DH(key_type)) {
-            if (buf_size < PSA_BITS_TO_BYTES(bits)) {
-
-                return MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL;
-            }
-            buf_size = PSA_BITS_TO_BYTES(bits);
-        }
+    if (mbedtls_ssl_get_psa_ffdh_info_from_tls_id(named_group, &bits, &key_type) == PSA_SUCCESS) {
+        alg = PSA_ALG_FFDH;
     }
 #endif
 
@@ -1585,22 +1578,17 @@
         return MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE;
     }
 
+    if (buf_size < PSA_BITS_TO_BYTES(bits)) {
+
+        return MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL;
+    }
+
     handshake->ecdh_psa_type = key_type;
     ssl->handshake->ecdh_bits = bits;
 
     key_attributes = psa_key_attributes_init();
     psa_set_key_usage_flags(&key_attributes, PSA_KEY_USAGE_DERIVE);
-
-    if (PSA_KEY_TYPE_IS_ECC(key_type)) {
-#if defined(PSA_WANT_ALG_ECDH)
-        psa_set_key_algorithm(&key_attributes, PSA_ALG_ECDH);
-#endif
-    } else {
-#if defined(PSA_WANT_ALG_FFDH)
-        psa_set_key_algorithm(&key_attributes, PSA_ALG_FFDH);
-#endif
-    }
-
+    psa_set_key_algorithm(&key_attributes, alg);
     psa_set_key_type(&key_attributes, handshake->ecdh_psa_type);
     psa_set_key_bits(&key_attributes, handshake->ecdh_bits);
 
@@ -1623,7 +1611,6 @@
         ret = PSA_TO_MBEDTLS_ERR(status);
         MBEDTLS_SSL_DEBUG_RET(1, "psa_export_public_key", ret);
         return ret;
-
     }
 
     *out_len = own_pubkey_len;