Make PK EC sign/verify actually restartable
diff --git a/include/mbedtls/ecdsa.h b/include/mbedtls/ecdsa.h
index ce94af8..7784252 100644
--- a/include/mbedtls/ecdsa.h
+++ b/include/mbedtls/ecdsa.h
@@ -91,6 +91,9 @@
 #if defined(MBEDTLS_ECDSA_DETERMINISTIC)
     mbedtls_ecdsa_restart_det_ctx *det; /*!<  ecdsa_sign_det() sub-context  */
 #endif
+#if defined(MBEDTLS_PK_C)
+    mbedtls_ecdsa_context *ecdsa;       /*!<  used by the PK layer          */
+#endif
 } mbedtls_ecdsa_restart_ctx;
 
 #else /* MBEDTLS_ECP_RESTARTABLE */
diff --git a/include/mbedtls/pk_internal.h b/include/mbedtls/pk_internal.h
index 01d0f21..e996b6c 100644
--- a/include/mbedtls/pk_internal.h
+++ b/include/mbedtls/pk_internal.h
@@ -58,6 +58,21 @@
                       int (*f_rng)(void *, unsigned char *, size_t),
                       void *p_rng );
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    /** Verify signature (restartable) */
+    int (*verify_rs_func)( void *ctx, mbedtls_md_type_t md_alg,
+                           const unsigned char *hash, size_t hash_len,
+                           const unsigned char *sig, size_t sig_len,
+                           void *rs_ctx );
+
+    /** 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,
+                         int (*f_rng)(void *, unsigned char *, size_t),
+                         void *p_rng, void *rs_ctx );
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+
     /** Decrypt message */
     int (*decrypt_func)( void *ctx, const unsigned char *input, size_t ilen,
                          unsigned char *output, size_t *olen, size_t osize,
diff --git a/library/ecdsa.c b/library/ecdsa.c
index e137a00..6730762 100644
--- a/library/ecdsa.c
+++ b/library/ecdsa.c
@@ -779,6 +779,9 @@
 #if defined(MBEDTLS_ECDSA_DETERMINISTIC)
     ctx->det = NULL;
 #endif
+#if defined(MBEDTLS_PK_C)
+    ctx->ecdsa = NULL;
+#endif
 }
 
 /*
@@ -801,6 +804,12 @@
     mbedtls_free( ctx->det );
     ctx->det = NULL;
 #endif
+
+#if defined(MBEDTLS_PK_C)
+    mbedtls_ecdsa_free( ctx->ecdsa );
+    mbedtls_free( ctx->ecdsa );
+    ctx->ecdsa = NULL;
+#endif
 }
 #endif /* MBEDTLS_ECP_RESTARTABLE */
 
