Merge pull request #4694 from gilles-peskine-arm/out_size-3.0

Add output size parameter to signature functions
diff --git a/ChangeLog.d/out_size.txt b/ChangeLog.d/out_size.txt
new file mode 100644
index 0000000..721bf6a
--- /dev/null
+++ b/ChangeLog.d/out_size.txt
@@ -0,0 +1,5 @@
+API changes
+   * The functions mbedtls_pk_sign(), mbedtls_pk_sign_restartable(),
+     mbedtls_ecdsa_write_signature() and
+     mbedtls_ecdsa_write_signature_restartable() now take an extra parameter
+     indicating the size of the output buffer for the signature.
diff --git a/docs/3.0-migration-guide.d/out_size.md b/docs/3.0-migration-guide.d/out_size.md
new file mode 100644
index 0000000..49d3246
--- /dev/null
+++ b/docs/3.0-migration-guide.d/out_size.md
@@ -0,0 +1,9 @@
+Extra parameter for the output buffer size
+------------------------------------------
+
+The following functions now take an extra parameter indicating the size of the output buffer:
+
+* `mbedtls_ecdsa_write_signature()`, `mbedtls_ecdsa_write_signature_restartable()`
+* `mbedtls_pk_sign()`, `mbedtls_pk_sign_restartable()`
+
+The requirements for the output buffer have not changed, but passing a buffer that is too small now reliably causes the functions to return an error, rather than overflowing the buffer.
diff --git a/include/mbedtls/ecdsa.h b/include/mbedtls/ecdsa.h
index 4586ddb..71b73ee 100644
--- a/include/mbedtls/ecdsa.h
+++ b/include/mbedtls/ecdsa.h
@@ -290,6 +290,7 @@
  *                  size of the curve used, plus 9. For example, 73 Bytes if
  *                  a 256-bit curve is used. A buffer length of
  *                  #MBEDTLS_ECDSA_MAX_LEN is always safe.
+ * \param sig_size  The size of the \p sig buffer in bytes.
  * \param slen      The address at which to store the actual length of
  *                  the signature written. Must not be \c NULL.
  * \param f_rng     The RNG function. This must not be \c NULL if
@@ -306,7 +307,7 @@
 int mbedtls_ecdsa_write_signature( mbedtls_ecdsa_context *ctx,
                                    mbedtls_md_type_t md_alg,
                            const unsigned char *hash, size_t hlen,
-                           unsigned char *sig, size_t *slen,
+                           unsigned char *sig, size_t sig_size, size_t *slen,
                            int (*f_rng)(void *, unsigned char *, size_t),
                            void *p_rng );
 
@@ -332,6 +333,7 @@
  *                  size of the curve used, plus 9. For example, 73 Bytes if
  *                  a 256-bit curve is used. A buffer length of
  *                  #MBEDTLS_ECDSA_MAX_LEN is always safe.
+ * \param sig_size  The size of the \p sig buffer in bytes.
  * \param slen      The address at which to store the actual length of
  *                  the signature written. Must not be \c NULL.
  * \param f_rng     The RNG function. This must not be \c NULL if
@@ -352,7 +354,7 @@
 int mbedtls_ecdsa_write_signature_restartable( mbedtls_ecdsa_context *ctx,
                            mbedtls_md_type_t md_alg,
                            const unsigned char *hash, size_t hlen,
-                           unsigned char *sig, size_t *slen,
+                           unsigned char *sig, size_t sig_size, size_t *slen,
                            int (*f_rng)(void *, unsigned char *, size_t),
                            void *p_rng,
                            mbedtls_ecdsa_restart_ctx *rs_ctx );
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index 9bbb2f6..f3553cc 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -63,6 +63,7 @@
 #define MBEDTLS_ERR_PK_UNKNOWN_NAMED_CURVE -0x3A00  /**< Elliptic curve is unsupported (only NIST curves are supported). */
 #define MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE -0x3980  /**< Unavailable feature, e.g. RSA disabled for RSA key. */
 #define MBEDTLS_ERR_PK_SIG_LEN_MISMATCH    -0x3900  /**< The buffer contains a valid signature followed by more data. */
