From 331c3421d1f0c6427c6b5d8083b80ecac749c646 Mon Sep 17 00:00:00 2001
From: Przemyslaw Stekiel <przemyslaw.stekiel@mobica.com>
Date: Thu, 21 Oct 2021 12:26:58 +0200
Subject: [PATCH] Address review comments

Signed-off-by: Przemyslaw Stekiel <przemyslaw.stekiel@mobica.com>
---
 programs/ssl/ssl_server2.c | 42 ++++++++++++++++++++++++++++----------
 tests/ssl-opt.sh           |  4 ++--
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 05cf085169..29c696f336 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -202,7 +202,7 @@ int main( void )
 #endif /* MBEDTLS_X509_CRT_PARSE_C */
 #if defined(MBEDTLS_USE_PSA_CRYPTO) && defined(MBEDTLS_X509_CRT_PARSE_C)
 #define USAGE_KEY_OPAQUE \
-    "    key_opaque=%%d       Handle your private key as if it were opaque\n" \
+    "    key_opaque=%%d       Handle your private keys as if they were opaque\n" \
     "                        default: 0 (disabled)\n"
 #else
 #define USAGE_KEY_OPAQUE ""
@@ -1320,8 +1320,9 @@ int main( int argc, char *argv[] )
     mbedtls_pk_context pkey;
     mbedtls_x509_crt srvcert2;
     mbedtls_pk_context pkey2;
-#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
     psa_key_id_t key_slot = 0; /* invalid key slot */
+    psa_key_id_t key_slot2 = 0; /* invalid key slot */
 #endif
     int key_cert_init = 0, key_cert_init2 = 0;
 #if defined(MBEDTLS_SSL_ASYNC_PRIVATE)
@@ -2519,24 +2520,39 @@ int main( int argc, char *argv[] )
                             (unsigned int) -ret );
             goto exit;
         }
+        key_cert_init2 = 2;
+#endif /* MBEDTLS_ECDSA_C */
+    }
+
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if( opt.key_opaque != 0 )
     {
-        if( ( ret = mbedtls_pk_wrap_as_opaque( &pkey2, &key_slot,
-                                               PSA_ALG_SHA_256 ) ) != 0 )
+        if ( mbedtls_pk_get_type( &pkey ) == MBEDTLS_PK_ECKEY )
         {
-            mbedtls_printf( " failed\n  !  "
-                            "mbedtls_pk_wrap_as_opaque returned -0x%x\n\n", (unsigned int)  -ret );
-            goto exit;
+            if( ( ret = mbedtls_pk_wrap_as_opaque( &pkey, &key_slot,
+                                                PSA_ALG_SHA_256 ) ) != 0 )
+            {
+                mbedtls_printf( " failed\n  !  "
+                                "mbedtls_pk_wrap_as_opaque returned -0x%x\n\n", (unsigned int)  -ret );
+                goto exit;
+            }
+        }
+
+        if ( mbedtls_pk_get_type( &pkey2 ) == MBEDTLS_PK_ECKEY )
+        {
+            if( ( ret = mbedtls_pk_wrap_as_opaque( &pkey2, &key_slot2,
+                                                PSA_ALG_SHA_256 ) ) != 0 )
+            {
+                mbedtls_printf( " failed\n  !  "
+                                "mbedtls_pk_wrap_as_opaque returned -0x%x\n\n", (unsigned int)  -ret );
+                goto exit;
+            }
         }
     }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
-        key_cert_init2 = 2;
-#endif /* MBEDTLS_ECDSA_C */
 #endif /* MBEDTLS_CERTS_C */
-    }
 
-    mbedtls_printf( " ok (key type: %s)\n", mbedtls_pk_get_name( &pkey2 ) );
+    mbedtls_printf( " ok (key types: %s - %s)\n", mbedtls_pk_get_name( &pkey ), mbedtls_pk_get_name( &pkey2 ) );
 #endif /* MBEDTLS_X509_CRT_PARSE_C */
 
 #if defined(MBEDTLS_DHM_C) && defined(MBEDTLS_FS_IO)
@@ -4031,6 +4047,10 @@ exit:
     mbedtls_pk_free( &pkey );
     mbedtls_x509_crt_free( &srvcert2 );
     mbedtls_pk_free( &pkey2 );
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_destroy_key( key_slot );
+    psa_destroy_key( key_slot2 );
+#endif
 #endif
 
 #if defined(MBEDTLS_DHM_C) && defined(MBEDTLS_FS_IO)
diff --git a/tests/ssl-opt.sh b/tests/ssl-opt.sh
index ef7238c5fc..187f7a0f1f 100755
--- a/tests/ssl-opt.sh
+++ b/tests/ssl-opt.sh
@@ -1643,7 +1643,7 @@ run_test    "Opaque key for server authentication" \
              key_file=data_files/server5.key" \
             0 \
             -c "Verifying peer X.509 certificate... ok" \
-            -s "key type: Opaque" \
+            -s "key types: RSA - Opaque" \
             -S "error" \
             -C "error"
 
@@ -1659,7 +1659,7 @@ run_test    "Opaque key for client/server authentication" \
             0 \
             -c "key type: Opaque" \
             -c "Verifying peer X.509 certificate... ok" \
-            -s "key type: Opaque" \
+            -s "key types: RSA - Opaque" \
             -s "Verifying peer X.509 certificate... ok" \
             -S "error" \
             -C "error"