Add tests for restartable PK sign/verify
diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index dc24cfd..8ca9a21 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -153,3 +153,11 @@
 RSA hash_len overflow (size_t vs unsigned int)
 depends_on:MBEDTLS_RSA_C:MBEDTLS_HAVE_INT64
 pk_rsa_overflow:
+
+ECDSA restartable sign/verify: ECDSA, restart disabled
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_SHA256_C
+pk_sign_verify_restart:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:"C9AFA9D845BA75166B5C215767B1D6934E50C3DB36E89B127B8A622B120F6721":"60FED4BA255A9D31C961EB74C6356D68C049B8923B61FA6CE669622E60F29FB6":"7903FE1008B8BC99A41AE9E95628BC64F2F1B20C2D7E9F5177A3C294D4462299":MBEDTLS_MD_SHA256:"test":"3045022100f1abb023518351cd71d881567b1ea663ed3efcf6c5132b354f28d3b0b7d383670220019f4113742a2b14bd25926b49c649155f267e60d3814b4c0cc84250e46f0083":0:0:0
+
+ECDSA restartable sign/verify: ECKEY, restart disabled
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_SHA256_C
+pk_sign_verify_restart:MBEDTLS_PK_ECKEY:MBEDTLS_ECP_DP_SECP256R1:"C9AFA9D845BA75166B5C215767B1D6934E50C3DB36E89B127B8A622B120F6721":"60FED4BA255A9D31C961EB74C6356D68C049B8923B61FA6CE669622E60F29FB6":"7903FE1008B8BC99A41AE9E95628BC64F2F1B20C2D7E9F5177A3C294D4462299":MBEDTLS_MD_SHA256:"test":"3045022100f1abb023518351cd71d881567b1ea663ed3efcf6c5132b354f28d3b0b7d383670220019f4113742a2b14bd25926b49c649155f267e60d3814b4c0cc84250e46f0083":0:0:0
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 5fa8a69..6e8c032 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -127,6 +127,14 @@
     mbedtls_rsa_context *rsa;
     mbedtls_pk_context pk;
     int msg_len;
+    void *rs_ctx = NULL;
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    mbedtls_ecdsa_restart_ctx ctx;
+
+    rs_ctx = &ctx;
+    mbedtls_ecdsa_restart_init( rs_ctx );
+    mbedtls_ecp_set_max_ops( 42 );
+#endif
 
     mbedtls_pk_init( &pk );
 
@@ -150,7 +158,13 @@
     TEST_ASSERT( mbedtls_pk_verify( &pk, digest, hash_result, 0,
                             result_str, mbedtls_pk_get_len( &pk ) ) == result );
 
+    TEST_ASSERT( mbedtls_pk_verify_restartable( &pk, digest, hash_result, 0,
+                    result_str, mbedtls_pk_get_len( &pk ), rs_ctx ) == result );
+
 exit:
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    mbedtls_ecdsa_restart_free( rs_ctx );
+#endif
     mbedtls_pk_free( &pk );
 }
 /* END_CASE */