+#define MBEDTLS_ERR_PK_BUFFER_TOO_SMALL    -0x3880  /**< The output buffer is too small. */
 
 #ifdef __cplusplus
 extern "C" {
@@ -495,6 +496,7 @@
  *                  #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.
@@ -511,7 +513,7 @@
  */
 int mbedtls_pk_sign( mbedtls_pk_context *ctx, mbedtls_md_type_t md_alg,
              const unsigned char *hash, size_t hash_len,
-             unsigned char *sig, size_t *sig_len,
+             unsigned char *sig, size_t sig_size, size_t *sig_len,
              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng );
 
 /**
@@ -532,6 +534,7 @@
  *                  #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.
@@ -545,7 +548,7 @@
 int mbedtls_pk_sign_restartable( mbedtls_pk_context *ctx,
              mbedtls_md_type_t md_alg,
              const unsigned char *hash, size_t hash_len,
-             unsigned char *sig, size_t *sig_len,
+             unsigned char *sig, size_t sig_size, size_t *sig_len,
              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
              mbedtls_pk_restart_ctx *rs_ctx );
 
diff --git a/library/ecdsa.c b/library/ecdsa.c
index 9220633..0b612ce 100644
--- a/library/ecdsa.c
+++ b/library/ecdsa.c
@@ -648,7 +648,8 @@
  * Convert a signature (given by context) to ASN.1
  */
 static int ecdsa_signature_to_asn1( const mbedtls_mpi *r, const mbedtls_mpi *s,
-                                    unsigned char *sig, size_t *slen )
+                                    unsigned char *sig, size_t sig_size,
+                                    size_t *slen )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     unsigned char buf[MBEDTLS_ECDSA_MAX_LEN] = {0};
@@ -662,6 +663,9 @@
     MBEDTLS_ASN1_CHK_ADD( len, mbedtls_asn1_write_tag( &p, buf,
                                        MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE ) );
 
+    if( len > sig_size )
+        return( MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL );
+
     memcpy( sig, p, len );
     *slen = len;
 
@@ -674,7 +678,7 @@
 int mbedtls_ecdsa_write_signature_restartable( mbedtls_ecdsa_context *ctx,
                            mbedtls_md_type_t md_alg,
                            const unsigned char *hash, size_t hlen,
-                           unsigned char *sig, size_t *slen,
+                           unsigned char *sig, size_t sig_size, size_t *slen,
                            int (*f_rng)(void *, unsigned char *, size_t),
                            void *p_rng,
                            mbedtls_ecdsa_restart_ctx *rs_ctx )
@@ -712,7 +716,7 @@
 #endif /* MBEDTLS_ECDSA_SIGN_ALT */
 #endif /* MBEDTLS_ECDSA_DETERMINISTIC */
 
-    MBEDTLS_MPI_CHK( ecdsa_signature_to_asn1( &r, &s, sig, slen ) );
+    MBEDTLS_MPI_CHK( ecdsa_signature_to_asn1( &r, &s, sig, sig_size, slen ) );
 
 cleanup:
     mbedtls_mpi_free( &r );
