Update mbedtls_pk_wrap_as_opaque() usage in SSL client2 & server2

Signed-off-by: Neil Armstrong <narmstrong@baylibre.com>
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index cd60986..3275a58 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -1709,8 +1709,22 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if( opt.key_opaque != 0 )
     {
-        if( ( ret = mbedtls_pk_wrap_as_opaque( &pkey, &key_slot,
-                                               PSA_ALG_ANY_HASH ) ) != 0 )
+        psa_algorithm_t psa_alg, psa_alg2;
+
+        if( mbedtls_pk_get_type( &pkey ) == MBEDTLS_PK_ECKEY )
+        {
+            psa_alg = PSA_ALG_ECDSA( PSA_ALG_ANY_HASH );
+            psa_alg2 = PSA_ALG_NONE;
+        }
+        else
+        {
+            psa_alg = PSA_ALG_RSA_PKCS1V15_SIGN( PSA_ALG_ANY_HASH );
+            psa_alg2 = PSA_ALG_RSA_PSS( PSA_ALG_ANY_HASH );
+        }
+
+        if( ( ret = mbedtls_pk_wrap_as_opaque( &pkey, &key_slot, psa_alg,
+                                               PSA_KEY_USAGE_SIGN_HASH,
+                                               psa_alg2 ) ) != 0 )
         {
             mbedtls_printf( " failed\n  !  "
                             "mbedtls_pk_wrap_as_opaque returned -0x%x\n\n", (unsigned int)  -ret );
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index a91af0e..46a9188 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -2575,11 +2575,29 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if( opt.key_opaque != 0 )
     {
+        psa_algorithm_t psa_alg, psa_alg2;
+        psa_key_usage_t psa_usage;
+
         if ( mbedtls_pk_get_type( &pkey ) == MBEDTLS_PK_ECKEY ||
              mbedtls_pk_get_type( &pkey ) == MBEDTLS_PK_RSA )
         {
+            if( mbedtls_pk_get_type( &pkey ) == MBEDTLS_PK_ECKEY )
+            {
+                psa_alg = PSA_ALG_ECDSA( PSA_ALG_ANY_HASH );
+                psa_alg2 = PSA_ALG_ECDH;
+                psa_usage = PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_DERIVE;
+            }
+            else
+            {
+                psa_alg = PSA_ALG_RSA_PKCS1V15_SIGN( PSA_ALG_ANY_HASH );
+                psa_alg2 = PSA_ALG_NONE;
+                psa_usage = PSA_KEY_USAGE_SIGN_HASH;
+            }
+
             if( ( ret = mbedtls_pk_wrap_as_opaque( &pkey, &key_slot,
-                                                PSA_ALG_ANY_HASH ) ) != 0 )
+                                                   psa_alg,
+                                                   psa_usage,
+                                                   psa_alg2 ) ) != 0 )
             {
                 mbedtls_printf( " failed\n  !  "
                                 "mbedtls_pk_wrap_as_opaque returned -0x%x\n\n", (unsigned int)  -ret );
@@ -2590,8 +2608,23 @@
         if ( mbedtls_pk_get_type( &pkey2 ) == MBEDTLS_PK_ECKEY ||
              mbedtls_pk_get_type( &pkey2 ) == MBEDTLS_PK_RSA )
         {
+            if( mbedtls_pk_get_type( &pkey2 ) == MBEDTLS_PK_ECKEY )
+            {
+                psa_alg = PSA_ALG_ECDSA( PSA_ALG_ANY_HASH );
+                psa_alg2 = PSA_ALG_ECDH;
+                psa_usage = PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_DERIVE;
+            }
+            else
+            {
+                psa_alg = PSA_ALG_RSA_PKCS1V15_SIGN( PSA_ALG_ANY_HASH );
+                psa_alg2 = PSA_ALG_NONE;
+                psa_usage = PSA_KEY_USAGE_SIGN_HASH;
+            }
+
             if( ( ret = mbedtls_pk_wrap_as_opaque( &pkey2, &key_slot2,
-                                                PSA_ALG_ANY_HASH ) ) != 0 )
+                                                   psa_alg,
+                                                   psa_usage,
+                                                   psa_alg2 ) ) != 0 )
             {
                 mbedtls_printf( " failed\n  !  "
                                 "mbedtls_pk_wrap_as_opaque returned -0x%x\n\n", (unsigned int)  -ret );