diff --git a/library/pk.c b/library/pk.c
index b5081f9..e439c7a 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -184,12 +184,20 @@
                const unsigned char *sig, size_t sig_len,
                void *rs_ctx )
 {
-    (void) rs_ctx; // XXX temporary
-
     if( ctx == NULL || ctx->pk_info == NULL ||
         pk_hashlen_helper( md_alg, &hash_len ) != 0 )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    if( ctx->pk_info->verify_rs_func != NULL )
+    {
+        return( ctx->pk_info->verify_rs_func( ctx->pk_ctx,
+                    md_alg, hash, hash_len, sig, sig_len, rs_ctx ) );
+    }
+#else
+    (void) rs_ctx;
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+
     if( ctx->pk_info->verify_func == NULL )
         return( MBEDTLS_ERR_PK_TYPE_MISMATCH );
 
@@ -276,12 +284,20 @@
              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
              void *rs_ctx )
 {
-    (void) rs_ctx; // XXX temporary
-
     if( ctx == NULL || ctx->pk_info == NULL ||
         pk_hashlen_helper( md_alg, &hash_len ) != 0 )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    if( ctx->pk_info->sign_rs_func != NULL )
+    {
+        return( ctx->pk_info->sign_rs_func( ctx->pk_ctx, md_alg,
+                    hash, hash_len, sig, sig_len, f_rng, p_rng, rs_ctx ) );
+    }
+#else
+    (void) rs_ctx;
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+
     if( ctx->pk_info->sign_func == NULL )
         return( MBEDTLS_ERR_PK_TYPE_MISMATCH );
 
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index db6274c..d3933a0 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -180,6 +180,10 @@
     rsa_can_do,
     rsa_verify_wrap,
     rsa_sign_wrap,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
     rsa_decrypt_wrap,
     rsa_encrypt_wrap,
     rsa_check_pair_wrap,
@@ -252,6 +256,118 @@
     return( ret );
 }
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+/* Forward declarations */
+static int ecdsa_verify_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
+                       const unsigned char *hash, size_t hash_len,
+                       const unsigned char *sig, size_t sig_len,
+                       void *rs_ctx );
+
+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,
+                   int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
+                   void *rs_ctx );
+
+static int eckey_verify_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
+                       const unsigned char *hash, size_t hash_len,
+                       const unsigned char *sig, size_t sig_len,
+                       void *p_rs_ctx )
+{
+    int ret;
+    mbedtls_ecdsa_context ecdsa, *p_ecdsa = &ecdsa;
+    mbedtls_ecdsa_restart_ctx *rs_ctx = p_rs_ctx;
+
+    mbedtls_ecdsa_init( &ecdsa );
+
+    /* set up our own sub-context if needed */
+    if( mbedtls_ecp_restart_enabled() &&
+        rs_ctx != NULL && rs_ctx->ecdsa == NULL )
+    {
+        rs_ctx->ecdsa = mbedtls_calloc( 1, sizeof( *rs_ctx->ecdsa ) );
+        if( rs_ctx->ecdsa == NULL )
+            return( MBEDTLS_ERR_PK_ALLOC_FAILED );
+
+        mbedtls_ecdsa_init( rs_ctx->ecdsa );
+        MBEDTLS_MPI_CHK( mbedtls_ecdsa_from_keypair( rs_ctx->ecdsa, ctx ) );
+    }
+
+    if( rs_ctx != NULL && rs_ctx->ecdsa != NULL )
+    {
+        /* redirect to our context */
+        p_ecdsa = rs_ctx->ecdsa;
+    }
+    else
+    {
+        MBEDTLS_MPI_CHK( mbedtls_ecdsa_from_keypair( p_ecdsa, ctx ) );
+    }
+
+    MBEDTLS_MPI_CHK( ecdsa_verify_rs_wrap( p_ecdsa, md_alg, hash, hash_len,
+                                           sig, sig_len, rs_ctx ) );
+
+cleanup:
+    /* clear our sub-context when not in progress (done or error) */
+    if( rs_ctx != NULL && ret != MBEDTLS_ERR_ECP_IN_PROGRESS ) {
+        mbedtls_ecdsa_free( rs_ctx->ecdsa );
+        mbedtls_free( rs_ctx->ecdsa );
+        rs_ctx->ecdsa = NULL;
+    }
+
+    mbedtls_ecdsa_free( &ecdsa );
+
+    return( ret );
+}
+
+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,
+                   int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
+                   void *p_rs_ctx )
+{
+    int ret;
+    mbedtls_ecdsa_context ecdsa, *p_ecdsa = &ecdsa;
+    mbedtls_ecdsa_restart_ctx *rs_ctx = p_rs_ctx;
+
+    mbedtls_ecdsa_init( &ecdsa );
+
+    /* set up our own sub-context if needed */
+    if( mbedtls_ecp_restart_enabled() &&
+        rs_ctx != NULL && rs_ctx->ecdsa == NULL )
+    {
+        rs_ctx->ecdsa = mbedtls_calloc( 1, sizeof( *rs_ctx->ecdsa ) );
+        if( rs_ctx->ecdsa == NULL )
+            return( MBEDTLS_ERR_PK_ALLOC_FAILED );
+
+        mbedtls_ecdsa_init( rs_ctx->ecdsa );
+        MBEDTLS_MPI_CHK( mbedtls_ecdsa_from_keypair( rs_ctx->ecdsa, ctx ) );
+    }
+
+    if( rs_ctx != NULL && rs_ctx->ecdsa != NULL )
+    {
+        /* redirect to our context */
+        p_ecdsa = rs_ctx->ecdsa;
+    }
+    else
+    {
+        MBEDTLS_MPI_CHK( mbedtls_ecdsa_from_keypair( p_ecdsa, ctx ) );
+    }
+
+    MBEDTLS_MPI_CHK( ecdsa_sign_rs_wrap( p_ecdsa, md_alg, hash, hash_len,
+                                         sig, sig_len, f_rng, p_rng, rs_ctx ) );
+
+cleanup:
+    /* clear our sub-context when not in progress (done or error) */
+    if( rs_ctx != NULL && ret != MBEDTLS_ERR_ECP_IN_PROGRESS ) {
+        mbedtls_ecdsa_free( rs_ctx->ecdsa );
+        mbedtls_free( rs_ctx->ecdsa );
+        rs_ctx->ecdsa = NULL;
+    }
+
+    mbedtls_ecdsa_free( &ecdsa );
+
+    return( ret );
+}
+#endif /* MBEDTLS_ECP_RESTARTABLE */
 #endif /* MBEDTLS_ECDSA_C */
 
 static int eckey_check_pair( const void *pub, const void *prv )