@@ -727,7 +731,7 @@
 int mbedtls_ecdsa_write_signature( mbedtls_ecdsa_context *ctx,
                                  mbedtls_md_type_t md_alg,
                                  const unsigned char *hash, size_t hlen,
-                                 unsigned char *sig, size_t *slen,
+                                 unsigned char *sig, size_t sig_size, size_t *slen,
                                  int (*f_rng)(void *, unsigned char *, size_t),
                                  void *p_rng )
 {
@@ -736,7 +740,8 @@
     ECDSA_VALIDATE_RET( sig  != NULL );
     ECDSA_VALIDATE_RET( slen != NULL );
     return( mbedtls_ecdsa_write_signature_restartable(
-                ctx, md_alg, hash, hlen, sig, slen, f_rng, p_rng, NULL ) );
+                ctx, md_alg, hash, hlen, sig, sig_size, slen,
+                f_rng, p_rng, NULL ) );
 }
 
 /*
diff --git a/library/pk.c b/library/pk.c
index 275d34b..ea4869c 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -396,7 +396,7 @@
 int mbedtls_pk_sign_restartable( mbedtls_pk_context *ctx,
              mbedtls_md_type_t md_alg,
              const unsigned char *hash, size_t hash_len,
-             unsigned char *sig, size_t *sig_len,
+             unsigned char *sig, size_t sig_size, size_t *sig_len,
              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
              mbedtls_pk_restart_ctx *rs_ctx )
 {
@@ -421,7 +421,9 @@
             return( ret );
 
         ret = ctx->pk_info->sign_rs_func( ctx->pk_ctx, md_alg,
-                hash, hash_len, sig, sig_len, f_rng, p_rng, rs_ctx->rs_ctx );
+                                          hash, hash_len,
+                                          sig, sig_size, sig_len,
+                                          f_rng, p_rng, rs_ctx->rs_ctx );
 
         if( ret != MBEDTLS_ERR_ECP_IN_PROGRESS )
             mbedtls_pk_restart_free( rs_ctx );
@@ -435,8 +437,10 @@
     if( ctx->pk_info->sign_func == NULL )
         return( MBEDTLS_ERR_PK_TYPE_MISMATCH );
 
-    return( ctx->pk_info->sign_func( ctx->pk_ctx, md_alg, hash, hash_len,
-                                     sig, sig_len, f_rng, p_rng ) );
+    return( ctx->pk_info->sign_func( ctx->pk_ctx, md_alg,
+                                     hash, hash_len,
+                                     sig, sig_size, sig_len,
+                                     f_rng, p_rng ) );
 }
 
 /*
@@ -444,11 +448,12 @@
  */
 int mbedtls_pk_sign( mbedtls_pk_context *ctx, mbedtls_md_type_t md_alg,
              const unsigned char *hash, size_t hash_len,
-             unsigned char *sig, size_t *sig_len,
+             unsigned char *sig, size_t sig_size, size_t *sig_len,
              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     return( mbedtls_pk_sign_restartable( ctx, md_alg, hash, hash_len,
-                                         sig, sig_len, f_rng, p_rng, NULL ) );
+                                         sig, sig_size, sig_len,
+                                         f_rng, p_rng, NULL ) );
 }
 
 /*
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 864e495..80c0aad 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -108,7 +108,7 @@
 
 static int rsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     mbedtls_rsa_context * rsa = (mbedtls_rsa_context *) ctx;
@@ -119,6 +119,8 @@
 #endif /* SIZE_MAX > UINT_MAX */
 
     *sig_len = mbedtls_rsa_get_len( rsa );
+    if( sig_size < *sig_len )
+        return( MBEDTLS_ERR_PK_BUFFER_TOO_SMALL );
 
     return( mbedtls_rsa_pkcs1_sign( rsa, f_rng, p_rng,
                                     md_alg, (unsigned int) hash_len,
@@ -247,7 +249,7 @@
 
 static int ecdsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng );
 
 static int eckey_verify_wrap( void *ctx, mbedtls_md_type_t md_alg,
@@ -269,7 +271,7 @@
 
 static int eckey_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
@@ -278,7 +280,8 @@
     mbedtls_ecdsa_init( &ecdsa );
 
     if( ( ret = mbedtls_ecdsa_from_keypair( &ecdsa, ctx ) ) == 0 )
-        ret = ecdsa_sign_wrap( &ecdsa, md_alg, hash, hash_len, sig, sig_len,
+        ret = ecdsa_sign_wrap( &ecdsa, md_alg, hash, hash_len,
+                               sig, sig_size, sig_len,
                                f_rng, p_rng );
 
     mbedtls_ecdsa_free( &ecdsa );
@@ -295,7 +298,7 @@
 
 static int ecdsa_sign_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
                    void *rs_ctx );
 
@@ -367,7 +370,7 @@
 
 static int eckey_sign_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
                        void *rs_ctx )
 {
@@ -383,7 +386,7 @@
         MBEDTLS_MPI_CHK( mbedtls_ecdsa_from_keypair( &rs->ecdsa_ctx, ctx ) );
 
     MBEDTLS_MPI_CHK( ecdsa_sign_rs_wrap( &rs->ecdsa_ctx, md_alg,
-                                         hash, hash_len, sig, sig_len,
+                                         hash, hash_len, sig, sig_size, sig_len,
                                          f_rng, p_rng, &rs->ecdsa_rs ) );
 
 cleanup:
@@ -652,11 +655,13 @@
 
 static int ecdsa_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     return( mbedtls_ecdsa_write_signature( (mbedtls_ecdsa_context *) ctx,
-                md_alg, hash, hash_len, sig, sig_len, f_rng, p_rng ) );
+                                           md_alg, hash, hash_len,
+                                           sig, sig_size, sig_len,
+                                           f_rng, p_rng ) );
 }
 
 #if defined(MBEDTLS_ECP_RESTARTABLE)
@@ -681,13 +686,13 @@
 
 static int ecdsa_sign_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
                    void *rs_ctx )
 {
     return( mbedtls_ecdsa_write_signature_restartable(
                 (mbedtls_ecdsa_context *) ctx,
-                md_alg, hash, hash_len, sig, sig_len, f_rng, p_rng,
+                md_alg, hash, hash_len, sig, sig_size, sig_len, f_rng, p_rng,
                 (mbedtls_ecdsa_restart_ctx *) rs_ctx ) );
 
 }
