- Updated sanity checks

diff --git a/library/dhm.c b/library/dhm.c
index 0e76f0b..b587fa6 100644
--- a/library/dhm.c
+++ b/library/dhm.c
@@ -63,6 +63,37 @@
 }
 
 /*
+ * Verify sanity of public value with regards to P
+ */
+static int dhm_verifypub( const mpi *P,  const mpi *pub_value )
+{
+    mpi X;
+
+    mpi_init( &X, NULL );
+    mpi_lset( &X, 1 );
+
+    /* Check G^Y or G^X is valid */
+    if( mpi_cmp_mpi( pub_value, &X ) <= 0 )
+    {
+        mpi_free( &X, NULL );
+        return( POLARSSL_ERR_DHM_BAD_INPUT_DATA );
+    }
+
+    /* Reset: x = P - 1 */
+    mpi_sub_int( &X, P, 1 );
+
+    if( mpi_cmp_mpi( pub_value, &X ) >= 0 )
+    {
+        mpi_free( &X, NULL );
+        return( POLARSSL_ERR_DHM_BAD_INPUT_DATA );
+    }
+
+    mpi_free( &X, NULL );
+
+    return( 0 );
+}
+
+/*
  * Parse the ServerKeyExchange parameters
  */
 int dhm_read_params( dhm_context *ctx,
@@ -89,6 +120,9 @@
     if( end != *p + n )
         return( POLARSSL_ERR_DHM_BAD_INPUT_DATA );
 
+    if( ( ret = dhm_verifypub( &ctx->P, &ctx->GY ) ) != 0 )
+        return( ret );
+
     return( 0 );
 }
 
@@ -105,12 +139,12 @@
     /*
      * Generate X as large as possible ( < P )
      */
-    n = x_size / sizeof( t_int );
+    n = x_size / sizeof( t_int ) + 1;
     MPI_CHK( mpi_grow( &ctx->X, n ) );
     MPI_CHK( mpi_lset( &ctx->X, 0 ) );
 
     p = (unsigned char *) ctx->X.p;
-    for( i = 0; i < x_size - 1; i++ )
+    for( i = 0; i < x_size; i++ )
         *p++ = (unsigned char) f_rng( p_rng );
 
     while( mpi_cmp_mpi( &ctx->X, &ctx->P ) >= 0 )
@@ -122,6 +156,9 @@
     MPI_CHK( mpi_exp_mod( &ctx->GX, &ctx->G, &ctx->X,
                           &ctx->P , &ctx->RP ) );
 
+    if( ( ret = dhm_verifypub( &ctx->P, &ctx->GX ) ) != 0 )
+        return( ret );
+
     /*
      * export P, G, GX
      */
@@ -184,13 +221,12 @@
     /*
      * generate X and calculate GX = G^X mod P
      */
-    n = x_size / sizeof( t_int );
+    n = x_size / sizeof( t_int ) + 1;
     MPI_CHK( mpi_grow( &ctx->X, n ) );
     MPI_CHK( mpi_lset( &ctx->X, 0 ) );
 
-    n = x_size - 1;
     p = (unsigned char *) ctx->X.p;
-    for( i = 0; i < n; i++ )
+    for( i = 0; i < x_size; i++ )
         *p++ = (unsigned char) f_rng( p_rng );
 
     while( mpi_cmp_mpi( &ctx->X, &ctx->P ) >= 0 )
@@ -199,6 +235,9 @@
     MPI_CHK( mpi_exp_mod( &ctx->GX, &ctx->G, &ctx->X,
                           &ctx->P , &ctx->RP ) );
 
+    if( dhm_verifypub( &ctx->P, &ctx->GX ) != 0 )
+        return( POLARSSL_ERR_DHM_MAKE_PUBLIC_FAILED );
+
     MPI_CHK( mpi_write_binary( &ctx->GX, output, olen ) );
 
 cleanup:
@@ -223,6 +262,9 @@
     MPI_CHK( mpi_exp_mod( &ctx->K, &ctx->GY, &ctx->X,
                           &ctx->P, &ctx->RP ) );
 
+    if( ( ret = dhm_verifypub( &ctx->P, &ctx->GY ) ) != 0 )
+        return( ret );
+
     *olen = mpi_size( &ctx->K );
 
     MPI_CHK( mpi_write_binary( &ctx->K, output, *olen ) );