replace pk_sign with pk_sign_ext

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 913280e..7dd8fb4 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -32,6 +32,7 @@
 
 #include "ssl_misc.h"
 #include "ssl_tls13_keys.h"
+#include "ssl_debug_helpers.h"
 
 int mbedtls_ssl_tls13_fetch_handshake_msg( mbedtls_ssl_context *ssl,
                                            unsigned hs_type,
@@ -981,6 +982,70 @@
 /*
  * STATE HANDLING: Output Certificate Verify
  */
+static uint16_t ssl_tls13_get_sig_alg_from_pk( mbedtls_ssl_context *ssl,
+                                               mbedtls_pk_context *own_key )
+{
+    mbedtls_pk_type_t sig = mbedtls_ssl_sig_from_pk( own_key );
+    /* Determine the size of the key */
+    size_t own_key_size = mbedtls_pk_get_bitlen( own_key );
+    uint16_t algorithm = MBEDTLS_TLS1_3_SIG_NONE;
+    ((void) own_key_size);
+
+    switch( sig )
+    {
+#if defined(MBEDTLS_ECDSA_C)
+        case MBEDTLS_SSL_SIG_ECDSA:
+            switch( own_key_size )
+            {
+                case 256:
+                    algorithm = MBEDTLS_TLS1_3_SIG_ECDSA_SECP256R1_SHA256;
+                    break;
+                case 384:
+                    algorithm = MBEDTLS_TLS1_3_SIG_ECDSA_SECP384R1_SHA384;
+                    break;
+                case 521:
+                    algorithm = MBEDTLS_TLS1_3_SIG_ECDSA_SECP521R1_SHA512;
+                    break;
+                default:
+                    MBEDTLS_SSL_DEBUG_MSG( 3,
+                                           ( "unknown key size: %"
+                                             MBEDTLS_PRINTF_SIZET " bits",
+                                             own_key_size ) );
+                    break;
+            }
+            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 ) )
+            {
+                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 ) )
+            {
+                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 ) )
+            {
+                algorithm = MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA512;
+            }
+            break;
+#endif /* MBEDTLS_X509_RSASSA_PSS_SUPPORT */
+        default:
+            MBEDTLS_SSL_DEBUG_MSG( 1,
+                                   ( "unkown signature type : %u", sig ) );
+            break;
+    }
+    return( algorithm );
+}
+
 static int ssl_tls13_write_certificate_verify_body( mbedtls_ssl_context *ssl,
                                                     unsigned char *buf,
                                                     unsigned char *end,
@@ -994,12 +1059,11 @@
     size_t handshake_hash_len;
     unsigned char verify_buffer[ SSL_VERIFY_STRUCT_MAX_SIZE ];
     size_t verify_buffer_len;
-    unsigned char signature_type;
-#if defined(MBEDTLS_ECDSA_C) || defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
-    size_t own_key_size;
-#endif /* MBEDTLS_ECDSA_C || MBEDTLS_X509_RSASSA_PSS_SUPPORT */
-    mbedtls_md_type_t md_alg;
+    mbedtls_pk_type_t pk_type = MBEDTLS_PK_NONE;
+    const void *options = NULL;
+    mbedtls_md_type_t md_alg = MBEDTLS_MD_NONE;
     uint16_t algorithm = MBEDTLS_TLS1_3_SIG_NONE;
+    mbedtls_pk_rsassa_pss_options pss_opts;
     size_t signature_len = 0;
     const mbedtls_md_info_t *md_info;
     unsigned char verify_hash[ MBEDTLS_MD_MAX_SIZE ];
@@ -1036,101 +1100,60 @@
      *    opaque signature<0..2^16-1>;
      *  } CertificateVerify;
      */
