Add rsa_pss_rsae_sha256 support

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index 24a3d9d..8b0f668 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -980,9 +980,9 @@
     unsigned char verify_buffer[ SSL_VERIFY_STRUCT_MAX_SIZE ];
     size_t verify_buffer_len;
     unsigned char signature_type;
-#if defined(MBEDTLS_ECDSA_C)
+#if defined(MBEDTLS_ECDSA_C) || defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
     size_t own_key_size;
-#endif /* MBEDTLS_ECDSA_C */
+#endif /* MBEDTLS_ECDSA_C || MBEDTLS_X509_RSASSA_PSS_SUPPORT */
     mbedtls_md_type_t md_alg;
     uint16_t algorithm = MBEDTLS_TLS1_3_SIG_NONE;
     size_t signature_len = 0;
@@ -1022,10 +1022,10 @@
      *  } CertificateVerify;
      */
     signature_type = mbedtls_ssl_sig_from_pk( own_key );
-#if defined(MBEDTLS_ECDSA_C)
+#if defined(MBEDTLS_ECDSA_C) || defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
     /* Determine the size of the key */
     own_key_size = mbedtls_pk_get_bitlen( own_key );
-#endif /* MBEDTLS_ECDSA_C */
+#endif /* MBEDTLS_ECDSA_C || MBEDTLS_X509_RSASSA_PSS_SUPPORT */
     switch( signature_type )
     {
 #if defined(MBEDTLS_ECDSA_C)
@@ -1054,10 +1054,50 @@
             break;
 #endif /* MBEDTLS_ECDSA_C */
 
+#if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
+        case MBEDTLS_SSL_SIG_RSA:
+            if( own_key_size <= 2048 &&
+                mbedtls_ssl_sig_alg_is_received( ssl,
+                                    MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA256 ) )
+            {
+                md_alg  = MBEDTLS_MD_SHA256;
+                algorithm = MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA256;
+            }
+            else if( own_key_size <= 3072 &&
+                     mbedtls_ssl_sig_alg_is_received( ssl,
+                                    MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA384 ) )
+            {
+                md_alg  = MBEDTLS_MD_SHA384;
+                algorithm = MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA384;
+            }
+            else if( own_key_size <= 4096 &&
+                     mbedtls_ssl_sig_alg_is_received( ssl,
+                                    MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA512 ) )
+            {
+                md_alg  = MBEDTLS_MD_SHA512;
+                algorithm = MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA512;
+            }
+            else
+            {
+                MBEDTLS_SSL_DEBUG_MSG( 3, ( "unknown key size: %"
+                                            MBEDTLS_PRINTF_SIZET " bits",
+                                            own_key_size ) );
+                break;
+            }
+
+            if( mbedtls_rsa_set_padding( mbedtls_pk_rsa( *own_key ),
+                                         MBEDTLS_RSA_PKCS_V21,
+                                         md_alg ) != 0 )
+            {
+                MBEDTLS_SSL_DEBUG_MSG( 1, ( "Set RSA padding Fail" ) );
+                return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
+            }
+            break;
+#endif /* MBEDTLS_X509_RSASSA_PSS_SUPPORT */
         default:
             MBEDTLS_SSL_DEBUG_MSG( 1,
                                    ( "unkown pk type : %d", signature_type ) );
-            break;
+            return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
     }
 
     if( algorithm == MBEDTLS_TLS1_3_SIG_NONE ||