diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index d712deb..3415c28 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -5004,10 +5004,7 @@
 
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO) &&                   \
-    ( defined(MBEDTLS_KEY_EXCHANGE_PSK_ENABLED) ||       \
-      defined(MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED) ||   \
-      defined(MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED) || \
-      defined(MBEDTLS_KEY_EXCHANGE_DHE_PSK_ENABLED) )
+    defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED )
 static int ssl_use_opaque_psk( mbedtls_ssl_context const *ssl )
 {
     if( ssl->conf->f_psk != NULL )
@@ -5102,14 +5099,8 @@
 #endif /* MBEDTLS_SSL_EXTENDED_MS_ENABLED */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO) &&                   \
-    ( defined(MBEDTLS_KEY_EXCHANGE_PSK_ENABLED) ||       \
-      defined(MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED) ||   \
-      defined(MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED) || \
-      defined(MBEDTLS_KEY_EXCHANGE_DHE_PSK_ENABLED) )
-    if( ( handshake->ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_PSK ||
-          handshake->ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_RSA_PSK ||
-          handshake->ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_PSK ||
-          handshake->ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_DHE_PSK ) &&
+    defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
+    if( mbedtls_ssl_ciphersuite_uses_psk( handshake->ciphersuite_info ) == 1 &&
         ssl_use_opaque_psk( ssl ) == 1 )
     {
         /* Perform PSK-to-MS expansion in a single step. */