@@ -770,7 +775,7 @@
 
 static int rsa_alt_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     mbedtls_rsa_alt_context *rsa_alt = (mbedtls_rsa_alt_context *) ctx;
@@ -783,6 +788,8 @@
     *sig_len = rsa_alt->key_len_func( rsa_alt->key );
     if( *sig_len > MBEDTLS_PK_SIGNATURE_MAX_SIZE )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+    if( *sig_len > sig_size )
+        return( MBEDTLS_ERR_PK_BUFFER_TOO_SMALL );
 
     return( rsa_alt->sign_func( rsa_alt->key, f_rng, p_rng,
                 md_alg, (unsigned int) hash_len, hash, sig ) );
@@ -822,7 +829,8 @@
 
     if( ( ret = rsa_alt_sign_wrap( (void *) prv, MBEDTLS_MD_NONE,
                                    hash, sizeof( hash ),
-                                   sig, &sig_len, f_rng, p_rng ) ) != 0 )
+                                   sig, sizeof( sig ), &sig_len,
+                                   f_rng, p_rng ) ) != 0 )
     {
         return( ret );
     }
@@ -1010,7 +1018,7 @@
 
 static int pk_opaque_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
                    const unsigned char *hash, size_t hash_len,
-                   unsigned char *sig, size_t *sig_len,
+                   unsigned char *sig, size_t sig_size, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
 #if !defined(MBEDTLS_ECDSA_C)
@@ -1019,41 +1027,28 @@
     ((void) hash);
     ((void) hash_len);
     ((void) sig);
+    ((void) sig_size);
     ((void) sig_len);
     ((void) f_rng);
     ((void) p_rng);
     return( MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE );
 #else /* !MBEDTLS_ECDSA_C */
     const psa_key_id_t *key = (const psa_key_id_t *) ctx;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_algorithm_t alg = PSA_ALG_ECDSA( mbedtls_psa_translate_md( md_alg ) );
-    size_t buf_len;
     psa_status_t status;
 
     /* PSA has its own RNG */
     (void) f_rng;
     (void) p_rng;
 
-    /* PSA needs an output buffer of known size, but our API doesn't provide
-     * that information. Assume that the buffer is large enough for a
-     * maximal-length signature with that key (otherwise the application is
-     * buggy anyway). */
-    status = psa_get_key_attributes( *key, &attributes );
-    if( status != PSA_SUCCESS )
-        return( mbedtls_psa_err_translate_pk( status ) );
-    buf_len = MBEDTLS_ECDSA_MAX_SIG_LEN( psa_get_key_bits( &attributes ) );
-    psa_reset_key_attributes( &attributes );
-    if( buf_len > MBEDTLS_PK_SIGNATURE_MAX_SIZE )
-        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
-
     /* make the signature */
     status = psa_sign_hash( *key, alg, hash, hash_len,
-                            sig, buf_len, sig_len );
+                            sig, sig_size, sig_len );
     if( status != PSA_SUCCESS )
         return( mbedtls_psa_err_translate_pk( status ) );
 
     /* transcode it to ASN.1 sequence */
-    return( pk_ecdsa_sig_asn1_from_psa( sig, sig_len, buf_len ) );
+    return( pk_ecdsa_sig_asn1_from_psa( sig, sig_len, sig_size ) );
 #endif /* !MBEDTLS_ECDSA_C */
 }
 
diff --git a/library/pk_wrap.h b/library/pk_wrap.h
index eab9939..6f5addf 100644
--- a/library/pk_wrap.h
+++ b/library/pk_wrap.h
@@ -49,7 +49,7 @@
     /** Make signature */
     int (*sign_func)( void *ctx, mbedtls_md_type_t md_alg,
                       const unsigned char *hash, size_t hash_len,
-                      unsigned char *sig, size_t *sig_len,
+                      unsigned char *sig, size_t sig_size, size_t *sig_len,
                       int (*f_rng)(void *, unsigned char *, size_t),
                       void *p_rng );
 