@@ -291,10 +407,18 @@
 #if defined(MBEDTLS_ECDSA_C)
     eckey_verify_wrap,
     eckey_sign_wrap,
-#else
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    eckey_verify_rs_wrap,
+    eckey_sign_rs_wrap,
+#endif
+#else /* MBEDTLS_ECDSA_C */
+    NULL,
+    NULL,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
     NULL,
     NULL,
 #endif
+#endif /* MBEDTLS_ECDSA_C */
     NULL,
     NULL,
     eckey_check_pair,
@@ -319,6 +443,10 @@
     eckeydh_can_do,
     NULL,
     NULL,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
     NULL,
     NULL,
     eckey_check_pair,
@@ -359,6 +487,40 @@
                 md_alg, hash, hash_len, sig, sig_len, f_rng, p_rng ) );
 }
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+static int ecdsa_verify_rs_wrap( void *ctx, mbedtls_md_type_t md_alg,
+                       const unsigned char *hash, size_t hash_len,
+                       const unsigned char *sig, size_t sig_len,
+                       void *rs_ctx )
+{
+    int ret;
+    ((void) md_alg);
+
+    ret = mbedtls_ecdsa_read_signature_restartable(
+            (mbedtls_ecdsa_context *) ctx,
+            hash, hash_len, sig, sig_len,
+            (mbedtls_ecdsa_restart_ctx *) rs_ctx );
+
+    if( ret == MBEDTLS_ERR_ECP_SIG_LEN_MISMATCH )
+        return( MBEDTLS_ERR_PK_SIG_LEN_MISMATCH );
+
+    return( ret );
+}
+
+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,
+                   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,
+                (mbedtls_ecdsa_restart_ctx *) rs_ctx ) );
+
+}
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+
 static void *ecdsa_alloc_wrap( void )
 {
     void *ctx = mbedtls_calloc( 1, sizeof( mbedtls_ecdsa_context ) );
@@ -382,6 +544,10 @@
     ecdsa_can_do,
     ecdsa_verify_wrap,
     ecdsa_sign_wrap,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    ecdsa_verify_rs_wrap,
+    ecdsa_sign_rs_wrap,
+#endif
     NULL,
     NULL,
     eckey_check_pair,   /* Compatible key structures */
@@ -496,6 +662,10 @@
     rsa_alt_can_do,
     NULL,
     rsa_alt_sign_wrap,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
     rsa_alt_decrypt_wrap,
     NULL,
 #if defined(MBEDTLS_RSA_C)
diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index 8ca9a21..bd3db18 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -161,3 +161,27 @@
 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
+
+ECDSA restartable sign/verify: ECDSA, max_ops=1
+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":1:1:10000
+
+ECDSA restartable sign/verify: ECKEY, max_ops=1
+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":1:1:10000
+
+ECDSA restartable sign/verify: ECDSA, max_ops=10000
+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":10000:0:0
+
+ECDSA restartable sign/verify: ECKEY, max_ops=10000
+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":10000:0:0
+
+ECDSA restartable sign/verify: ECDSA, max_ops=250
+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":250:2:64
+
+ECDSA restartable sign/verify: ECKEY, max_ops=1
+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":250:2:64