diff --git a/library/ssl_srv.c b/library/ssl_srv.c
index be41593..4f91b8b 100644
--- a/library/ssl_srv.c
+++ b/library/ssl_srv.c
@@ -3108,9 +3108,8 @@
         }
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-        // Handle only ECDHE keys using PSA crypto.
-        if ( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA ||
-             ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA )
+        if( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA ||
+            ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA )
         {
             psa_status_t status = PSA_ERROR_GENERIC_ERROR;
             psa_key_attributes_t key_attributes;
@@ -3194,28 +3193,16 @@
             len += header_size;
         }
         else
-        {
-            if( ( ret = mbedtls_ecdh_make_params(
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+        if( ( ret = mbedtls_ecdh_make_params(
                     &ssl->handshake->ecdh_ctx, &len,
                     ssl->out_msg + ssl->out_msglen,
                     MBEDTLS_SSL_OUT_CONTENT_LEN - ssl->out_msglen,
                     ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
-            {
-                MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_make_params", ret );
-                return( ret );
-            }
-        }
-#else
-        if( ( ret = mbedtls_ecdh_make_params(
-                  &ssl->handshake->ecdh_ctx, &len,
-                  ssl->out_msg + ssl->out_msglen,
-                  MBEDTLS_SSL_OUT_CONTENT_LEN - ssl->out_msglen,
-                  ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_make_params", ret );
             return( ret );
         }
-#endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE_WITH_SERVER_SIGNATURE_ENABLED)
         dig_signed = ssl->out_msg + ssl->out_msglen;
@@ -3885,6 +3872,68 @@
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED */
+#if defined(MBEDTLS_USE_PSA_CRYPTO) &&                           \
+        ( defined(MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED) ||     \
+          defined(MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED) )
+    if( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA ||
+        ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA )
+    {
+        size_t data_len = (size_t)( *p++ );
+        size_t buf_len = (size_t)( end - p );
+        psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+        mbedtls_ssl_handshake_params *handshake = ssl->handshake;
+
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "Read the peer's public key." ) );
+
+        /*
+        * We must have at least two bytes (1 for length, at least 1 for data)
+        */
+        if( buf_len < 2 )
+        {
+            MBEDTLS_SSL_DEBUG_MSG( 1, ( "Invalid buffer length" ) );
+            return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
+        }
+
+        if( data_len < 1 || data_len > buf_len )
+        {
+            MBEDTLS_SSL_DEBUG_MSG( 1, ( "Invalid data length" ) );
+            return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
+        }
+
+        /* Store peer's ECDH public key. */
+        memcpy( handshake->ecdh_psa_peerkey, p, data_len );
+        handshake->ecdh_psa_peerkey_len = data_len;
+
+        /* Compute ECDH shared secret. */
+        status = psa_raw_key_agreement(
+                    PSA_ALG_ECDH, handshake->ecdh_psa_privkey,
+                    handshake->ecdh_psa_peerkey, handshake->ecdh_psa_peerkey_len,
+                    handshake->premaster, sizeof( handshake->premaster ),
+                    &handshake->pmslen );
+        if( status != PSA_SUCCESS )
+        {
+            ret = psa_ssl_status_to_mbedtls( status );
+            MBEDTLS_SSL_DEBUG_RET( 1, "psa_raw_key_agreement", ret );
+            (void) psa_destroy_key( handshake->ecdh_psa_privkey );
+            handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
+            return( ret );
+        }
+
+        status = psa_destroy_key( handshake->ecdh_psa_privkey );
+
+        if( status != PSA_SUCCESS )
+        {
+            ret = psa_ssl_status_to_mbedtls( status );
+            MBEDTLS_SSL_DEBUG_RET( 1, "psa_destroy_key", ret );
+            return( ret );
+        }
+        handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
+    }
+    else
+
+#endif /* MBEDTLS_USE_PSA_CRYPTO &&
+            ( MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||
+              MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED ) */
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED) ||                     \
     defined(MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED) ||                   \
     defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) ||                      \
@@ -3894,89 +3943,6 @@
         ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDH_RSA ||
         ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
     {
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-        // Handle only ECDHE keys using PSA crypto.
-        if ( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA ||
-             ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA )
-        {
-            size_t data_len = (size_t)( *p++ );
-            size_t buf_len = (size_t)( end - p );
-            psa_status_t status = PSA_ERROR_GENERIC_ERROR;
-            mbedtls_ssl_handshake_params *handshake = ssl->handshake;
-
-            MBEDTLS_SSL_DEBUG_MSG( 1, ( "Read the peer's public key." ) );
-
-            /*
-            * We must have at least two bytes (1 for length, at least 1 for data)
-            */
-            if( buf_len < 2 )
-            {
-                MBEDTLS_SSL_DEBUG_MSG( 1, ( "Invalid buffer length" ) );
-                return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
-            }
-
-            if( data_len < 1 || data_len > buf_len )
-            {
-                MBEDTLS_SSL_DEBUG_MSG( 1, ( "Invalid data length" ) );
-                return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
-            }
-
-            /* Store peer's ECDH public key. */
-            memcpy( handshake->ecdh_psa_peerkey, p, data_len );
-            handshake->ecdh_psa_peerkey_len = data_len;
-
-            /* Compute ECDH shared secret. */
-            status = psa_raw_key_agreement(
-                        PSA_ALG_ECDH, handshake->ecdh_psa_privkey,
-                        handshake->ecdh_psa_peerkey, handshake->ecdh_psa_peerkey_len,
-                        handshake->premaster, sizeof( handshake->premaster ),
-                        &handshake->pmslen );
-            if( status != PSA_SUCCESS )
-            {
-                ret = psa_ssl_status_to_mbedtls( status );
-                MBEDTLS_SSL_DEBUG_RET( 1, "psa_raw_key_agreement", ret );
-                (void) psa_destroy_key( handshake->ecdh_psa_privkey );
-                handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
-                return( ret );
-            }
-
-            status = psa_destroy_key( handshake->ecdh_psa_privkey );
-
-            if( status != PSA_SUCCESS )
-            {
-                ret = psa_ssl_status_to_mbedtls( status );
-                MBEDTLS_SSL_DEBUG_RET( 1, "psa_destroy_key", ret );
-                return( ret );
-            }
-            handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
-
-        }
-        else
-        {
-            if( ( ret = mbedtls_ecdh_read_public( &ssl->handshake->ecdh_ctx,
-                                p, end - p) ) != 0 )
-            {
-                MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_read_public", ret );
-                return( MBEDTLS_ERR_SSL_DECODE_ERROR );
-            }
-
-            MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
-                                    MBEDTLS_DEBUG_ECDH_QP );
-
-            if( ( ret = mbedtls_ecdh_calc_secret( &ssl->handshake->ecdh_ctx,
-                                        &ssl->handshake->pmslen,
-                                        ssl->handshake->premaster,
-                                        MBEDTLS_MPI_MAX_SIZE,
-                                        ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
-            {
-                MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_calc_secret", ret );
-                return( MBEDTLS_ERR_SSL_DECODE_ERROR );
-            }
-
-            MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
-                                    MBEDTLS_DEBUG_ECDH_Z );
-        }
-#else /* MBEDTLS_USE_PSA_CRYPTO */
         if( ( ret = mbedtls_ecdh_read_public( &ssl->handshake->ecdh_ctx,
                                       p, end - p) ) != 0 )
         {
@@ -3999,7 +3965,6 @@
 
         MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
                                 MBEDTLS_DEBUG_ECDH_Z );
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||