Merge pull request #315 from gilles-peskine-arm/pk_signature_max_size

Define MBEDTLS_PK_SIGNATURE_MAX_SIZE
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index d750004..6343563 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -102,6 +102,58 @@
 } mbedtls_pk_rsassa_pss_options;
 
 /**
+ * \brief           Maximum size of a signature made by mbedtls_pk_sign().
+ */
+/* We need to set MBEDTLS_PK_SIGNATURE_MAX_SIZE to the maximum signature
+ * size among the supported signature types. Do it by starting at 0,
+ * then incrementally increasing to be large enough for each supported
+ * signature mechanism.
+ *
+ * The resulting value can be 0, for example if MBEDTLS_ECDH_C is enabled
+ * (which allows the pk module to be included) but neither MBEDTLS_ECDSA_C
+ * nor MBEDTLS_RSA_C nor any opaque signature mechanism (PSA or RSA_ALT).
+ */
+#define MBEDTLS_PK_SIGNATURE_MAX_SIZE 0
+
+#if ( defined(MBEDTLS_RSA_C) || defined(MBEDTLS_PK_RSA_ALT_SUPPORT) ) && \
+    MBEDTLS_MPI_MAX_SIZE > MBEDTLS_PK_SIGNATURE_MAX_SIZE
+/* For RSA, the signature can be as large as the bignum module allows.
+ * For RSA_ALT, the signature size is not necessarily tied to what the
+ * bignum module can do, but in the absence of any specific setting,
+ * we use that (rsa_alt_sign_wrap in pk_wrap will check). */
+#undef MBEDTLS_PK_SIGNATURE_MAX_SIZE
+#define MBEDTLS_PK_SIGNATURE_MAX_SIZE MBEDTLS_MPI_MAX_SIZE
+#endif
+
+#if defined(MBEDTLS_ECDSA_C) &&                                 \
+    MBEDTLS_ECDSA_MAX_LEN > MBEDTLS_PK_SIGNATURE_MAX_SIZE
+/* For ECDSA, the ecdsa module exports a constant for the maximum
+ * signature size. */
+#undef MBEDTLS_PK_SIGNATURE_MAX_SIZE
+#define MBEDTLS_PK_SIGNATURE_MAX_SIZE MBEDTLS_ECDSA_MAX_LEN
+#endif
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if PSA_ASYMMETRIC_SIGNATURE_MAX_SIZE > MBEDTLS_PK_SIGNATURE_MAX_SIZE
+/* PSA_ASYMMETRIC_SIGNATURE_MAX_SIZE is the maximum size of a signature made
+ * through the PSA API in the PSA representation. */
+#undef MBEDTLS_PK_SIGNATURE_MAX_SIZE
+#define MBEDTLS_PK_SIGNATURE_MAX_SIZE PSA_ASYMMETRIC_SIGNATURE_MAX_SIZE
+#endif
+
+#if PSA_VENDOR_ECDSA_SIGNATURE_MAX_SIZE + 11 > MBEDTLS_PK_SIGNATURE_MAX_SIZE
+/* The Mbed TLS representation is different for ECDSA signatures:
+ * PSA uses the raw concatenation of r and s,
+ * whereas Mbed TLS uses the ASN.1 representation (SEQUENCE of two INTEGERs).
+ * Add the overhead of ASN.1: up to (1+2) + 2 * (1+2+1) for the
+ * types, lengths (represented by up to 2 bytes), and potential leading
+ * zeros of the INTEGERs and the SEQUENCE. */
+#undef MBEDTLS_PK_SIGNATURE_MAX_SIZE
+#define MBEDTLS_PK_SIGNATURE_MAX_SIZE ( PSA_VENDOR_ECDSA_SIGNATURE_MAX_SIZE + 11 )
+#endif
+#endif /* defined(MBEDTLS_USE_PSA_CRYPTO) */
+
+/**
  * \brief           Types for interfacing with the debug module
  */
 typedef enum
@@ -442,8 +494,13 @@
  * \param md_alg    Hash algorithm used (see notes)
  * \param hash      Hash of the message to sign
  * \param hash_len  Hash length or 0 (see notes)
