Fix mbedtls_mpi_random when N has leading zeros

mbedtls_mpi_random() uses mbedtls_mpi_cmp_mpi_ct(), which requires its
two arguments to have the same storage size. This was not the case
when the upper bound passed to mbedtls_mpi_random() had leading zero
limbs.

Fix this by forcing the result MPI to the desired size. Since this is
not what mbedtls_mpi_fill_random() does, don't call it from
mbedtls_mpi_random(), but instead call a new auxiliary function.

Add tests to cover this and other conditions with varying sizes for
the two arguments.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/library/bignum.c b/library/bignum.c
index fba9acc..25fb9ae 100644
--- a/library/bignum.c
+++ b/library/bignum.c
@@ -2395,6 +2395,30 @@
     return( ret );
 }
 
+/* Fill X with n_bytes random bytes.
+ * X must already have room for those bytes.
+ * n_bytes must not be 0.
+ */
+static int mpi_fill_random_internal(
+    mbedtls_mpi *X, size_t n_bytes,
+    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    const size_t limbs = CHARS_TO_LIMBS( n_bytes );
+    const size_t overhead = ( limbs * ciL ) - n_bytes;
+
+    if( X->n < limbs )
+        return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
+    MBEDTLS_MPI_CHK( mbedtls_mpi_lset( X, 0 ) );
+    MBEDTLS_MPI_CHK( mbedtls_mpi_grow( X, limbs ) );
+
+    MBEDTLS_MPI_CHK( f_rng( p_rng, (unsigned char *) X->p + overhead, n_bytes ) );
+    mpi_bigendian_to_host( X->p, limbs );
+
+cleanup:
+    return( ret );
+}
+
 /*
  * Fill X with size bytes of random.
  *
@@ -2408,8 +2432,6 @@
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     size_t const limbs = CHARS_TO_LIMBS( size );
-    size_t const overhead = ( limbs * ciL ) - size;
-    unsigned char *Xp;
 
     MPI_VALIDATE_RET( X     != NULL );
     MPI_VALIDATE_RET( f_rng != NULL );
@@ -2421,12 +2443,10 @@
         mbedtls_mpi_init( X );
         MBEDTLS_MPI_CHK( mbedtls_mpi_grow( X, limbs ) );
     }
-    MBEDTLS_MPI_CHK( mbedtls_mpi_lset( X, 0 ) );
+    if( size == 0 )
+        return( 0 );
 
-    Xp = (unsigned char*) X->p;
-    MBEDTLS_MPI_CHK( f_rng( p_rng, Xp + overhead, size ) );
-
-    mpi_bigendian_to_host( X->p, limbs );
+    ret = mpi_fill_random_internal( X, size, f_rng, p_rng );
 
 cleanup:
     return( ret );
@@ -2450,6 +2470,16 @@
     if( mbedtls_mpi_cmp_int( N, min ) <= 0 )
         return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
 
+    /* Ensure that target MPI has exactly the same number of limbs
+     * as the upper bound, even if the upper bound has leading zeros.
+     * This is necessary for the mbedtls_mpi_lt_mpi_ct() check. */
+    if( X->n != N->n )
+    {
+        mbedtls_mpi_free( X );
+        mbedtls_mpi_init( X );
+        MBEDTLS_MPI_CHK( mbedtls_mpi_grow( X, N->n ) );
+    }
+
     /*
      * Match the procedure given in RFC 6979 §3.3 (deterministic ECDSA)
      * when f_rng is a suitably parametrized instance of HMAC_DRBG:
@@ -2460,7 +2490,7 @@
      */
     do
     {
-        MBEDTLS_MPI_CHK( mbedtls_mpi_fill_random( X, n_bytes, f_rng, p_rng ) );
+        MBEDTLS_MPI_CHK( mpi_fill_random_internal( X, n_bytes, f_rng, p_rng ) );
         MBEDTLS_MPI_CHK( mbedtls_mpi_shift_r( X, 8 * n_bytes - n_bits ) );
 
         /*
diff --git a/tests/suites/test_suite_mpi.data b/tests/suites/test_suite_mpi.data
index fb2c553..3488854 100644
--- a/tests/suites/test_suite_mpi.data
+++ b/tests/suites/test_suite_mpi.data
@@ -1132,6 +1132,33 @@
 MPI random in range: 3..4
 mpi_random_many:1:"04":1000
 
+MPI random in range: smaller result
+mpi_random_grown:1:"aaaaaaaaaaaaaaaabbbbbbbbbbbbbbbb":1
+
+MPI random in range: same size result (32-bit limbs)
+mpi_random_grown:1:"aaaaaaaaaaaaaaaa":2
+
+MPI random in range: same size result (64-bit limbs)
+mpi_random_grown:1:"aaaaaaaaaaaaaaaa":1
+
+MPI random in range: larger result
+mpi_random_grown:1:"aaaaaaaaaaaaaaaa":3
+
+MPI random in range: leading 0 limb in upper bound #0
+mpi_random_grown:1:"00aaaaaaaaaaaaaaaa":0
+
+MPI random in range: leading 0 limb in upper bound #1
+mpi_random_grown:1:"00aaaaaaaaaaaaaaaa":1
+
+MPI random in range: leading 0 limb in upper bound #2
+mpi_random_grown:1:"00aaaaaaaaaaaaaaaa":2
+
+MPI random in range: leading 0 limb in upper bound #3
+mpi_random_grown:1:"00aaaaaaaaaaaaaaaa":3
+
+MPI random in range: leading 0 limb in upper bound #4
+mpi_random_grown:1:"00aaaaaaaaaaaaaaaa":4
+
 MPI random bad arguments: min < 0
 mpi_random_fail:-1:"04":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
 
diff --git a/tests/suites/test_suite_mpi.function b/tests/suites/test_suite_mpi.function
index 0ca6b64..f3c9107 100644
--- a/tests/suites/test_suite_mpi.function
+++ b/tests/suites/test_suite_mpi.function
@@ -1538,6 +1538,29 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void mpi_random_grown( int min, data_t *bound_bytes, int nlimbs )
+{
+    mbedtls_mpi upper_bound;
+    mbedtls_mpi result;
+
+    mbedtls_mpi_init( &upper_bound );
+    mbedtls_mpi_init( &result );
+
+    TEST_EQUAL( 0, mbedtls_mpi_grow( &result, nlimbs ) );
+    TEST_EQUAL( 0, mbedtls_mpi_read_binary( &upper_bound,
+                                            bound_bytes->x, bound_bytes->len ) );
+    TEST_EQUAL( 0, mbedtls_mpi_random( &result, min, &upper_bound,
+                                       mbedtls_test_rnd_std_rand, NULL ) );
+    TEST_ASSERT( mbedtls_mpi_cmp_mpi( &result, &upper_bound ) < 0 );
+    TEST_ASSERT( mbedtls_mpi_cmp_int( &result, min ) >= 0 );
+
+exit:
+    mbedtls_mpi_free( &upper_bound );
+    mbedtls_mpi_free( &result );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void mpi_random_fail( int min, data_t *bound_bytes, int expected_ret )
 {
     mbedtls_mpi upper_bound;