@@ -63,7 +63,7 @@
     /** Make signature (restartable) */
     int (*sign_rs_func)( void *ctx, mbedtls_md_type_t md_alg,
                          const unsigned char *hash, size_t hash_len,
-                         unsigned char *sig, size_t *sig_len,
+                         unsigned char *sig, size_t sig_size, size_t *sig_len,
                          int (*f_rng)(void *, unsigned char *, size_t),
                          void *p_rng, void *rs_ctx );
 #endif /* MBEDTLS_ECDSA_C && MBEDTLS_ECP_RESTARTABLE */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5884bc7..d089fe0 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -266,6 +266,8 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         case MBEDTLS_ERR_PK_SIG_LEN_MISMATCH:
             return( PSA_ERROR_INVALID_SIGNATURE );
+        case MBEDTLS_ERR_PK_BUFFER_TOO_SMALL:
+            return( PSA_ERROR_BUFFER_TOO_SMALL );
 
         case MBEDTLS_ERR_PLATFORM_HW_ACCEL_FAILED:
             return( PSA_ERROR_HARDWARE_FAILURE );
diff --git a/library/ssl_cli.c b/library/ssl_cli.c
index 30e64c4..f641970 100644
--- a/library/ssl_cli.c
+++ b/library/ssl_cli.c
@@ -3948,6 +3948,11 @@
     mbedtls_md_type_t md_alg = MBEDTLS_MD_NONE;
     size_t hashlen;
     void *rs_ctx = NULL;
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t out_buf_len = ssl->out_buf_len - ( ssl->out_msg - ssl->out_buf );
+#else
+    size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN - ( ssl->out_msg - ssl->out_buf );
+#endif
 
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> write certificate verify" ) );
 