@@ -253,12 +267,118 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_ECP_RESTARTABLE:MBEDTLS_ECDSA_C:MBEDTLS_ECDSA_DETERMINISTIC */
+void pk_sign_verify_restart( int pk_type, int grp_id, char *d_str,
+                              char *QX_str, char *QY_str,
+                              int md_alg, char *msg, char *sig_str,
+                              int max_ops, int min_restart, int max_restart )
+{
+    int ret, cnt_restart;
+    mbedtls_ecdsa_restart_ctx rs_ctx;
+    mbedtls_pk_context prv, pub;
+    unsigned char hash[MBEDTLS_MD_MAX_SIZE];
+    unsigned char sig[MBEDTLS_ECDSA_MAX_LEN];
+    unsigned char sig_check[MBEDTLS_ECDSA_MAX_LEN];
+    size_t hlen, slen, slen_check;
+    const mbedtls_md_info_t *md_info;
+
+    mbedtls_ecdsa_restart_init( &rs_ctx );
+    mbedtls_pk_init( &prv );
+    mbedtls_pk_init( &pub );
+    memset( hash, 0, sizeof( hash ) );
+    memset( sig, 0, sizeof( sig ) );
+    memset( sig_check, 0, sizeof( sig_check ) );
+
+    TEST_ASSERT( mbedtls_pk_setup( &prv, mbedtls_pk_info_from_type( pk_type ) ) == 0 );
+    TEST_ASSERT( mbedtls_ecp_group_load( &mbedtls_pk_ec( prv )->grp, grp_id ) == 0 );
+    TEST_ASSERT( mbedtls_mpi_read_string( &mbedtls_pk_ec( prv )->d, 16, d_str ) == 0 );
+
+    TEST_ASSERT( mbedtls_pk_setup( &pub, mbedtls_pk_info_from_type( pk_type ) ) == 0 );
+    TEST_ASSERT( mbedtls_ecp_group_load( &mbedtls_pk_ec( pub )->grp, grp_id ) == 0 );
+    TEST_ASSERT( mbedtls_ecp_point_read_string( &mbedtls_pk_ec( pub )->Q, 16, QX_str, QY_str ) == 0 );
+
+    slen_check = unhexify( sig_check, sig_str );
+
+    md_info = mbedtls_md_info_from_type( md_alg );
+    TEST_ASSERT( md_info != NULL );
+
+    hlen = mbedtls_md_get_size( md_info );
+    mbedtls_md( md_info, (const unsigned char *) msg, strlen( msg ), hash );
+
+    mbedtls_ecp_set_max_ops( max_ops );
+
+    slen = sizeof( sig );
+    cnt_restart = 0;
+    do {
+        ret = mbedtls_pk_sign_restartable( &prv, md_alg, hash, hlen,
+                                            sig, &slen, NULL, NULL, &rs_ctx );
+    } while( ret == MBEDTLS_ERR_ECP_IN_PROGRESS && ++cnt_restart );
+
+    TEST_ASSERT( ret == 0 );
+    TEST_ASSERT( slen == slen_check );
+    TEST_ASSERT( memcmp( sig, sig_check, slen ) == 0 );
+
+    TEST_ASSERT( cnt_restart >= min_restart );
+    TEST_ASSERT( cnt_restart <= max_restart );
+
+    cnt_restart = 0;
+    do {
+        ret = mbedtls_pk_verify_restartable( &pub, md_alg,
+                                 hash, hlen, sig, slen, &rs_ctx );
+    } while( ret == MBEDTLS_ERR_ECP_IN_PROGRESS && ++cnt_restart );
+
+    TEST_ASSERT( ret == 0 );
+    TEST_ASSERT( cnt_restart >= min_restart );
+    TEST_ASSERT( cnt_restart <= max_restart );
+
+    hash[0]++;
+    do {
+        ret = mbedtls_pk_verify_restartable( &pub, md_alg,
+                                 hash, hlen, sig, slen, &rs_ctx );
+    } while( ret == MBEDTLS_ERR_ECP_IN_PROGRESS );
+    TEST_ASSERT( ret != 0 );
+    hash[0]--;
+
+    sig[0]++;
+    do {
+        ret = mbedtls_pk_verify_restartable( &pub, md_alg,
+                                 hash, hlen, sig, slen, &rs_ctx );
+    } while( ret == MBEDTLS_ERR_ECP_IN_PROGRESS );
+    TEST_ASSERT( ret != 0 );
+    sig[0]--;
+
+    /* Do we leak memory when aborting? try verify then sign */
+    ret = mbedtls_pk_verify_restartable( &pub, md_alg,
+                             hash, hlen, sig, slen, &rs_ctx );
+    TEST_ASSERT( ret == 0 || ret == MBEDTLS_ERR_ECP_IN_PROGRESS );
+    mbedtls_ecdsa_restart_free( &rs_ctx );
+
+    slen = sizeof( sig );
+    ret = mbedtls_pk_sign_restartable( &prv, md_alg, hash, hlen,
+                                        sig, &slen, NULL, NULL, &rs_ctx );
+    TEST_ASSERT( ret == 0 || ret == MBEDTLS_ERR_ECP_IN_PROGRESS );
+
+exit:
+    mbedtls_ecdsa_restart_free( &rs_ctx );
+    mbedtls_pk_free( &prv );
+    mbedtls_pk_free( &pub );
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_SHA256_C */
 void pk_sign_verify( int type, int sign_ret, int verify_ret )
 {
     mbedtls_pk_context pk;
     unsigned char hash[50], sig[5000];
     size_t sig_len;
+    void *rs_ctx = NULL;
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    mbedtls_ecdsa_restart_ctx ctx;
+
+    rs_ctx = &ctx;
+    mbedtls_ecdsa_restart_init( rs_ctx );
+    mbedtls_ecp_set_max_ops( 42000 );
+#endif
 
     mbedtls_pk_init( &pk );
 
@@ -268,13 +388,49 @@
     TEST_ASSERT( mbedtls_pk_setup( &pk, mbedtls_pk_info_from_type( type ) ) == 0 );
     TEST_ASSERT( pk_genkey( &pk ) == 0 );
 
-    TEST_ASSERT( mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, sizeof hash,
-                          sig, &sig_len, rnd_std_rand, NULL ) == sign_ret );
+    TEST_ASSERT( mbedtls_pk_sign_restartable( &pk, MBEDTLS_MD_SHA256,
+                 hash, sizeof hash, sig, &sig_len,
+                 rnd_std_rand, NULL, rs_ctx ) == sign_ret );
 
     TEST_ASSERT( mbedtls_pk_verify( &pk, MBEDTLS_MD_SHA256,
                             hash, sizeof hash, sig, sig_len ) == verify_ret );
 
+    if( verify_ret == 0 )
+    {
+        hash[0]++;
+        TEST_ASSERT( mbedtls_pk_verify( &pk, MBEDTLS_MD_SHA256,
+                                hash, sizeof hash, sig, sig_len ) != 0 );
+        hash[0]--;
+
+        sig[0]++;
+        TEST_ASSERT( mbedtls_pk_verify( &pk, MBEDTLS_MD_SHA256,
+                                hash, sizeof hash, sig, sig_len ) != 0 );
+        sig[0]--;
+    }
+
+    TEST_ASSERT( mbedtls_pk_sign( &pk, MBEDTLS_MD_SHA256, hash, sizeof hash,
+                          sig, &sig_len, rnd_std_rand, NULL ) == sign_ret );
+
+    TEST_ASSERT( mbedtls_pk_verify_restartable( &pk, MBEDTLS_MD_SHA256,
+                 hash, sizeof hash, sig, sig_len, rs_ctx ) == verify_ret );
+
+    if( verify_ret == 0 )
+    {
+        hash[0]++;
+        TEST_ASSERT( mbedtls_pk_verify_restartable( &pk, MBEDTLS_MD_SHA256,
+                     hash, sizeof hash, sig, sig_len, rs_ctx ) != 0 );
+        hash[0]--;
+
+        sig[0]++;
+        TEST_ASSERT( mbedtls_pk_verify_restartable( &pk, MBEDTLS_MD_SHA256,
+                     hash, sizeof hash, sig, sig_len, rs_ctx ) != 0 );
+        sig[0]--;
+    }
+
 exit:
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    mbedtls_ecdsa_restart_free( rs_ctx );
+#endif
     mbedtls_pk_free( &pk );
 }
 /* END_CASE */