- * \param sig       Place to write the signature
- * \param sig_len   Number of bytes written
+ * \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_len   On successful return,
+ *                  the number of bytes written to \p sig.
  * \param f_rng     RNG function
  * \param p_rng     RNG parameter
  *
@@ -474,16 +531,21 @@
  *
  * \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 md_alg    Hash algorithm used (see notes for mbedtls_pk_sign())
  * \param hash      Hash of the message to sign
- * \param hash_len  Hash length or 0 (see notes)
- * \param sig       Place to write the signature
- * \param sig_len   Number of bytes written
+ * \param hash_len  Hash length or 0 (see notes for mbedtls_pk_sign())
+ * \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_len   On successful return,
+ *                  the number of bytes written to \p sig.
  * \param f_rng     RNG function
  * \param p_rng     RNG parameter
  * \param rs_ctx    Restart context (NULL to disable restart)
  *
- * \return          See \c mbedtls_pk_sign(), or
+ * \return          See \c mbedtls_pk_sign().
  * \return          #MBEDTLS_ERR_ECP_IN_PROGRESS if maximum number of
  *                  operations was reached: see \c mbedtls_ecp_set_max_ops().
  */
diff --git a/include/psa/crypto_sizes.h b/include/psa/crypto_sizes.h
index bcca724..3332247 100644
--- a/include/psa/crypto_sizes.h
+++ b/include/psa/crypto_sizes.h
@@ -247,21 +247,6 @@
  */
 #define PSA_ALG_TLS12_PSK_TO_MS_MAX_PSK_LEN 128
 
-/** \def PSA_ASYMMETRIC_SIGNATURE_MAX_SIZE
- *
- * Maximum size of an asymmetric signature.
- *
- * This macro must expand to a compile-time constant integer. This value
- * should be the maximum size of a MAC supported by the implementation,
- * in bytes, and must be no smaller than this maximum.
- */
-#define PSA_ASYMMETRIC_SIGNATURE_MAX_SIZE                               \
-    PSA_BITS_TO_BYTES(                                                  \
-        PSA_VENDOR_RSA_MAX_KEY_BITS > PSA_VENDOR_ECC_MAX_CURVE_BITS ?   \
-        PSA_VENDOR_RSA_MAX_KEY_BITS :                                   \
-        PSA_VENDOR_ECC_MAX_CURVE_BITS                                   \
-        )
-
 /** The maximum size of a block cipher supported by the implementation. */
 #define PSA_MAX_BLOCK_CIPHER_BLOCK_SIZE 16
 
@@ -457,6 +442,22 @@
      PSA_KEY_TYPE_IS_ECC(key_type) ? PSA_ECDSA_SIGNATURE_SIZE(key_bits) : \
      ((void)alg, 0))
 
+#define PSA_VENDOR_ECDSA_SIGNATURE_MAX_SIZE     \
+    PSA_ECDSA_SIGNATURE_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)
+
+/** \def PSA_ASYMMETRIC_SIGNATURE_MAX_SIZE
+ *
+ * Maximum size of an asymmetric signature.
+ *
+ * This macro must expand to a compile-time constant integer. This value
+ * should be the maximum size of a signature supported by the implementation,
+ * in bytes, and must be no smaller than this maximum.
+ */
+#define PSA_ASYMMETRIC_SIGNATURE_MAX_SIZE                               \
+    (PSA_BITS_TO_BYTES(PSA_VENDOR_RSA_MAX_KEY_BITS) > PSA_VENDOR_ECDSA_SIGNATURE_MAX_SIZE ? \
+     PSA_BITS_TO_BYTES(PSA_VENDOR_RSA_MAX_KEY_BITS) :                   \
+     PSA_VENDOR_ECDSA_SIGNATURE_MAX_SIZE)
+
 /** Sufficient output buffer size for psa_asymmetric_encrypt().
  *
  * This macro returns a sufficient buffer size for a ciphertext produced using
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 5a699c0..702c3bb 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -774,6 +774,8 @@
 #endif /* SIZE_MAX > UINT_MAX */
 
     *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 );
 
     return( rsa_alt->sign_func( rsa_alt->key, f_rng, p_rng, MBEDTLS_RSA_PRIVATE,
                 md_alg, (unsigned int) hash_len, hash, sig ) );
