From 0020df9cf929119eb322784a8608039533f645b9 Mon Sep 17 00:00:00 2001
From: Janos Follath <janos.follath@arm.com>
Date: Sat, 26 Nov 2022 17:23:16 +0000
Subject: [PATCH] mpi_mod_io: test with various buffer sizes

Signed-off-by: Janos Follath <janos.follath@arm.com>
---
 tests/suites/test_suite_bignum_mod.function | 70 ++++++++++++++++++---
 1 file changed, 60 insertions(+), 10 deletions(-)

diff --git a/tests/suites/test_suite_bignum_mod.function b/tests/suites/test_suite_bignum_mod.function
index df6bb45f6..8fdd7b986 100644
--- a/tests/suites/test_suite_bignum_mod.function
+++ b/tests/suites/test_suite_bignum_mod.function
@@ -188,7 +188,8 @@ void mpi_mod_io( char * input_N, data_t * input_A, int endian )
     mbedtls_mpi_uint *N = NULL;
     mbedtls_mpi_uint *R = NULL;
     mbedtls_mpi_uint *R_COPY = NULL;
-    unsigned char *r_buff = NULL;
+    unsigned char *obuf = NULL;
+    unsigned char *ref_buf = NULL;
     mbedtls_mpi_mod_modulus m;
     mbedtls_mpi_mod_residue r;
     mbedtls_mpi_mod_residue r_copy;
@@ -204,7 +205,6 @@ void mpi_mod_io( char * input_N, data_t * input_A, int endian )
     /* Allocate the memory for intermediate data structures */
     ASSERT_ALLOC( R, n_bytes );
     ASSERT_ALLOC( R_COPY, n_bytes );
-    ASSERT_ALLOC( r_buff, a_bytes );
 
     /* Test that input's size is not greater to modulo's */
     TEST_LE_U( a_bytes, n_bytes );
@@ -219,22 +219,72 @@ void mpi_mod_io( char * input_N, data_t * input_A, int endian )
     TEST_EQUAL( 0, mbedtls_mpi_mod_read( &r, &m, input_A->x, input_A->len,
                                          endian ) );
 
-    TEST_EQUAL( 0, mbedtls_mpi_mod_write( &r, &m, r_buff, a_bytes,
-                                          endian ) );
-
-    /* Make sure that writing didn't change the value of r */
-    TEST_EQUAL( 0, mbedtls_mpi_mod_residue_setup( &r_copy, &m, R_COPY, n_limbs ) );
+    /* Read a copy for checking that writing didn't change the value of r */
+    TEST_EQUAL( 0, mbedtls_mpi_mod_residue_setup( &r_copy, &m,
+                                                  R_COPY, n_limbs ) );
     TEST_EQUAL( 0, mbedtls_mpi_mod_read( &r_copy, &m, input_A->x, input_A->len,
                                          endian ) );
-    ASSERT_COMPARE( r.p, r.limbs, r_copy.p, r_copy.limbs );
 
-    ASSERT_COMPARE( r_buff, a_bytes, input_A->x, a_bytes );
+    /* Get number of bytes without leading zeroes */
+    size_t a_bytes_trimmed = a_bytes;
+    while( a_bytes_trimmed > 0 )
+    {
+        unsigned char* r_byte_array = (unsigned char*) r.p;
+        if( r_byte_array[--a_bytes_trimmed] != 0 )
+            break;
+    }
+    a_bytes_trimmed++;
+
+    /* Test write with three output buffer sizes: tight, same as input and
+     * longer than the input */
+    size_t obuf_sizes[3];
+    const size_t obuf_sizes_len = sizeof( obuf_sizes ) / sizeof( obuf_sizes[0] );
+    obuf_sizes[0] = a_bytes_trimmed;
+    obuf_sizes[1] = a_bytes;
+    obuf_sizes[2] = a_bytes + 8;
+
+    for( size_t i = 0; i < obuf_sizes_len; i++ )
+    {
+        ASSERT_ALLOC( obuf, obuf_sizes[i] );
+        TEST_EQUAL( 0, mbedtls_mpi_mod_write( &r, &m, obuf, obuf_sizes[i], endian ) );
+
+        /* Make sure that writing didn't corrupt the value of r */
+        ASSERT_COMPARE( r.p, r.limbs, r_copy.p, r_copy.limbs );
+
+        /* Set up reference output for checking the result */
+        ASSERT_ALLOC( ref_buf, obuf_sizes[i] );
+        switch( endian )
+        {
+            case MBEDTLS_MPI_MOD_EXT_REP_LE:
+                memcpy( ref_buf, input_A->x, a_bytes_trimmed );
+                break;
+            case MBEDTLS_MPI_MOD_EXT_REP_BE:
+                {
+                    size_t a_offset = input_A->len - a_bytes_trimmed;
+                    size_t ref_offset = obuf_sizes[i] - a_bytes_trimmed;
+                    memcpy( ref_buf + ref_offset, input_A->x + a_offset,
+                            a_bytes_trimmed );
+                }
+                break;
+            default:
+                TEST_ASSERT( 0 );
+        }
+
+        /* Check the result */
+        ASSERT_COMPARE( obuf, obuf_sizes[i], ref_buf, obuf_sizes[i] );
+
+        mbedtls_free( ref_buf );
+        ref_buf = NULL;
+        mbedtls_free( obuf );
+        obuf = NULL;
+    }
+
 exit:
     mbedtls_mpi_mod_modulus_free( &m );
     mbedtls_free( N );
     mbedtls_free( R );
     mbedtls_free( R_COPY );
-    mbedtls_free( r_buff );
+    mbedtls_free( obuf );
 }
 /* END_CASE */
 /* END MERGE SLOT 7 */