Dynamically allocate/free restart subcontext in PK

include/mbedtls/pk.h

@@ -135,7 +135,8 @@ typedef struct
  */
 typedef struct
 {
-    mbedtls_ecdsa_restart_ctx ecdsa;    /* temporary */
+    const mbedtls_pk_info_t *   pk_info; /**< Public key informations        */
+    void *                      rs_ctx;  /**< Underlying restart context     */
 } mbedtls_pk_restart_ctx;
 #else
 /* Now we can declare functions that take a pointer to that */

include/mbedtls/pk_internal.h

@@ -94,6 +94,14 @@ struct mbedtls_pk_info_t
     /** Free the given context */
     void (*ctx_free_func)( void *ctx );
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    /** Allocate the restart context */
+    void * (*rs_alloc_func)( void );
+
+    /** Free the restart context */
+    void (*rs_free_func)( void *rs_ctx );
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+
     /** Interface with the debug module */
     void (*debug_func)( const void *ctx, mbedtls_pk_debug_item *items );
 

library/pk.c

@@ -79,7 +79,8 @@ void mbedtls_pk_free( mbedtls_pk_context *ctx )
  */
 void mbedtls_pk_restart_init( mbedtls_pk_restart_ctx *ctx )
 {
-    mbedtls_ecdsa_restart_init( &ctx->ecdsa );
+    ctx->pk_info = NULL;
+    ctx->rs_ctx = NULL;
 }
 
 /*
@@ -87,10 +88,16 @@ void mbedtls_pk_restart_init( mbedtls_pk_restart_ctx *ctx )
  */
 void mbedtls_pk_restart_free( mbedtls_pk_restart_ctx *ctx )
 {
-    if( ctx == NULL )
+    if( ctx == NULL || ctx->pk_info == NULL ||
+        ctx->pk_info->rs_free_func == NULL )
+    {
         return;
+    }
 
-    mbedtls_ecdsa_restart_free( &ctx->ecdsa );
+    ctx->pk_info->rs_free_func( ctx->rs_ctx );
+
+    ctx->pk_info = NULL;
+    ctx->rs_ctx = NULL;
 }
 #endif /* MBEDTLS_ECP_RESTARTABLE */
 
@@ -196,6 +203,30 @@ static inline int pk_hashlen_helper( mbedtls_md_type_t md_alg, size_t *hash_len
     return( 0 );
 }
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+/*
+ * Helper to set up a restart context if needed
+ */
+static int pk_restart_setup( mbedtls_pk_restart_ctx *ctx,
+                               const mbedtls_pk_info_t *info )
+{
+    /* Don't do anything it already set up */
+    if( ctx->pk_info != NULL )
+        return( 0 );
+
+    /* Should never happen when we're called */
+    if( info->rs_alloc_func == NULL || info->rs_free_func == NULL )
+        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+
+    if( ( ctx->rs_ctx = info->rs_alloc_func() ) == NULL )
+        return( MBEDTLS_ERR_PK_ALLOC_FAILED );
+
+    ctx->pk_info = info;
+
+    return( 0 );
+}
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+
 /*
  * Verify a signature (restartable)
  */
@@ -210,10 +241,20 @@ int mbedtls_pk_verify_restartable( mbedtls_pk_context *ctx,
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
 #if defined(MBEDTLS_ECP_RESTARTABLE)
-    if( ctx->pk_info->verify_rs_func != NULL )
+    if( rs_ctx != NULL && 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 ) );
+        int ret;
+
+        if( ( ret = pk_restart_setup( rs_ctx, ctx->pk_info ) ) != 0 )
+            return( ret );
+
+        ret = ctx->pk_info->verify_rs_func( ctx->pk_ctx,
+                   md_alg, hash, hash_len, sig, sig_len, rs_ctx->rs_ctx );
+
+        if( ret != MBEDTLS_ERR_ECP_IN_PROGRESS )
+            mbedtls_pk_restart_free( rs_ctx );
+
+        return( ret );
     }
 #else
     (void) rs_ctx;
@@ -310,10 +351,20 @@ int mbedtls_pk_sign_restartable( mbedtls_pk_context *ctx,
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
 #if defined(MBEDTLS_ECP_RESTARTABLE)
-    if( ctx->pk_info->sign_rs_func != NULL )
+    if( rs_ctx != NULL && 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 ) );
+        int ret;
+
+        if( ( ret = pk_restart_setup( rs_ctx, ctx->pk_info ) ) != 0 )
+            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 );
+
+        if( ret != MBEDTLS_ERR_ECP_IN_PROGRESS )
+            mbedtls_pk_restart_free( rs_ctx );
+
+        return( ret );
     }
 #else
     (void) rs_ctx;

library/pk_wrap.c

@@ -189,6 +189,10 @@ const mbedtls_pk_info_t mbedtls_rsa_info = {
     rsa_check_pair_wrap,
     rsa_alloc_wrap,
     rsa_free_wrap,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
     rsa_debug,
 };
 #endif /* MBEDTLS_RSA_C */
@@ -401,6 +405,24 @@ static void eckey_debug( const void *ctx, mbedtls_pk_debug_item *items )
     items->value = &( ((mbedtls_ecp_keypair *) ctx)->Q );
 }
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+static void *eckey_rs_alloc( void )
+{
+    void *ctx = mbedtls_calloc( 1, sizeof( mbedtls_ecdsa_restart_ctx ) );
+
+    if( ctx != NULL )
+        mbedtls_ecdsa_restart_init( ctx );
+
+    return( ctx );
+}
+
+static void eckey_rs_free( void *ctx )
+{
+    mbedtls_ecdsa_restart_free( ctx );
+    mbedtls_free( ctx );
+}
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+
 const mbedtls_pk_info_t mbedtls_eckey_info = {
     MBEDTLS_PK_ECKEY,
     "EC",
@@ -426,6 +448,10 @@ const mbedtls_pk_info_t mbedtls_eckey_info = {
     eckey_check_pair,
     eckey_alloc_wrap,
     eckey_free_wrap,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    eckey_rs_alloc,
+    eckey_rs_free,
+#endif
     eckey_debug,
 };
 
@@ -454,6 +480,10 @@ const mbedtls_pk_info_t mbedtls_eckeydh_info = {
     eckey_check_pair,
     eckey_alloc_wrap,       /* Same underlying key structure */
     eckey_free_wrap,        /* Same underlying key structure */
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
     eckey_debug,            /* Same underlying key structure */
 };
 #endif /* MBEDTLS_ECP_C */
@@ -555,6 +585,10 @@ const mbedtls_pk_info_t mbedtls_ecdsa_info = {
     eckey_check_pair,   /* Compatible key structures */
     ecdsa_alloc_wrap,
     ecdsa_free_wrap,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    eckey_rs_alloc,
+    eckey_rs_free,
+#endif
     eckey_debug,        /* Compatible key structures */
 };
 #endif /* MBEDTLS_ECDSA_C */
@@ -677,6 +711,10 @@ const mbedtls_pk_info_t mbedtls_rsa_alt_info = {
 #endif
     rsa_alt_alloc_wrap,
     rsa_alt_free_wrap,
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
     NULL,
 };