From d69439aa61dacd39979a9d3f263422eedbe3fb81 Mon Sep 17 00:00:00 2001
From: Jerry Yu <jerry.h.yu@arm.com>
Date: Thu, 24 Feb 2022 15:52:15 +0800
Subject: [PATCH] add mbedtls_pk_sign_ext

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
---
 include/mbedtls/pk.h | 39 ++++++++++++++++++++++++
 library/pk.c         | 70 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index 9ad7a1df6..b5ca6ea5d 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -535,6 +535,45 @@ int mbedtls_pk_sign( mbedtls_pk_context *ctx, mbedtls_md_type_t md_alg,
              unsigned char *sig, size_t sig_size, size_t *sig_len,
              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng );
 
+/**
+ * \brief           Make signature with options, including padding if relevant.
+ *
+ * \param type      Signature type (inc. possible padding type) to verify
+ * \param options   Pointer to type-specific options, or NULL
+ * \param ctx       The PK context to use. It must have been set up
+ *                  with a private key.
+ * \param md_alg    Hash algorithm used (see notes)
+ * \param hash      Hash of the message to sign
+ * \param hash_len  Hash length
+ * \param sig       Place to write the signature.
+ *                  It must have enough room for the signature.
+ *                  #MBEDTLS_PK_SIGNATURE_MAX_SIZE is always enough.
+ *                  You may use a smaller buffer if it is large enough
+ *                  given the key type.
+ * \param sig_size  The size of the \p sig buffer in bytes.
+ * \param sig_len   On successful return,
+ *                  the number of bytes written to \p sig.
+ * \param f_rng     RNG function, must not be \c NULL.
+ * \param p_rng     RNG parameter
+ *
+ * \return          0 on success, or a specific error code.
+ *
+ * \note            For RSA keys, the default padding type is PKCS#1 v1.5.
+ *                  There is no interface in the PK module to make RSASSA-PSS
+ *                  signatures yet.
+ *
+ * \note            For RSA, md_alg may be MBEDTLS_MD_NONE if hash_len != 0.
+ *                  For ECDSA, md_alg may never be MBEDTLS_MD_NONE.
+ *
+ * \note            For RSA, md_alg may be MBEDTLS_MD_NONE if hash_len != 0.
+ *                  For ECDSA, md_alg may never be MBEDTLS_MD_NONE.
+ */
+int mbedtls_pk_sign_ext( mbedtls_pk_type_t type, const void *options,
+             mbedtls_pk_context *ctx, mbedtls_md_type_t md_alg,
+             const unsigned char *hash, size_t hash_len,
+             unsigned char *sig, size_t sig_size, size_t *sig_len,
+             int (*f_rng)(void *, unsigned char *, size_t), void *p_rng );
+
 /**
  * \brief           Restartable version of \c mbedtls_pk_sign()
  *
diff --git a/library/pk.c b/library/pk.c
index 79eccaada..cb2535401 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -518,6 +518,76 @@ int mbedtls_pk_sign( mbedtls_pk_context *ctx, mbedtls_md_type_t md_alg,
                                          f_rng, p_rng, NULL ) );
 }
 
+/*
+ * Make a signature with options
+ */
+int mbedtls_pk_sign_ext( mbedtls_pk_type_t type, const void *options,
+             mbedtls_pk_context *ctx, mbedtls_md_type_t md_alg,
+             const unsigned char *hash, size_t hash_len,
+             unsigned char *sig, size_t sig_size, size_t *sig_len,
+             int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
+{
+    PK_VALIDATE_RET( ctx != NULL );
+    PK_VALIDATE_RET( ( md_alg == MBEDTLS_MD_NONE && hash_len == 0 ) ||
+                     hash != NULL );
+    PK_VALIDATE_RET( sig != NULL );
+
+    *sig_len = 0;
+    if( ctx->pk_info == NULL )
+        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+
+    if( ! mbedtls_pk_can_do( ctx, type ) )
+        return( MBEDTLS_ERR_PK_TYPE_MISMATCH );
+
+    if( type != MBEDTLS_PK_RSASSA_PSS )
+    {
+        /* General case: no options */
+        if( options != NULL )
+            return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+
+        return( mbedtls_pk_sign_restartable( ctx, md_alg, hash, hash_len,
+                                             sig, sig_size, sig_len,
+                                             f_rng, p_rng, NULL ) );
+    }
+
+#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PKCS1_V21)
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    const mbedtls_pk_rsassa_pss_options *pss_opts;
+
+#if SIZE_MAX > UINT_MAX
+    if( md_alg == MBEDTLS_MD_NONE && UINT_MAX < hash_len )
+        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+#endif /* SIZE_MAX > UINT_MAX */
+
+    if( options == NULL )
+        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+
+    pss_opts = (const mbedtls_pk_rsassa_pss_options *) options;
+    if( sig_size < mbedtls_pk_get_len( ctx ) )
+        return( MBEDTLS_ERR_RSA_VERIFY_FAILED );
+
+    if( mbedtls_rsa_set_padding( mbedtls_pk_rsa( *ctx ),
+                                 MBEDTLS_RSA_PKCS_V21,
+                                 md_alg ) != 0 )
+    {
+        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+    }
+
+    *sig_len = mbedtls_pk_get_len( ctx );
+    ret = mbedtls_rsa_rsassa_pss_sign_ext( mbedtls_pk_rsa( *ctx ),
+                                           f_rng, p_rng,
+                                           md_alg,
+                                           (unsigned int) hash_len,
+                                           hash,pss_opts->expected_salt_len,
+                                           sig
+                                         );
+    return( ret );
+#else
+    return( MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE );
+#endif /* MBEDTLS_RSA_C && MBEDTLS_PKCS1_V21 */
+
+}
+
 /*
  * Decrypt message
  */