@@ -1017,6 +1019,8 @@
         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_asymmetric_sign( *key, alg, hash, hash_len,
diff --git a/programs/pkey/pk_sign.c b/programs/pkey/pk_sign.c
index 47a098a..79fb273 100644
--- a/programs/pkey/pk_sign.c
+++ b/programs/pkey/pk_sign.c
@@ -70,7 +70,7 @@
     mbedtls_entropy_context entropy;
     mbedtls_ctr_drbg_context ctr_drbg;
     unsigned char hash[32];
-    unsigned char buf[MBEDTLS_MPI_MAX_SIZE];
+    unsigned char buf[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
     char filename[512];
     const char *pers = "mbedtls_pk_sign";
     size_t olen = 0;
diff --git a/programs/pkey/pk_verify.c b/programs/pkey/pk_verify.c
index a6bfe3f..72caf71 100644
--- a/programs/pkey/pk_verify.c
+++ b/programs/pkey/pk_verify.c
@@ -65,7 +65,7 @@
     size_t i;
     mbedtls_pk_context pk;
     unsigned char hash[32];
-    unsigned char buf[MBEDTLS_MPI_MAX_SIZE];
+    unsigned char buf[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
     char filename[512];
 
     mbedtls_pk_init( &pk );
diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index ea5fc4f..caa4c77 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -8,21 +8,41 @@
 depends_on:MBEDTLS_RSA_C
 valid_parameters_pkwrite:"308204a20201000282010100a9021f3d406ad555538bfd36ee82652e15615e89bfb8e84590dbee881652d3f143504796125964876bfd2be046f973beddcf92e1915bed66a06f8929794580d0836ad54143775f397c09044782b0573970eda3ec15191ea8330847c10542a9fd4cc3b4dfdd061f4d1051406773130f40f86d81255f0ab153c6307e1539acf95aee7f929ea6055be7139785b52392d9d42406d50925897507dda61a8f3f0919bead652c64eb959bdcfe415e17a6da6c5b69cc02ba142c16249c4adccdd0f7526773f12da023fd7ef431ca2d70ca890b04db2ea64f706e9ecebd5889e253599e6e5a9265e2883f0c9419a3dde5e89d9513ed29dbab7012dc5aca6b17ab528254b10203010001028201001689f5e89142ae18a6ffb0513715a4b0b4a13b9e5b3729a2bd62d738c6e15cea7bf3a4d85ab2193a0628c9452bb1f0c1af8b132789df1c95e72778bf5330f5b0d915d242d5e0818e85001ed5fa93d1ce13455deb0a15438562e8e3c8d60ec1e4c9ebff9f2b36b9cde9332cc79f0d17a7ae79cc1353cd75409ad9b4b6d7ee3d82af6f3207656cf2ac98947c15c398db0cebf8dc3eef5398269480cdd09411b960273ae3f364da09af849f24aa87346c58618ea91d9d6cd1d3932c80dbfc1f0a4166a9036911999ca27761079f0ce02db02c1c909ff9b4278578d7bb1b54b2b7082fc9e864b6b394e331c0d11a9a68255565b6dd477f4119c5809839520700711102818100d7db987ad86de6a9b0749fb5da80bacde3bebd72dcc83f60a27db74f927ac3661386577bfce5b4a00ad024682401d6aad29713c8e223b53415305ca07559821099b187fdd1bad3dc4dec9da96f5fa6128331e8f7d89f1e1a788698d1a27256dc7cd392f04e531a9e38e7265bf4fd7eec01e7835e9b1a0dd8923e440381be1c2702818100c87025fff7a493c623404966fbc8b32ed164ca620ad1a0ad11ef42fd12118456017856a8b42e5d4ad36104e9dc9f8a2f3003c3957ffddb20e2f4e3fc3cf2cdddae01f57a56de4fd24b91ab6d3e5cc0e8af0473659594a6bbfdaacf958f19c8d508eac12d8977616af6877106288093d37904a139220c1bc278ea56edc086976702818043e708685c7cf5fa9b4f948e1856366d5e1f3a694f9a8e954f884c89f3823ac5798ee12657bfcaba2dac9c47464c6dc2fecc17a531be19da706fee336bb6e47b645dbc71d3eff9856bddeb1ac9b644ffbdd58d7ba9e1240f1faaf797ba8a4d58becbaf85789e1bd979fcfccc209d3db7f0416bc9eef09b3a6d86b8ce8199d4310281804f4b86ccffe49d0d8ace98fb63ea9f708b284ba483d130b6a75cb76cb4e4372d6b41774f20912319420ca4cbfc1b25a8cb5f01d6381f6ebc50ed3ef08010327f5ba2acc1ac7220b3fa6f7399314db2879b0db0b5647abd87abb01295815a5b086491b2c0d81c616ed67ef8a8ce0727f446711d7323d4147b5828a52143c43b4b028180540756beba83c20a0bda11d6dec706a71744ff28090cec079dffb507d82828038fe657f61496a20317f779cb683ce8196c29a6fe28839a282eef4de57773be56808b0c3e2ac7747e2b200b2fbf20b55258cd24622a1ce0099de098ab0855106ae087f08b0c8c346d81619400c1b4838e33ed9ff90f05db8fccf8fb7ab881ca12"
 
-PK utils: RSA
+PK utils: RSA 512-bit
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME
-pk_utils:MBEDTLS_PK_RSA:512:64:"RSA"
+pk_utils:MBEDTLS_PK_RSA:512:512:64:"RSA"
 
-PK utils: ECKEY
+PK utils: ECKEY SECP192R1
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP192R1_ENABLED
-pk_utils:MBEDTLS_PK_ECKEY:192:24:"EC"
+pk_utils:MBEDTLS_PK_ECKEY:MBEDTLS_ECP_DP_SECP192R1:192:24:"EC"
 
-PK utils: ECKEY_DH
+PK utils: ECKEY_DH SECP192R1
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP192R1_ENABLED
-pk_utils:MBEDTLS_PK_ECKEY_DH:192:24:"EC_DH"
+pk_utils:MBEDTLS_PK_ECKEY_DH:MBEDTLS_ECP_DP_SECP192R1:192:24:"EC_DH"
 
-PK utils: ECDSA
+PK utils: ECKEY_DH Curve25519
+depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+pk_utils:MBEDTLS_PK_ECKEY_DH:MBEDTLS_ECP_DP_CURVE25519:255:32:"EC_DH"
+
+PK utils: ECKEY_DH Curve448
+depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_CURVE448_ENABLED
+pk_utils:MBEDTLS_PK_ECKEY_DH:MBEDTLS_ECP_DP_CURVE448:448:56:"EC_DH"
+
+PK utils: ECDSA SECP192R1
 depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP192R1_ENABLED
-pk_utils:MBEDTLS_PK_ECDSA:192:24:"ECDSA"
+pk_utils:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP192R1:192:24:"ECDSA"
+
+PK utils: ECDSA SECP256R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+pk_utils:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:256:32:"ECDSA"
+
+PK utils: ECDSA SECP384R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP384R1_ENABLED
+pk_utils:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP384R1:384:48:"ECDSA"
+
+PK utils: ECDSA SECP521R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP521R1_ENABLED
+pk_utils:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP521R1:521:66:"ECDSA"
 
 PK PSA utilities: setup/free, info functions, unsupported operations
 pk_psa_utils:
@@ -83,21 +103,41 @@
 depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
 pk_ec_test_vec:MBEDTLS_PK_ECKEY:MBEDTLS_ECP_DP_SECP256R1:"0437cc56d976091e5a723ec7592dff206eee7cf9069174d0ad14b5f768225962924ee500d82311ffea2fd2345d5d16bd8a88c26b770d55cd8a2a0efa01c8b4edff":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855":"30430220685a6994daa6a14e4411b5267edc2a00beee907f2dddd956b2a5a1df791c15f8021f675db4538c000c734489ac737fddd5a739c5a23cd6c6eceea70c286ca4fac9":0
 
-ECDSA sign-verify
+ECDSA sign-verify: SECP192R1
 depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP192R1_ENABLED
-pk_sign_verify:MBEDTLS_PK_ECDSA:0:0
+pk_sign_verify:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP192R1:0:0
 
-EC(DSA) sign-verify
+ECDSA sign-verify: SECP256R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+pk_sign_verify:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:0:0
+
+ECDSA sign-verify: SECP384R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP384R1_ENABLED
+pk_sign_verify:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP384R1:0:0
+
+ECDSA sign-verify: SECP521R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP521R1_ENABLED
+pk_sign_verify:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP521R1:0:0
+
+ECDSA sign-verify: BP256R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_BP256R1_ENABLED
+pk_sign_verify:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_BP256R1:0:0
+
+ECDSA sign-verify: BP512R1
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_BP512R1_ENABLED
+pk_sign_verify:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_BP512R1:0:0
+
+EC(DSA) sign-verify: SECP192R1
 depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP192R1_ENABLED
-pk_sign_verify:MBEDTLS_PK_ECKEY:0:0
+pk_sign_verify:MBEDTLS_PK_ECKEY:MBEDTLS_ECP_DP_SECP192R1:0:0
 
-EC_DH (no) sign-verify
+EC_DH (no) sign-verify: SECP192R1
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP192R1_ENABLED
-pk_sign_verify:MBEDTLS_PK_ECKEY_DH:MBEDTLS_ERR_PK_TYPE_MISMATCH:MBEDTLS_ERR_PK_TYPE_MISMATCH
+pk_sign_verify:MBEDTLS_PK_ECKEY_DH:MBEDTLS_ECP_DP_SECP192R1:MBEDTLS_ERR_PK_TYPE_MISMATCH:MBEDTLS_ERR_PK_TYPE_MISMATCH
 
 RSA sign-verify
 depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15:MBEDTLS_GENPRIME
-pk_sign_verify:MBEDTLS_PK_RSA:0:0
+pk_sign_verify:MBEDTLS_PK_RSA:512:0:0
 
 RSA encrypt test vector
 depends_on:MBEDTLS_PKCS1_V15
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index b349075..ccf1736 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -27,13 +27,27 @@
 #define RSA_KEY_SIZE 512
 #define RSA_KEY_LEN   64
 
-static int pk_genkey( mbedtls_pk_context *pk )
+/** Generate a key of the desired type.
+ *
+ * \param pk        The PK object to fill. It must have been initialized
+ *                  with mbedtls_pk_setup().
+ * \param parameter - For RSA keys, the key size in bits.
+ *                  - For EC keys, the curve (\c MBEDTLS_ECP_DP_xxx).
+ *
+ * \return          The status from the underlying type-specific key
+ *                  generation function.
+ * \return          -1 if the key type is not recognized.
+ */
+static int pk_genkey( mbedtls_pk_context *pk, int parameter )
 {
     ((void) pk);
+    (void) parameter;
 
 #if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
     if( mbedtls_pk_get_type( pk ) == MBEDTLS_PK_RSA )
-        return mbedtls_rsa_gen_key( mbedtls_pk_rsa( *pk ), rnd_std_rand, NULL, RSA_KEY_SIZE, 3 );
+        return mbedtls_rsa_gen_key( mbedtls_pk_rsa( *pk ),
+                                    rnd_std_rand, NULL,
+                                    parameter, 3 );
 #endif
 #if defined(MBEDTLS_ECP_C)
     if( mbedtls_pk_get_type( pk ) == MBEDTLS_PK_ECKEY ||
@@ -42,7 +56,7 @@
     {
         int ret;
         if( ( ret = mbedtls_ecp_group_load( &mbedtls_pk_ec( *pk )->grp,
-                                      MBEDTLS_ECP_DP_SECP192R1 ) ) != 0 )
+                                            parameter ) ) != 0 )
             return( ret );
 
         return mbedtls_ecp_gen_keypair( &mbedtls_pk_ec( *pk )->grp, &mbedtls_pk_ec( *pk )->d,
@@ -608,18 +622,18 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void pk_utils( int type, int size, int len, char * name )
+void pk_utils( int type, int parameter, int bitlen, int len, char * name )
 {
     mbedtls_pk_context pk;
 
     mbedtls_pk_init( &pk );
 
     TEST_ASSERT( mbedtls_pk_setup( &pk, mbedtls_pk_info_from_type( type ) ) == 0 );
-    TEST_ASSERT( pk_genkey( &pk ) == 0 );
+    TEST_ASSERT( pk_genkey( &pk, parameter ) == 0 );
 
     TEST_ASSERT( (int) mbedtls_pk_get_type( &pk ) == type );
     TEST_ASSERT( mbedtls_pk_can_do( &pk, type ) );
-    TEST_ASSERT( mbedtls_pk_get_bitlen( &pk ) == (unsigned) size );
+    TEST_ASSERT( mbedtls_pk_get_bitlen( &pk ) == (unsigned) bitlen );
     TEST_ASSERT( mbedtls_pk_get_len( &pk ) == (unsigned) len );
     TEST_ASSERT( strcmp( mbedtls_pk_get_name( &pk), name ) == 0 );
 
@@ -897,11 +911,12 @@
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_SHA256_C */
-void pk_sign_verify( int type, int sign_ret, int verify_ret )
+void pk_sign_verify( int type, int parameter, int sign_ret, int verify_ret )
 {
     mbedtls_pk_context pk;
-    unsigned char hash[50], sig[5000];
     size_t sig_len;
+    unsigned char hash[MBEDTLS_MD_MAX_SIZE];
+    unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
     void *rs_ctx = NULL;
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     mbedtls_pk_restart_ctx ctx;
@@ -921,11 +936,15 @@
     memset( sig, 0, sizeof sig );
 
     TEST_ASSERT( mbedtls_pk_setup( &pk, mbedtls_pk_info_from_type( type ) ) == 0 );
-    TEST_ASSERT( pk_genkey( &pk ) == 0 );
+    TEST_ASSERT( pk_genkey( &pk, parameter ) == 0 );
 
     TEST_ASSERT( mbedtls_pk_sign_restartable( &pk, MBEDTLS_MD_SHA256,
                  hash, sizeof hash, sig, &sig_len,
                  rnd_std_rand, NULL, rs_ctx ) == sign_ret );
+    if( sign_ret == 0 )
+        TEST_ASSERT( sig_len <= MBEDTLS_PK_SIGNATURE_MAX_SIZE );
+    else
+        sig_len = MBEDTLS_PK_SIGNATURE_MAX_SIZE;
 
     TEST_ASSERT( mbedtls_pk_verify( &pk, MBEDTLS_MD_SHA256,
                             hash, sizeof hash, sig, sig_len ) == verify_ret );
@@ -945,6 +964,10 @@
 
     TEST_ASSERT( mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, sizeof hash,
                           sig, &sig_len, rnd_std_rand, NULL ) == sign_ret );
+    if( sign_ret == 0 )
+        TEST_ASSERT( sig_len <= MBEDTLS_PK_SIGNATURE_MAX_SIZE );
+    else
+        sig_len = MBEDTLS_PK_SIGNATURE_MAX_SIZE;
 
     TEST_ASSERT( mbedtls_pk_verify_restartable( &pk, MBEDTLS_MD_SHA256,
                  hash, sizeof hash, sig, sig_len, rs_ctx ) == verify_ret );
@@ -1153,7 +1176,7 @@
     /* Initiliaze PK RSA context with random key */
     TEST_ASSERT( mbedtls_pk_setup( &rsa,
                               mbedtls_pk_info_from_type( MBEDTLS_PK_RSA ) ) == 0 );
-    TEST_ASSERT( pk_genkey( &rsa ) == 0 );
+    TEST_ASSERT( pk_genkey( &rsa, RSA_KEY_SIZE ) == 0 );
 
     /* Extract key to the raw rsa context */
     TEST_ASSERT( mbedtls_rsa_copy( &raw, mbedtls_pk_rsa( rsa ) ) == 0 );