1 files changed, 47 insertions, 13 deletions

diff --git a/library/dhm.c b/library/dhm.c
index cff0958..28ac310 100644
--- a/library/dhm.c
+++ b/library/dhm.c
@@ -94,6 +94,9 @@ static int dhm_read_bignum( mbedtls_mpi *X,
  *
  * Parameter should be: 2 <= public_param <= P - 2
  *
+ * This means that we need to return an error if
+ *              public_param < 2 or public_param > P-2
+ *
  * For more information on the attack, see:
  *  http://www.cl.cam.ac.uk/~rja14/Papers/psandqs.pdf
  *  http://web.nvd.nist.gov/view/vuln/detail?vulnId=CVE-2005-2643
@@ -101,17 +104,17 @@ static int dhm_read_bignum( mbedtls_mpi *X,
 static int dhm_check_range( const mbedtls_mpi *param, const mbedtls_mpi *P )
 {
     mbedtls_mpi L, U;
-    int ret = MBEDTLS_ERR_DHM_BAD_INPUT_DATA;
+    int ret = 0;
 
     mbedtls_mpi_init( &L ); mbedtls_mpi_init( &U );
 
     MBEDTLS_MPI_CHK( mbedtls_mpi_lset( &L, 2 ) );
     MBEDTLS_MPI_CHK( mbedtls_mpi_sub_int( &U, P, 2 ) );
 
-    if( mbedtls_mpi_cmp_mpi( param, &L ) >= 0 &&
-        mbedtls_mpi_cmp_mpi( param, &U ) <= 0 )
+    if( mbedtls_mpi_cmp_mpi( param, &L ) < 0 ||
+        mbedtls_mpi_cmp_mpi( param, &U ) > 0 )
     {
-        ret = 0;
+        ret = MBEDTLS_ERR_DHM_BAD_INPUT_DATA;
     }
 
 cleanup:
@@ -188,10 +191,15 @@ int mbedtls_dhm_make_params( mbedtls_dhm_context *ctx, int x_size,
     /*
      * export P, G, GX
      */
-#define DHM_MPI_EXPORT(X,n)                     \
-    MBEDTLS_MPI_CHK( mbedtls_mpi_write_binary( X, p + 2, n ) ); \
-    *p++ = (unsigned char)( n >> 8 );           \
-    *p++ = (unsigned char)( n      ); p += n;
+#define DHM_MPI_EXPORT( X, n )                                          \
+    do {                                                                \
+        MBEDTLS_MPI_CHK( mbedtls_mpi_write_binary( ( X ),               \
+                                                   p + 2,               \
+                                                   ( n ) ) );           \
+        *p++ = (unsigned char)( ( n ) >> 8 );                           \
+        *p++ = (unsigned char)( ( n )      );                           \
+        p += ( n );                                                     \
+    } while( 0 )
 
     n1 = mbedtls_mpi_size( &ctx->P  );
     n2 = mbedtls_mpi_size( &ctx->G  );
@@ -202,7 +210,7 @@ int mbedtls_dhm_make_params( mbedtls_dhm_context *ctx, int x_size,
     DHM_MPI_EXPORT( &ctx->G , n2 );
     DHM_MPI_EXPORT( &ctx->GX, n3 );
 
-    *olen  = p - output;
+    *olen = p - output;
 
     ctx->len = n1;
 
@@ -215,6 +223,28 @@ cleanup:
 }
 
 /*
+ * Set prime modulus and generator
+ */
+int mbedtls_dhm_set_group( mbedtls_dhm_context *ctx,
+                           const mbedtls_mpi *P,
+                           const mbedtls_mpi *G )
+{
+    int ret;
+
+    if( ctx == NULL || P == NULL || G == NULL )
+        return( MBEDTLS_ERR_DHM_BAD_INPUT_DATA );
+
+    if( ( ret = mbedtls_mpi_copy( &ctx->P, P ) ) != 0 ||
+        ( ret = mbedtls_mpi_copy( &ctx->G, G ) ) != 0 )
+    {
+        return( MBEDTLS_ERR_DHM_SET_GROUP_FAILED + ret );
+    }
+
+    ctx->len = mbedtls_mpi_size( &ctx->P );
+    return( 0 );
+}
+
+/*
  * Import the peer's public value G^Y
  */
 int mbedtls_dhm_read_public( mbedtls_dhm_context *ctx,
@@ -401,10 +431,11 @@ cleanup:
  */
 void mbedtls_dhm_free( mbedtls_dhm_context *ctx )
 {
-    mbedtls_mpi_free( &ctx->pX); mbedtls_mpi_free( &ctx->Vf ); mbedtls_mpi_free( &ctx->Vi );
-    mbedtls_mpi_free( &ctx->RP ); mbedtls_mpi_free( &ctx->K ); mbedtls_mpi_free( &ctx->GY );
-    mbedtls_mpi_free( &ctx->GX ); mbedtls_mpi_free( &ctx->X ); mbedtls_mpi_free( &ctx->G );
-    mbedtls_mpi_free( &ctx->P );
+    mbedtls_mpi_free( &ctx->pX ); mbedtls_mpi_free( &ctx->Vf );
+    mbedtls_mpi_free( &ctx->Vi ); mbedtls_mpi_free( &ctx->RP );
+    mbedtls_mpi_free( &ctx->K  ); mbedtls_mpi_free( &ctx->GY );
+    mbedtls_mpi_free( &ctx->GX ); mbedtls_mpi_free( &ctx->X  );
+    mbedtls_mpi_free( &ctx->G  ); mbedtls_mpi_free( &ctx->P  );
 
     mbedtls_zeroize( ctx, sizeof( mbedtls_dhm_context ) );
 }
@@ -543,7 +574,10 @@ static int load_file( const char *path, unsigned char **buf, size_t *n )
     if( fread( *buf, 1, *n, f ) != *n )
     {
         fclose( f );
+
+        mbedtls_zeroize( *buf, *n + 1 );
         mbedtls_free( *buf );
+
         return( MBEDTLS_ERR_DHM_FILE_IO_ERROR );
     }