libmbedtls: add mbedtls_mpi_init_mempool()

782fddd10ccb ("libmbedtls: add mbedtls_mpi_init_mempool()") from
branch import/mbedtls-2.6.1

Adds mbedtls_mpi_init_mempool() which initializes a mbedtls_mpi struct
to use the mempool mbedtls_mpi_mempool if configured for memory
allocation. All local memory allocation is changed to use
mbedtls_mpi_init_mempool() instead of mbedtls_mpi_init(). This will give
a stack like alloc/free pattern for which the mempool is optimized.

Acked-by: Jerome Forissier <jerome.forissier@linaro.org>
Signed-off-by: Jens Wiklander <jens.wiklander@linaro.org>

2 files changed, 83 insertions, 26 deletions

diff --git a/lib/libmbedtls/mbedtls/include/mbedtls/bignum.h b/lib/libmbedtls/mbedtls/include/mbedtls/bignum.h
index 29291a4f..355fcf86 100644
--- a/lib/libmbedtls/mbedtls/include/mbedtls/bignum.h
+++ b/lib/libmbedtls/mbedtls/include/mbedtls/bignum.h
@@ -178,12 +178,15 @@ extern "C" {
  */
 typedef struct
 {
-    int s;              /*!<  integer sign      */
+    short s;              /*!<  integer sign      */
+    short use_mempool;
     size_t n;           /*!<  total # of limbs  */
     mbedtls_mpi_uint *p;          /*!<  pointer to limbs  */
 }
 mbedtls_mpi;
 
+extern void *mbedtls_mpi_mempool;
+
 /**
  * \brief           Initialize one MPI (make internal references valid)
  *                  This just makes it ready to be set or freed,
@@ -192,6 +195,7 @@ mbedtls_mpi;
  * \param X         One MPI to initialize.
  */
 void mbedtls_mpi_init( mbedtls_mpi *X );
+void mbedtls_mpi_init_mempool( mbedtls_mpi *X );
 
 /**
  * \brief          Unallocate one MPI
diff --git a/lib/libmbedtls/mbedtls/library/bignum.c b/lib/libmbedtls/mbedtls/library/bignum.c
index 4603c8cd..79f5da43 100644
--- a/lib/libmbedtls/mbedtls/library/bignum.c
+++ b/lib/libmbedtls/mbedtls/library/bignum.c
@@ -58,6 +58,8 @@
 #define mbedtls_free       free
 #endif
 
+#include <mempool.h>
+
 /* Implementation that should never be optimized out by the compiler */
 static void mbedtls_mpi_zeroize( mbedtls_mpi_uint *v, size_t n ) {
     volatile mbedtls_mpi_uint *p = v; while( n-- ) *p++ = 0;
@@ -76,19 +78,32 @@ static void mbedtls_mpi_zeroize( mbedtls_mpi_uint *v, size_t n ) {
 #define BITS_TO_LIMBS(i)  ( (i) / biL + ( (i) % biL != 0 ) )
 #define CHARS_TO_LIMBS(i) ( (i) / ciL + ( (i) % ciL != 0 ) )
 
+void *mbedtls_mpi_mempool;
+
 /*
  * Initialize one MPI
  */
-void mbedtls_mpi_init( mbedtls_mpi *X )
+static void mpi_init( mbedtls_mpi *X, short use_mempool)
 {
     if( X == NULL )
         return;
 
     X->s = 1;
+    X->use_mempool = use_mempool;
     X->n = 0;
     X->p = NULL;
 }
 
+void mbedtls_mpi_init( mbedtls_mpi *X )
+{
+    mpi_init( X, 0 /*use_mempool*/ );
+}
+
+void mbedtls_mpi_init_mempool( mbedtls_mpi *X )
+{
+    mpi_init( X, !!mbedtls_mpi_mempool /*use_mempool*/ );
+}
+
 /*
  * Unallocate one MPI
  */
@@ -100,7 +115,10 @@ void mbedtls_mpi_free( mbedtls_mpi *X )
     if( X->p != NULL )
     {
         mbedtls_mpi_zeroize( X->p, X->n );
-        mbedtls_free( X->p );
+        if( X->use_mempool )
+            mempool_free( mbedtls_mpi_mempool, X->p );
+        else
+            mbedtls_free( X->p );
     }
 
     X->s = 1;
@@ -120,14 +138,28 @@ int mbedtls_mpi_grow( mbedtls_mpi *X, size_t nblimbs )
 
     if( X->n < nblimbs )
     {
-        if( ( p = (mbedtls_mpi_uint*)mbedtls_calloc( nblimbs, ciL ) ) == NULL )
-            return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
+        if( X->use_mempool )
+	{
+            p = mempool_alloc( mbedtls_mpi_mempool, nblimbs * ciL );
+            if( p == NULL )
+                return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
+            memset( p, 0, nblimbs * ciL );
+        }
+        else
+        {
+            p = (mbedtls_mpi_uint*)mbedtls_calloc( nblimbs, ciL );
+            if( p == NULL )
+                return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
+        }
 
         if( X->p != NULL )
         {
             memcpy( p, X->p, X->n * ciL );
             mbedtls_mpi_zeroize( X->p, X->n );
-            mbedtls_free( X->p );
+            if( X->use_mempool )
+                mempool_free( mbedtls_mpi_mempool, X->p);
+            else
+                mbedtls_free( X->p );
         }
 
         X->n = nblimbs;
@@ -158,14 +190,28 @@ int mbedtls_mpi_shrink( mbedtls_mpi *X, size_t nblimbs )
     if( i < nblimbs )
         i = nblimbs;
 
-    if( ( p = (mbedtls_mpi_uint*)mbedtls_calloc( i, ciL ) ) == NULL )
-        return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
+    if( X->use_mempool )
+    {
+        p = mempool_alloc( mbedtls_mpi_mempool, nblimbs * ciL );
+        if( p == NULL )
+            return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
+        memset(p, 0, nblimbs * ciL);
+    }
+    else
+    {
+        p = (mbedtls_mpi_uint*)mbedtls_calloc( nblimbs, ciL );
+        if( p == NULL )
+            return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
+    }
 
     if( X->p != NULL )
     {
         memcpy( p, X->p, i * ciL );
         mbedtls_mpi_zeroize( X->p, X->n );
-        mbedtls_free( X->p );
+        if( X->use_mempool )
+            mempool_free( mbedtls_mpi_mempool, X->p );
+        else
+            mbedtls_free( X->p );
     }
 
     X->n = i;
@@ -431,7 +477,7 @@ int mbedtls_mpi_read_string( mbedtls_mpi *X, int radix, const char *s )
     if( radix < 2 || radix > 16 )
         return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
 
-    mbedtls_mpi_init( &T );
+    mbedtls_mpi_init_mempool( &T );
 
     slen = strlen( s );
 
@@ -548,7 +594,7 @@ int mbedtls_mpi_write_string( const mbedtls_mpi *X, int radix,
     }
 
     p = buf;
-    mbedtls_mpi_init( &T );
+    mbedtls_mpi_init_mempool( &T );
 
     if( X->s == -1 )
         *p++ = '-';
@@ -970,7 +1016,7 @@ int mbedtls_mpi_sub_abs( mbedtls_mpi *X, const mbedtls_mpi *A, const mbedtls_mpi
     if( mbedtls_mpi_cmp_abs( A, B ) < 0 )
         return( MBEDTLS_ERR_MPI_NEGATIVE_VALUE );
 
-    mbedtls_mpi_init( &TB );
+    mbedtls_mpi_init_mempool( &TB );
 
     if( X == B )
     {
@@ -1176,7 +1222,7 @@ int mbedtls_mpi_mul_mpi( mbedtls_mpi *X, const mbedtls_mpi *A, const mbedtls_mpi
     size_t i, j;
     mbedtls_mpi TA, TB;
 
-    mbedtls_mpi_init( &TA ); mbedtls_mpi_init( &TB );
+    mbedtls_mpi_init_mempool( &TA ); mbedtls_mpi_init_mempool( &TB );
 
     if( X == A ) { MBEDTLS_MPI_CHK( mbedtls_mpi_copy( &TA, A ) ); A = &TA; }
     if( X == B ) { MBEDTLS_MPI_CHK( mbedtls_mpi_copy( &TB, B ) ); B = &TB; }
@@ -1328,8 +1374,9 @@ int mbedtls_mpi_div_mpi( mbedtls_mpi *Q, mbedtls_mpi *R, const mbedtls_mpi *A, c
     if( mbedtls_mpi_cmp_int( B, 0 ) == 0 )
         return( MBEDTLS_ERR_MPI_DIVISION_BY_ZERO );
 
-    mbedtls_mpi_init( &X ); mbedtls_mpi_init( &Y ); mbedtls_mpi_init( &Z );
-    mbedtls_mpi_init( &T1 ); mbedtls_mpi_init( &T2 );
+    mbedtls_mpi_init_mempool( &X ); mbedtls_mpi_init_mempool( &Y );
+    mbedtls_mpi_init_mempool( &Z ); mbedtls_mpi_init_mempool( &T1 );
+    mbedtls_mpi_init_mempool( &T2 );
 
     if( mbedtls_mpi_cmp_abs( A, B ) < 0 )
     {
@@ -1626,8 +1673,8 @@ int mbedtls_mpi_exp_mod( mbedtls_mpi *X, const mbedtls_mpi *A, const mbedtls_mpi
      * Init temps and window size
      */
     mbedtls_mpi_montg_init( &mm, N );
-    mbedtls_mpi_init( &RR ); mbedtls_mpi_init( &T );
-    mbedtls_mpi_init( &Apos );
+    mbedtls_mpi_init_mempool( &RR ); mbedtls_mpi_init_mempool( &T );
+    mbedtls_mpi_init_mempool( &Apos );
     memset( W, 0, sizeof( W ) );
 
     i = mbedtls_mpi_bitlen( E );
@@ -1820,7 +1867,8 @@ int mbedtls_mpi_gcd( mbedtls_mpi *G, const mbedtls_mpi *A, const mbedtls_mpi *B
     size_t lz, lzt;
     mbedtls_mpi TG, TA, TB;
 
-    mbedtls_mpi_init( &TG ); mbedtls_mpi_init( &TA ); mbedtls_mpi_init( &TB );
+    mbedtls_mpi_init_mempool( &TG ); mbedtls_mpi_init_mempool( &TA );
+    mbedtls_mpi_init_mempool( &TB );
 
     MBEDTLS_MPI_CHK( mbedtls_mpi_copy( &TA, A ) );
     MBEDTLS_MPI_CHK( mbedtls_mpi_copy( &TB, B ) );
@@ -1898,9 +1946,11 @@ int mbedtls_mpi_inv_mod( mbedtls_mpi *X, const mbedtls_mpi *A, const mbedtls_mpi
     if( mbedtls_mpi_cmp_int( N, 1 ) <= 0 )
         return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
 
-    mbedtls_mpi_init( &TA ); mbedtls_mpi_init( &TU ); mbedtls_mpi_init( &U1 ); mbedtls_mpi_init( &U2 );
-    mbedtls_mpi_init( &G ); mbedtls_mpi_init( &TB ); mbedtls_mpi_init( &TV );
-    mbedtls_mpi_init( &V1 ); mbedtls_mpi_init( &V2 );
+    mbedtls_mpi_init_mempool( &TA ); mbedtls_mpi_init_mempool( &TU );
+    mbedtls_mpi_init_mempool( &U1 ); mbedtls_mpi_init_mempool( &U2 );
+    mbedtls_mpi_init_mempool( &G ); mbedtls_mpi_init_mempool( &TB );
+    mbedtls_mpi_init_mempool( &TV ); mbedtls_mpi_init_mempool( &V1 );
+    mbedtls_mpi_init_mempool( &V2 );
 
     MBEDTLS_MPI_CHK( mbedtls_mpi_gcd( &G, A, N ) );
 
@@ -2053,8 +2103,9 @@ static int mpi_miller_rabin( const mbedtls_mpi *X,
     size_t i, j, k, n, s;
     mbedtls_mpi W, R, T, A, RR;
 
-    mbedtls_mpi_init( &W ); mbedtls_mpi_init( &R ); mbedtls_mpi_init( &T ); mbedtls_mpi_init( &A );
-    mbedtls_mpi_init( &RR );
+    mbedtls_mpi_init_mempool( &W ); mbedtls_mpi_init_mempool( &R );
+    mbedtls_mpi_init_mempool( &T ); mbedtls_mpi_init_mempool( &A );
+    mbedtls_mpi_init_mempool( &RR );
 
     /*
      * W = |X| - 1
@@ -2193,7 +2244,7 @@ int mbedtls_mpi_gen_prime( mbedtls_mpi *X, size_t nbits, int dh_flag,
     if( nbits < 3 || nbits > MBEDTLS_MPI_MAX_BITS )
         return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
 
-    mbedtls_mpi_init( &Y );
+    mbedtls_mpi_init_mempool( &Y );
 
     n = BITS_TO_LIMBS( nbits );
 
@@ -2291,8 +2342,10 @@ int mbedtls_mpi_self_test( int verbose )
     int ret, i;
     mbedtls_mpi A, E, N, X, Y, U, V;
 
-    mbedtls_mpi_init( &A ); mbedtls_mpi_init( &E ); mbedtls_mpi_init( &N ); mbedtls_mpi_init( &X );
-    mbedtls_mpi_init( &Y ); mbedtls_mpi_init( &U ); mbedtls_mpi_init( &V );
+    mbedtls_mpi_init_mempool( &A ); mbedtls_mpi_init_mempool( &E );
+    mbedtls_mpi_init_mempool( &N ); mbedtls_mpi_init_mempool( &X );
+    mbedtls_mpi_init_mempool( &Y ); mbedtls_mpi_init_mempool( &U );
+    mbedtls_mpi_init_mempool( &V );
 
     MBEDTLS_MPI_CHK( mbedtls_mpi_read_string( &A, 16,
         "EFE021C2645FD1DC586E69184AF4A31E" \