-    signature_type = mbedtls_ssl_sig_from_pk( own_key );
-#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 || MBEDTLS_X509_RSASSA_PSS_SUPPORT */
-    switch( signature_type )
-    {
-#if defined(MBEDTLS_ECDSA_C)
-        case MBEDTLS_SSL_SIG_ECDSA:
-            switch( own_key_size )
-            {
-                case 256:
-                    md_alg  = MBEDTLS_MD_SHA256;
-                    algorithm = MBEDTLS_TLS1_3_SIG_ECDSA_SECP256R1_SHA256;
-                    break;
-                case 384:
-                    md_alg  = MBEDTLS_MD_SHA384;
-                    algorithm = MBEDTLS_TLS1_3_SIG_ECDSA_SECP384R1_SHA384;
-                    break;
-                case 521:
-                    md_alg  = MBEDTLS_MD_SHA512;
-                    algorithm = MBEDTLS_TLS1_3_SIG_ECDSA_SECP521R1_SHA512;
-                    break;
-                default:
-                    MBEDTLS_SSL_DEBUG_MSG( 3,
-                                           ( "unknown key size: %"
-                                             MBEDTLS_PRINTF_SIZET " bits",
-                                             own_key_size ) );
-                    break;
-            }
-            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;
-                MBEDTLS_SSL_DEBUG_MSG( 1,
-                            ( "signature algorthm is 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;
-                MBEDTLS_SSL_DEBUG_MSG( 1,
-                            ( "signature algorthm is 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;
-                MBEDTLS_SSL_DEBUG_MSG( 1,
-                            ( "signature algorthm is 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 ) );
-            return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
-    }
-
+    algorithm = ssl_tls13_get_sig_alg_from_pk( ssl, own_key );
     if( algorithm == MBEDTLS_TLS1_3_SIG_NONE ||
         ! mbedtls_ssl_sig_alg_is_received( ssl, algorithm ) )
     {
         MBEDTLS_SSL_DEBUG_MSG( 1,
                     ( "signature algorithm not in received or offered list." ) );
+
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "Signature algorithm is %s",
+                                    mbedtls_ssl_sig_alg_to_str( algorithm ) ) );
+
         MBEDTLS_SSL_PEND_FATAL_ALERT( MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE,
                                       MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
         return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
     }
 
+    switch( algorithm )
+    {
+#if defined(MBEDTLS_ECDSA_C)
+    case MBEDTLS_TLS1_3_SIG_ECDSA_SECP256R1_SHA256:
+        md_alg = MBEDTLS_MD_SHA256;
+        pk_type = MBEDTLS_PK_ECDSA;
+        break;
+    case MBEDTLS_TLS1_3_SIG_ECDSA_SECP384R1_SHA384:
+        md_alg = MBEDTLS_MD_SHA384;
+        pk_type = MBEDTLS_PK_ECDSA;
+        break;
+    case MBEDTLS_TLS1_3_SIG_ECDSA_SECP521R1_SHA512:
+        md_alg = MBEDTLS_MD_SHA512;
+        pk_type = MBEDTLS_PK_ECDSA;
+        break;
+#endif /* MBEDTLS_ECDSA_C */
+#if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
+    case MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA256:
+        md_alg = MBEDTLS_MD_SHA256;
+        pss_opts.expected_salt_len = MBEDTLS_RSA_SALT_LEN_ANY;
+        options = &pss_opts;
+        pk_type = MBEDTLS_PK_RSASSA_PSS;
+        break;
+    case MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA384:
+        md_alg = MBEDTLS_MD_SHA384;
+        pss_opts.expected_salt_len = MBEDTLS_RSA_SALT_LEN_ANY;
+        options = &pss_opts;
+        pk_type = MBEDTLS_PK_RSASSA_PSS;
+        break;
+    case MBEDTLS_TLS1_3_SIG_RSA_PSS_RSAE_SHA512:
+        md_alg = MBEDTLS_MD_SHA512;
+        pss_opts.expected_salt_len = MBEDTLS_RSA_SALT_LEN_ANY;
+        options = &pss_opts;
+        pk_type = MBEDTLS_PK_RSASSA_PSS;
+        break;
+#endif /* MBEDTLS_X509_RSASSA_PSS_SUPPORT */
+    default:
+        break;
+    }
     /* Check there is space for the algorithm identifier (2 bytes) and the
      * signature length (2 bytes).
      */
@@ -1150,10 +1173,10 @@
     verify_hash_len = mbedtls_md_get_size( md_info );
     MBEDTLS_SSL_DEBUG_BUF( 3, "verify hash", verify_hash, verify_hash_len );
 
-    if( ( ret = mbedtls_pk_sign( own_key, md_alg,
-                                 verify_hash, verify_hash_len,
-                                 p + 2, (size_t)( end - ( p + 2 ) ), &signature_len,
-                                 ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
+    if( ( ret = mbedtls_pk_sign_ext( pk_type, options,
+                        own_key, md_alg, verify_hash, verify_hash_len,
+                        p + 2, (size_t)( end - ( p + 2 ) ), &signature_len,
+                        ssl->conf->f_rng, ssl->conf->p_rng ) ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_pk_sign", ret );
         return( ret );