@@ -4046,7 +4051,9 @@
 
     if( ( ret = mbedtls_pk_sign_restartable( mbedtls_ssl_own_key( ssl ),
                          md_alg, hash_start, hashlen,
-                         ssl->out_msg + 6 + offset, &n,
+                         ssl->out_msg + 6 + offset,
+                         out_buf_len - 6 - offset,
+                         &n,
                          ssl->conf->f_rng, ssl->conf->p_rng, rs_ctx ) ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_pk_sign", ret );
diff --git a/library/ssl_srv.c b/library/ssl_srv.c
index 4715129..ab6d2d5 100644
--- a/library/ssl_srv.c
+++ b/library/ssl_srv.c
@@ -2922,6 +2922,14 @@
     (void) signature_len;
 #endif /* MBEDTLS_KEY_EXCHANGE_WITH_SERVER_SIGNATURE_ENABLED */
 
+#if defined(MBEDTLS_KEY_EXCHANGE_WITH_SERVER_SIGNATURE_ENABLED)
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t out_buf_len = ssl->out_buf_len - ( ssl->out_msg - ssl->out_buf );
+#else
+    size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN - ( ssl->out_msg - ssl->out_buf );
+#endif
+#endif
+
     ssl->out_msglen = 4; /* header (type:1, length:3) to be written later */
 
     /*
@@ -3224,6 +3232,7 @@
         if( ( ret = mbedtls_pk_sign( mbedtls_ssl_own_key( ssl ),
                                      md_alg, hash, hashlen,
                                      ssl->out_msg + ssl->out_msglen + 2,
+                                     out_buf_len - ssl->out_msglen - 2,
                                      signature_len,
                                      ssl->conf->f_rng,
                                      ssl->conf->p_rng ) ) != 0 )
diff --git a/library/x509write_crt.c b/library/x509write_crt.c
index c90f8bc..c8169f1 100644
--- a/library/x509write_crt.c
+++ b/library/x509write_crt.c
@@ -474,7 +474,7 @@
     }
 
     if( ( ret = mbedtls_pk_sign( ctx->issuer_key, ctx->md_alg,
-                                 hash, 0, sig, &sig_len,
+                                 hash, 0, sig, sizeof( sig ), &sig_len,
                                  f_rng, p_rng ) ) != 0 )
     {
         return( ret );
diff --git a/library/x509write_csr.c b/library/x509write_csr.c
index 9f0ad93..555f296 100644
--- a/library/x509write_csr.c
+++ b/library/x509write_csr.c
@@ -136,7 +136,7 @@
 static int x509write_csr_der_internal( mbedtls_x509write_csr *ctx,
                                  unsigned char *buf,
                                  size_t size,
-                                 unsigned char *sig,
+                                 unsigned char *sig, size_t sig_size,
                                  int (*f_rng)(void *, unsigned char *, size_t),
                                  void *p_rng )
 {
@@ -235,7 +235,8 @@
     if( ret != 0 )
         return( ret );
 #endif
-    if( ( ret = mbedtls_pk_sign( ctx->key, ctx->md_alg, hash, 0, sig, &sig_len,
+    if( ( ret = mbedtls_pk_sign( ctx->key, ctx->md_alg, hash, 0,
+                                 sig, sig_size, &sig_len,
                                  f_rng, p_rng ) ) != 0 )
     {
         return( ret );
@@ -304,7 +305,9 @@
         return( MBEDTLS_ERR_X509_ALLOC_FAILED );
     }
 
-    ret = x509write_csr_der_internal( ctx, buf, size, sig, f_rng, p_rng );
+    ret = x509write_csr_der_internal( ctx, buf, size,
+                                      sig, MBEDTLS_PK_SIGNATURE_MAX_SIZE,
+                                      f_rng, p_rng );
 
     mbedtls_free( sig );
 
diff --git a/programs/pkey/ecdsa.c b/programs/pkey/ecdsa.c
index f4756a4..6b6e951 100644
--- a/programs/pkey/ecdsa.c
+++ b/programs/pkey/ecdsa.c
@@ -180,7 +180,7 @@
 
     if( ( ret = mbedtls_ecdsa_write_signature( &ctx_sign, MBEDTLS_MD_SHA256,
                                        hash, sizeof( hash ),
-                                       sig, &sig_len,
+                                       sig, sizeof( sig ), &sig_len,
                                        mbedtls_ctr_drbg_random, &ctr_drbg ) ) != 0 )
     {
         mbedtls_printf( " failed\n  ! mbedtls_ecdsa_write_signature returned %d\n", ret );
diff --git a/programs/pkey/pk_sign.c b/programs/pkey/pk_sign.c
index 08c4ddc..2ab428d 100644
--- a/programs/pkey/pk_sign.c
+++ b/programs/pkey/pk_sign.c
@@ -119,8 +119,9 @@
         goto exit;
     }
 
-    if( ( ret = mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, 0, buf, &olen,
-                         mbedtls_ctr_drbg_random, &ctr_drbg ) ) != 0 )
+    if( ( ret = mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, 0,
+                                 buf, sizeof( buf ), &olen,
+                                 mbedtls_ctr_drbg_random, &ctr_drbg ) ) != 0 )
     {
         mbedtls_printf( " failed\n  ! mbedtls_pk_sign returned -0x%04x\n", (unsigned int) -ret );
         goto exit;
diff --git a/programs/pkey/rsa_sign_pss.c b/programs/pkey/rsa_sign_pss.c
index 91cb789..7ec7335 100644
--- a/programs/pkey/rsa_sign_pss.c
+++ b/programs/pkey/rsa_sign_pss.c
@@ -135,8 +135,9 @@
         goto exit;
     }
 
-    if( ( ret = mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, 0, buf, &olen,
-                         mbedtls_ctr_drbg_random, &ctr_drbg ) ) != 0 )
+    if( ( ret = mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, 0,
+                                 buf, sizeof( buf ), &olen,
+                                 mbedtls_ctr_drbg_random, &ctr_drbg ) ) != 0 )
     {
         mbedtls_printf( " failed\n  ! mbedtls_pk_sign returned %d\n\n", ret );
         goto exit;
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index cf608b9..d2aa48a 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -1131,7 +1131,7 @@
             ret = mbedtls_pk_sign( key_slot->pk,
                                    ctx->md_alg,
                                    ctx->input, ctx->input_len,
-                                   output, output_len,
+                                   output, output_size, output_len,
                                    config_data->f_rng, config_data->p_rng );
             break;
         default:
diff --git a/programs/test/benchmark.c b/programs/test/benchmark.c
index 6b3cc25..5985caf 100644
--- a/programs/test/benchmark.c
+++ b/programs/test/benchmark.c
@@ -1084,7 +1084,7 @@
                                               curve_info->name );
             TIME_PUBLIC( title, "sign",
                     ret = mbedtls_ecdsa_write_signature( &ecdsa, MBEDTLS_MD_SHA256, buf, curve_info->bit_size,
-                                                tmp, &sig_len, myrand, NULL ) );
+                                                tmp, sizeof( tmp ), &sig_len, myrand, NULL ) );
 
             mbedtls_ecdsa_free( &ecdsa );
         }
@@ -1100,7 +1100,7 @@
 
             if( mbedtls_ecdsa_genkey( &ecdsa, curve_info->grp_id, myrand, NULL ) != 0 ||
                 mbedtls_ecdsa_write_signature( &ecdsa, MBEDTLS_MD_SHA256, buf, curve_info->bit_size,
-                                               tmp, &sig_len, myrand, NULL ) != 0 )
+                                               tmp, sizeof( tmp ), &sig_len, myrand, NULL ) != 0 )
             {
                 mbedtls_exit( 1 );
             }
diff --git a/tests/suites/test_suite_ecdsa.function b/tests/suites/test_suite_ecdsa.function
index 8803762..4496812 100644
--- a/tests/suites/test_suite_ecdsa.function
+++ b/tests/suites/test_suite_ecdsa.function
@@ -205,7 +205,7 @@
     /* generate and write signature, then read and verify it */
     TEST_ASSERT( mbedtls_ecdsa_write_signature( &ctx, MBEDTLS_MD_SHA256,
                  hash, sizeof( hash ),
-                 sig, &sig_len, &mbedtls_test_rnd_pseudo_rand,
+                 sig, sizeof( sig ), &sig_len, &mbedtls_test_rnd_pseudo_rand,
                  &rnd_info ) == 0 );
     TEST_ASSERT( mbedtls_ecdsa_read_signature( &ctx, hash, sizeof( hash ),
                  sig, sig_len ) == 0 );
@@ -269,7 +269,7 @@
     /* generate and write signature, then read and verify it */
     TEST_ASSERT( mbedtls_ecdsa_write_signature( &ctx, MBEDTLS_MD_SHA256,
                  hash, sizeof( hash ),
-                 sig, &sig_len, &mbedtls_test_rnd_pseudo_rand,
+                 sig, sizeof( sig ), &sig_len, &mbedtls_test_rnd_pseudo_rand,
                  &rnd_info ) == 0 );
     TEST_ASSERT( mbedtls_ecdsa_read_signature( &ctx, hash, sizeof( hash ),
                  sig, sig_len ) == 0 );
@@ -404,8 +404,8 @@
     cnt_restart = 0;
     do {
         ret = mbedtls_ecdsa_write_signature_restartable( &ctx,
-                md_alg, hash, hlen, sig, &slen, mbedtls_test_rnd_std_rand, NULL,
-                &rs_ctx );
+                md_alg, hash, hlen, sig, sizeof( sig ), &slen,
+                mbedtls_test_rnd_std_rand, NULL, &rs_ctx );
     } while( ret == MBEDTLS_ERR_ECP_IN_PROGRESS && ++cnt_restart );
 
     TEST_ASSERT( ret == 0 );
@@ -420,8 +420,8 @@
     if( min_restart > 0 )
     {
         ret = mbedtls_ecdsa_write_signature_restartable( &ctx,
-                md_alg, hash, hlen, sig, &slen, mbedtls_test_rnd_std_rand, NULL,
-                &rs_ctx );
+                md_alg, hash, hlen, sig, sizeof( sig ), &slen,
+                mbedtls_test_rnd_std_rand, NULL, &rs_ctx );
         TEST_ASSERT( ret == MBEDTLS_ERR_ECP_IN_PROGRESS );
     }
 
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 6e93bc6..56cc45b 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -226,7 +226,7 @@
     TEST_ASSERT( mbedtls_pk_sign_restartable( &pk,
                                               MBEDTLS_MD_NONE,
                                               NULL, 0,
-                                              buf, &len,
+                                              buf, sizeof( buf ), &len,
                                               mbedtls_test_rnd_std_rand, NULL,
                                               NULL ) ==
                  MBEDTLS_ERR_PK_BAD_INPUT_DATA );
