From 77af79a324f8d340ef5595d92e1782eb3125ccc7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20P=C3=A9gouri=C3=A9-Gonnard?=
 <manuel.pegourie-gonnard@arm.com>
Date: Tue, 14 Mar 2017 10:58:00 +0100
Subject: [PATCH] Add proper allocation of restart context

We'll need to store MPIs and other things that allocate memory in this
context, so we need a place to free it. We can't rely on doing it before
returning from ecp_mul() as we might return MBEDTLS_ERR_ECP_IN_PROGRESS (thus
preserving the context) and never be called again (for example, TLS handshake
aborted for another reason). So, ecp_group_free() looks like a good place to
do this, if the restart context is part of struct ecp_group.

This means it's not possible to use the same ecp_group structure in different
threads concurrently, but:
- that's already the case (and documented) for other reasons
- this feature is precisely intended for environments that lack threading

An alternative option would be for the caller to have to allocate/free the
restart context and pass it explicitly, but this means creating new functions
that take a context argument, and putting a burden on the user.
---
 include/mbedtls/ecp.h                | 13 +++++++
 library/ecp.c                        | 55 +++++++++++++++++++++++-----
 tests/suites/test_suite_ecp.function |  4 ++
 3 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/include/mbedtls/ecp.h b/include/mbedtls/ecp.h
index 7bcc69c017..d44e658437 100644
--- a/include/mbedtls/ecp.h
+++ b/include/mbedtls/ecp.h
@@ -112,6 +112,16 @@ typedef struct
 }
 mbedtls_ecp_point;
 
+#if defined(MBEDTLS_ECP_EARLY_RETURN)
+/**
+ * \brief           ECP context for resuming operations after returning
+ *                  \c MBEDTLS_ERR_ECP_IN_PROGRESS
+ *
+ * \note            Opaque struct
+ */
+typedef struct mbedtls_ecp_restart mbedtls_ecp_restart_ctx;
+#endif
+
 /**
  * \brief           ECP group structure
  *
@@ -153,6 +163,9 @@ typedef struct
     void *t_data;                       /*!< unused                         */
     mbedtls_ecp_point *T;       /*!<  pre-computed points for ecp_mul_comb()        */
     size_t T_size;      /*!<  number for pre-computed points                */
+#if defined(MBEDTLS_ECP_EARLY_RETURN)
+    mbedtls_ecp_restart_ctx *rs;        /*!<  context for resuming operation */
+#endif
 }
 mbedtls_ecp_group;
 
diff --git a/library/ecp.c b/library/ecp.c
index 71f8700b57..60aa0a3cd9 100644
--- a/library/ecp.c
+++ b/library/ecp.c
@@ -101,20 +101,28 @@ void mbedtls_ecp_set_max_ops( unsigned max_ops )
 }
 
 /*
- * Saved context type for restarting operations.
- *
- * XXX: this is a temporary place for the definition
+ * Restart context type for interrupted operations
  */
-typedef struct {
+struct mbedtls_ecp_restart {
     unsigned char fake_it;  /* for tests: should we fake early return? */
-} ecp_restart_context;
+};
 
 /*
- * Saved context fro restarting operations.
- *
- * XXX: temporary place for the allocation
+ * Init restart context
  */
-static ecp_restart_context ecp_restart;
+static void ecp_restart_init( mbedtls_ecp_restart_ctx *ctx )
+{
+    memset( ctx, 0, sizeof( mbedtls_ecp_restart_ctx ) );
+}
+
+/*
+ * Free the components of a restart context
+ */
+static void ecp_restart_free( mbedtls_ecp_restart_ctx *ctx )
+{
+    if( ctx == NULL )
+        return;
+}
 #endif /* MBEDTLS_ECP_EARLY_RETURN */
 
 #if defined(MBEDTLS_ECP_DP_SECP192R1_ENABLED) ||   \
@@ -378,6 +386,11 @@ void mbedtls_ecp_group_free( mbedtls_ecp_group *grp )
         mbedtls_free( grp->T );
     }
 
+#if defined(MBEDTLS_ECP_EARLY_RETURN)
+    ecp_restart_free( grp->rs );
+    mbedtls_free( grp->rs );
+#endif
+
     mbedtls_zeroize( grp, sizeof( mbedtls_ecp_group ) );
 }
 
@@ -1501,8 +1514,22 @@ static int ecp_mul_comb( mbedtls_ecp_group *grp, mbedtls_ecp_point *R,
     size_t d;
     mbedtls_ecp_point *T = NULL;
 
+    /* set up restart context if needed */
 #if defined(MBEDTLS_ECP_EARLY_RETURN)
-    if( ecp_restart.fake_it++ != 0 && ecp_max_ops != 0 )
+    if( ecp_max_ops != 0 && grp->rs == NULL )
+    {
+        grp->rs = mbedtls_calloc( 1, sizeof( mbedtls_ecp_restart_ctx ) );
+        if( grp->rs == NULL )
+            return( MBEDTLS_ERR_ECP_ALLOC_FAILED );
+        ecp_restart_init( grp->rs );
+
+        grp->rs->fake_it = 1;
+    }
+#endif
+
+    /* XXX: temporary */
+#if defined(MBEDTLS_ECP_EARLY_RETURN)
+    if( grp->rs && grp->rs->fake_it++ != 0 )
         return( MBEDTLS_ERR_ECP_IN_PROGRESS );
 #endif
 
@@ -1560,6 +1587,14 @@ cleanup:
     if( ret != 0 )
         mbedtls_ecp_point_free( R );
 
+#if defined(MBEDTLS_ECP_EARLY_RETURN)
+    if( grp->rs != NULL && ret != MBEDTLS_ERR_ECP_IN_PROGRESS ) {
+        ecp_restart_free( grp->rs );
+        mbedtls_free( grp->rs );
+        grp->rs = NULL;
+    }
+#endif
+
     return( ret );
 }
 
diff --git a/tests/suites/test_suite_ecp.function b/tests/suites/test_suite_ecp.function
index c8daef990d..74e23875da 100644
--- a/tests/suites/test_suite_ecp.function
+++ b/tests/suites/test_suite_ecp.function
@@ -145,6 +145,10 @@ void ecp_test_vect_restart( int id,
     }
     while( ret != 0 );
 
+    /* Do we leak memory when not finishing an operation? */
+    ret = mbedtls_ecp_mul( &grp, &R, &dB, &R, NULL, NULL );
+    TEST_ASSERT( ret == 0 || ret == MBEDTLS_ERR_ECP_IN_PROGRESS );
+
 exit:
     mbedtls_ecp_group_free( &grp ); mbedtls_ecp_point_free( &R );
     mbedtls_mpi_free( &dA ); mbedtls_mpi_free( &xA ); mbedtls_mpi_free( &yA );