Use PSA version of key agreement only for ECDHE keys

This should be cleaned when server-side static ECDH (1.2) support is added (#5320).

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/library/ssl_srv.c b/library/ssl_srv.c
index 55d7802..23fc49d 100644
--- a/library/ssl_srv.c
+++ b/library/ssl_srv.c
@@ -3879,55 +3879,86 @@
         ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
     {
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-        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 )
+        // 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 )
         {
-            MBEDTLS_SSL_DEBUG_MSG( 1, ( "Invalid buffer length" ) );
-            return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
-        }
+            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;
 
-        if( data_len < 1 || data_len > buf_len )
+            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 );
+                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
         {
-            MBEDTLS_SSL_DEBUG_MSG( 1, ( "Invalid data length" ) );
-            return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
+            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 );
         }
-
-        /* 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 );
-            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
+#else /* MBEDTLS_USE_PSA_CRYPTO */
         if( ( ret = mbedtls_ecdh_read_public( &ssl->handshake->ecdh_ctx,
                                       p, end - p) ) != 0 )
         {
@@ -3950,7 +3981,7 @@
 
         MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
                                 MBEDTLS_DEBUG_ECDH_Z );
-#endif
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||