@@ -234,7 +234,7 @@
     TEST_ASSERT( mbedtls_pk_sign_restartable( &pk,
                                               MBEDTLS_MD_NONE,
                                               NULL, 0,
-                                              buf, &len,
+                                              buf, sizeof( buf ), &len,
                                               mbedtls_test_rnd_std_rand, NULL,
                                               NULL ) ==
                  MBEDTLS_ERR_PK_BAD_INPUT_DATA );
@@ -242,7 +242,7 @@
     TEST_ASSERT( mbedtls_pk_sign( &pk,
                                   MBEDTLS_MD_NONE,
                                   NULL, 0,
-                                  buf, &len,
+                                  buf, sizeof( buf ), &len,
                                   mbedtls_test_rnd_std_rand, NULL ) ==
                  MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
@@ -555,8 +555,9 @@
     cnt_restart = 0;
     do {
         ret = mbedtls_pk_sign_restartable( &prv, md_alg, hash, hlen,
-                                            sig, &slen, mbedtls_test_rnd_std_rand,
-                                            NULL, &rs_ctx );
+                                           sig, sizeof( sig ), &slen,
+                                           mbedtls_test_rnd_std_rand, NULL,
+                                           &rs_ctx );
     } while( ret == MBEDTLS_ERR_ECP_IN_PROGRESS && ++cnt_restart );
 
     TEST_ASSERT( ret == 0 );
