Adapt key checking functions for Curve25519
diff --git a/library/ecp.c b/library/ecp.c
index 967da94..75f5b4d 100644
--- a/library/ecp.c
+++ b/library/ecp.c
@@ -177,6 +177,14 @@
 }
 
 /*
+ * Tell if a group structure is associated to a Montgomery curve
+ */
+static inline int ecp_is_montgomery( const ecp_group *grp )
+{
+    return( grp->G.X.p != NULL && grp->G.Y.p == NULL );
+}
+
+/*
  * Initialize (the components of) a point
  */
 void ecp_point_init( ecp_point *pt )
@@ -1498,22 +1506,31 @@
 }
 
 /*
- * Check that a point is valid as a public key (SEC1 3.2.3.1)
+ * Check that a point is valid as a public key
  */
 int ecp_check_pubkey( const ecp_group *grp, const ecp_point *pt )
 {
     int ret;
     mpi YY, RHS;
 
-    if( mpi_cmp_int( &pt->Z, 0 ) == 0 )
-        return( POLARSSL_ERR_ECP_INVALID_KEY );
-
-    /*
-     * pt coordinates must be normalized for our checks
-     */
+    /* Must use affine coordinates */
     if( mpi_cmp_int( &pt->Z, 1 ) != 0 )
         return( POLARSSL_ERR_ECP_INVALID_KEY );
 
+    if( ecp_is_montgomery( grp ) )
+    {
+        /* Just check X is the correct number of bytes */
+        if( mpi_size( &pt->X ) > ( grp->nbits + 7 ) / 8 )
+            return( POLARSSL_ERR_ECP_INVALID_KEY );
+
+        return( 0 );
+    }
+
+    /*
+     * Ok, so we have a short Weierstrass curve as in SEC1 3.2.3.1
+     */
+
+    /* pt coordinates must be normalized for our checks */
     if( mpi_cmp_int( &pt->X, 0 ) < 0 ||
         mpi_cmp_int( &pt->Y, 0 ) < 0 ||
         mpi_cmp_mpi( &pt->X, &grp->P ) >= 0 ||
@@ -1522,6 +1539,12 @@
 
     mpi_init( &YY ); mpi_init( &RHS );
 
+    if( mpi_cmp_int( &pt->X, 0 ) < 0 ||
+        mpi_cmp_int( &pt->Y, 0 ) < 0 ||
+        mpi_cmp_mpi( &pt->X, &grp->P ) >= 0 ||
+        mpi_cmp_mpi( &pt->Y, &grp->P ) >= 0 )
+        return( POLARSSL_ERR_ECP_INVALID_KEY );
+
     /*
      * YY = Y^2
      * RHS = X (X^2 + A) + B = X^3 + A X + B
@@ -1543,13 +1566,26 @@
 }
 
 /*
- * Check that an mpi is valid as a private key (SEC1 3.2)
+ * Check that an mpi is valid as a private key
  */
 int ecp_check_privkey( const ecp_group *grp, const mpi *d )
 {
-    /* We want 1 <= d <= N-1 */
-    if ( mpi_cmp_int( d, 1 ) < 0 || mpi_cmp_mpi( d, &grp->N ) >= 0 )
-        return( POLARSSL_ERR_ECP_INVALID_KEY );
+    if( ecp_is_montgomery( grp ) )
+    {
+        /* see the Curve25519 paper */
+        if( mpi_get_bit( d, 0 ) != 0 ||
+            mpi_get_bit( d, 1 ) != 0 ||
+            mpi_get_bit( d, 2 ) != 0 ||
+            mpi_msb( d ) - 1 != grp->nbits ) /* mpi_msb is one-based! */
+            return( POLARSSL_ERR_ECP_INVALID_KEY );
+    }
+    else
+    {
+        /* see SEC1 3.2 */
+        if( mpi_cmp_int( d, 1 ) < 0 ||
+            mpi_cmp_mpi( d, &grp->N ) >= 0 )
+            return( POLARSSL_ERR_ECP_INVALID_KEY );
+    }
 
     return( 0 );
 }
diff --git a/library/ecp_curves.c b/library/ecp_curves.c
index 10a179f..542589f 100644
--- a/library/ecp_curves.c
+++ b/library/ecp_curves.c
@@ -357,6 +357,12 @@
     MPI_CHK( mpi_sub_int( &grp->P, &grp->P, 19 ) );
     grp->pbits = mpi_msb( &grp->P );
 
+    /* Y intentionaly not set, since we use x/z coordinates.
+     * This is used as a marker to identify Montgomery curves! */
+    MPI_CHK( mpi_lset( &grp->G.X, 9 ) );
+    MPI_CHK( mpi_lset( &grp->G.Z, 1 ) );
+    mpi_free( &grp->G.Y );
+
     /* Actually, the required msb for private keys */
     grp->nbits = 254;
 
diff --git a/tests/suites/test_suite_ecp.data b/tests/suites/test_suite_ecp.data
index 17ad8aa..2edc519 100644
--- a/tests/suites/test_suite_ecp.data
+++ b/tests/suites/test_suite_ecp.data
@@ -159,6 +159,12 @@
 ECP small check pubkey #10
 ecp_small_check_pub:10:25:1:POLARSSL_ERR_ECP_INVALID_KEY
 
+ECP check pubkey Montgomery #1 (too big)
+ecp_check_pub_mx:POLARSSL_ECP_DP_M255:"010000000000000000000000000000000000000000000000000000000000000000":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check pubkey Montgomery #2 (biggest)
+ecp_check_pub_mx:POLARSSL_ECP_DP_M255:"FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF":0
+
 ECP write binary #0 (zero, bad format)
 depends_on:POLARSSL_ECP_DP_SECP192R1_ENABLED
 ecp_write_binary:POLARSSL_ECP_DP_SECP192R1:"01":"01":"00":POLARSSL_ECP_PF_UNKNOWN:"00":1:POLARSSL_ERR_ECP_BAD_INPUT_DATA
@@ -271,9 +277,49 @@
 depends_on:POLARSSL_ECP_DP_SECP521R1_ENABLED
 ecp_tls_write_read_group:POLARSSL_ECP_DP_SECP521R1
 
-ECP check privkey
+ECP check privkey #1 (short weierstrass, too small)
 depends_on:POLARSSL_ECP_DP_SECP192R1_ENABLED
-ecp_check_privkey:POLARSSL_ECP_DP_SECP192R1
+ecp_check_privkey:POLARSSL_ECP_DP_SECP192R1:"00":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check privkey #2 (short weierstrass, smallest)
+depends_on:POLARSSL_ECP_DP_SECP192R1_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_SECP192R1:"01":0
+
+ECP check privkey #3 (short weierstrass, biggest)
+depends_on:POLARSSL_ECP_DP_SECP192R1_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_SECP192R1:"FFFFFFFFFFFFFFFFFFFFFFFF99DEF836146BC9B1B4D22830":0
+
+ECP check privkey #4 (short weierstrass, too big)
+depends_on:POLARSSL_ECP_DP_SECP192R1_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_SECP192R1:"FFFFFFFFFFFFFFFFFFFFFFFF99DEF836146BC9B1B4D22831":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check privkey #5 (montgomery, too big)
+depends_on:POLARSSL_ECP_DP_M255_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_M255:"C000000000000000000000000000000000000000000000000000000000000000":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check privkey #6 (montgomery, not big enough)
+depends_on:POLARSSL_ECP_DP_M255_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_M255:"3FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF0":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check privkey #7 (montgomery, msb OK)
+depends_on:POLARSSL_ECP_DP_M255_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_M255:"4000000000000000000000000000000000000000000000000000000000000000":0
+
+ECP check privkey #8 (montgomery, bit 0 set)
+depends_on:POLARSSL_ECP_DP_M255_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_M255:"4000000000000000000000000000000000000000000000000000000000000001":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check privkey #9 (montgomery, bit 1 set)
+depends_on:POLARSSL_ECP_DP_M255_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_M255:"4000000000000000000000000000000000000000000000000000000000000002":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check privkey #10 (montgomery, bit 2 set)
+depends_on:POLARSSL_ECP_DP_M255_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_M255:"4000000000000000000000000000000000000000000000000000000000000004":POLARSSL_ERR_ECP_INVALID_KEY
+
+ECP check privkey #11 (montgomery, OK)
+depends_on:POLARSSL_ECP_DP_M255_ENABLED
+ecp_check_privkey:POLARSSL_ECP_DP_M255:"7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF8":0
 
 ECP gen keypair
 depends_on:POLARSSL_ECP_DP_SECP192R1_ENABLED
diff --git a/tests/suites/test_suite_ecp.function b/tests/suites/test_suite_ecp.function
index cb93e85..626d2bf 100644
--- a/tests/suites/test_suite_ecp.function
+++ b/tests/suites/test_suite_ecp.function
@@ -189,6 +189,27 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void ecp_check_pub_mx( int grp_id, char *key_hex, int ret )
+{
+    ecp_group grp;
+    ecp_point P;
+
+    ecp_group_init( &grp );
+    ecp_point_init( &P );
+
+    TEST_ASSERT( ecp_use_known_dp( &grp, grp_id ) == 0 );
+
+    TEST_ASSERT( mpi_read_string( &P.X, 16, key_hex ) == 0 );
+    TEST_ASSERT( mpi_lset( &P.Z, 1 ) == 0 );
+
+    TEST_ASSERT( ecp_check_pubkey( &grp, &P ) == ret );
+
+    ecp_group_free( &grp );
+    ecp_point_free( &P );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void ecp_test_vect( int id, char *dA_str, char *xA_str, char *yA_str,
                     char *dB_str, char *xB_str, char *yB_str, char *xZ_str,
                     char *yZ_str )
@@ -490,7 +511,7 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void ecp_check_privkey( int id )
+void ecp_check_privkey( int id, char *key_hex, int ret )
 {
     ecp_group grp;
     mpi d;
@@ -499,12 +520,9 @@
     mpi_init( &d );
 
     TEST_ASSERT( ecp_use_known_dp( &grp, id ) == 0 );
+    TEST_ASSERT( mpi_read_string( &d, 16, key_hex ) == 0 );
 
-    TEST_ASSERT( mpi_lset( &d, 0 ) == 0 );
-    TEST_ASSERT( ecp_check_privkey( &grp, &d ) == POLARSSL_ERR_ECP_INVALID_KEY );
-
-    TEST_ASSERT( mpi_copy( &d, &grp.N ) == 0 );
-    TEST_ASSERT( ecp_check_privkey( &grp, &d ) == POLARSSL_ERR_ECP_INVALID_KEY );
+    TEST_ASSERT( ecp_check_privkey( &grp, &d ) == ret );
 
     ecp_group_free( &grp );
     mpi_free( &d );