Add support for RSA wrap in pk_psa_sign() test

Signed-off-by: Neil Armstrong <narmstrong@baylibre.com>
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index fe2a001..7f2b2b8 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1043,14 +1043,13 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_SHA256_C:MBEDTLS_USE_PSA_CRYPTO:MBEDTLS_ECDSA_C */
+/* BEGIN_CASE depends_on:MBEDTLS_SHA256_C:MBEDTLS_USE_PSA_CRYPTO */
 void pk_psa_sign( int parameter_arg,
                   int psa_type_arg, int expected_bits_arg )
 {
-    mbedtls_ecp_group_id grpid = parameter_arg;
     mbedtls_pk_context pk;
     unsigned char hash[32];
-    unsigned char sig[MBEDTLS_ECDSA_MAX_LEN];
+    unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
     unsigned char pkey_legacy[200];
     unsigned char pkey_psa[200];
     unsigned char *pkey_legacy_start, *pkey_psa_start;
@@ -1063,7 +1062,7 @@
 
     /*
      * This tests making signatures with a wrapped PSA key:
-     * - generate a fresh ECP legacy PK context
+     * - generate a fresh ECP/RSA legacy PK context
      * - wrap it in a PK context and make a signature this way
      * - extract the public key
      * - parse it to a PK context and verify the signature this way
@@ -1071,13 +1070,38 @@
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    /* Create legacy EC public/private key in PK context. */
-    mbedtls_pk_init( &pk );
-    TEST_ASSERT( mbedtls_pk_setup( &pk,
-                      mbedtls_pk_info_from_type( MBEDTLS_PK_ECKEY ) ) == 0 );
-    TEST_ASSERT( mbedtls_ecp_gen_key( grpid,
-                                      (mbedtls_ecp_keypair*) pk.pk_ctx,
-                                      mbedtls_test_rnd_std_rand, NULL ) == 0 );
+#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
+    if( PSA_KEY_TYPE_IS_RSA( psa_type_arg ) )
+    {
+        /* Create legacy RSA public/private key in PK context. */
+        mbedtls_pk_init( &pk );
+        TEST_ASSERT( mbedtls_pk_setup( &pk,
+                        mbedtls_pk_info_from_type( MBEDTLS_PK_RSA ) ) == 0 );
+        TEST_ASSERT( mbedtls_rsa_gen_key( mbedtls_pk_rsa( pk ),
+                        mbedtls_test_rnd_std_rand, NULL,
+                        parameter_arg, 3 ) == 0 );
+    }
+    else
+#endif /* MBEDTLS_RSA_C && MBEDTLS_GENPRIME */
+#if defined(MBEDTLS_ECDSA_C)
+    if( PSA_KEY_TYPE_IS_ECC_KEY_PAIR( psa_type_arg ) )
+    {
+        mbedtls_ecp_group_id grpid = parameter_arg;
+
+        /* Create legacy EC public/private key in PK context. */
+        mbedtls_pk_init( &pk );
+        TEST_ASSERT( mbedtls_pk_setup( &pk,
+                        mbedtls_pk_info_from_type( MBEDTLS_PK_ECKEY ) ) == 0 );
+        TEST_ASSERT( mbedtls_ecp_gen_key( grpid,
+                        (mbedtls_ecp_keypair*) pk.pk_ctx,
+                        mbedtls_test_rnd_std_rand, NULL ) == 0 );
+    }
+    else
+#endif /* MBEDTLS_ECDSA_C */
+    {
+        (void) parameter_arg;
+        TEST_ASSUME( ! "Opaque PK key not supported in this configuration" );
+    }
 
     /* Export underlying public key for re-importing in a legacy context. */
     ret = mbedtls_pk_write_pubkey_der( &pk, pkey_legacy,