@@ -603,8 +604,9 @@
 
         slen = sizeof( sig );
         ret = mbedtls_pk_sign_restartable( &prv, md_alg, hash, hlen,
-                                            sig, &slen, mbedtls_test_rnd_std_rand,
-                                            NULL, &rs_ctx );
+                                           sig, sizeof sig, &slen,
+                                           mbedtls_test_rnd_std_rand, NULL,
+                                           &rs_ctx );
         TEST_ASSERT( ret == MBEDTLS_ERR_ECP_IN_PROGRESS );
     }
 
@@ -645,8 +647,10 @@
     TEST_ASSERT( pk_genkey( &pk, parameter ) == 0 );
 
     TEST_ASSERT( mbedtls_pk_sign_restartable( &pk, MBEDTLS_MD_SHA256,
-                 hash, hash_len, sig, &sig_len,
-                 mbedtls_test_rnd_std_rand, NULL, rs_ctx ) == sign_ret );
+                                              hash, hash_len,
+                                              sig, sizeof sig, &sig_len,
+                                              mbedtls_test_rnd_std_rand, NULL,
+                                              rs_ctx ) == sign_ret );
     if( sign_ret == 0 )
         TEST_ASSERT( sig_len <= MBEDTLS_PK_SIGNATURE_MAX_SIZE );
     else
@@ -669,7 +673,7 @@
     }
 
     TEST_ASSERT( mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, hash_len,
-                                  sig, &sig_len,
+                                  sig, sizeof sig, &sig_len,
                                   mbedtls_test_rnd_std_rand,
                                   NULL ) == sign_ret );
     if( sign_ret == 0 )
@@ -848,8 +852,9 @@
     TEST_ASSERT( mbedtls_pk_verify( &pk, MBEDTLS_MD_NONE, hash, hash_len,
                     sig, sig_len ) == MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
-    TEST_ASSERT( mbedtls_pk_sign( &pk, MBEDTLS_MD_NONE, hash, hash_len, sig,
-                                  &sig_len, mbedtls_test_rnd_std_rand, NULL )
+    TEST_ASSERT( mbedtls_pk_sign( &pk, MBEDTLS_MD_NONE, hash, hash_len,
+                                  sig, sizeof sig, &sig_len,
+                                  mbedtls_test_rnd_std_rand, NULL )
                  == MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
 exit:
@@ -903,12 +908,14 @@
 
     /* Test signature */
 #if SIZE_MAX > UINT_MAX
-    TEST_ASSERT( mbedtls_pk_sign( &alt, MBEDTLS_MD_NONE, hash, SIZE_MAX, sig,
-                                  &sig_len, mbedtls_test_rnd_std_rand, NULL )
+    TEST_ASSERT( mbedtls_pk_sign( &alt, MBEDTLS_MD_NONE, hash, SIZE_MAX,
+                                  sig, sizeof sig, &sig_len,
+                                  mbedtls_test_rnd_std_rand, NULL )
                  == MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 #endif /* SIZE_MAX > UINT_MAX */
-    TEST_ASSERT( mbedtls_pk_sign( &alt, MBEDTLS_MD_NONE, hash, sizeof hash, sig,
-                                  &sig_len, mbedtls_test_rnd_std_rand, NULL )
+    TEST_ASSERT( mbedtls_pk_sign( &alt, MBEDTLS_MD_NONE, hash, sizeof hash,
+                                  sig, sizeof sig, &sig_len,
+                                  mbedtls_test_rnd_std_rand, NULL )
                  == 0 );
     TEST_ASSERT( sig_len == RSA_KEY_LEN );
     TEST_ASSERT( mbedtls_pk_verify( &rsa, MBEDTLS_MD_NONE,
@@ -996,7 +1003,7 @@
     memset( sig, 0, sizeof sig );
 
     TEST_ASSERT( mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256,
-                 hash, sizeof hash, sig, &sig_len,
+                 hash, sizeof hash, sig, sizeof sig, &sig_len,
                  NULL, NULL ) == 0 );
 
     /* Export underlying public key for re-importing in a psa context. */