From 8dfc8c41b7fb12a42d2828e88943850ec69e9480 Mon Sep 17 00:00:00 2001
From: Janos Follath <janos.follath@arm.com>
Date: Sat, 26 Nov 2022 15:39:02 +0000
Subject: [PATCH] mbedtls_mpi_mod_write: prevent data corruption

The function wasn't converting back data to internal representation when
writing it out.

Signed-off-by: Janos Follath <janos.follath@arm.com>
---
 library/bignum_mod.c                        | 16 ++++++++++++++--
 tests/suites/test_suite_bignum_mod.function | 10 ++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/library/bignum_mod.c b/library/bignum_mod.c
index 7f7c71512e..4fe6e48547 100644
--- a/library/bignum_mod.c
+++ b/library/bignum_mod.c
@@ -231,6 +231,7 @@ int mbedtls_mpi_mod_write( const mbedtls_mpi_mod_residue *r,
                            mbedtls_mpi_mod_ext_rep ext_rep )
 {
     int ret = MBEDTLS_ERR_MPI_BAD_INPUT_DATA;
+    int conv_ret = 0;
 
     /* Do our best to check if r and m have been set up */
     if ( r->limbs == 0 || m->limbs == 0 )
@@ -238,12 +239,23 @@ int mbedtls_mpi_mod_write( const mbedtls_mpi_mod_residue *r,
     if ( r->limbs != m->limbs )
         goto cleanup;
 
-    if ( m->int_rep == MBEDTLS_MPI_MOD_REP_MONTGOMERY)
-        ret = mbedtls_mpi_mod_raw_from_mont_rep( r->p, m );
+    if ( m->int_rep == MBEDTLS_MPI_MOD_REP_MONTGOMERY )
+    {
+        conv_ret = mbedtls_mpi_mod_raw_from_mont_rep( r->p, m );
+        if( conv_ret != 0 )
+            goto cleanup;
+    }
 
     ret = mbedtls_mpi_mod_raw_write( r->p, m, buf, buflen, ext_rep );
 
+    if ( m->int_rep == MBEDTLS_MPI_MOD_REP_MONTGOMERY )
+        conv_ret = mbedtls_mpi_mod_raw_to_mont_rep( r->p, m );
+
 cleanup:
+
+    if ( ret == 0 )
+        ret = conv_ret;
+
     return ( ret );
 }
 /* END MERGE SLOT 7 */
diff --git a/tests/suites/test_suite_bignum_mod.function b/tests/suites/test_suite_bignum_mod.function
index 7042ed3d2b..df6bb45f6f 100644
--- a/tests/suites/test_suite_bignum_mod.function
+++ b/tests/suites/test_suite_bignum_mod.function
@@ -187,9 +187,11 @@ 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;
     mbedtls_mpi_mod_modulus m;
     mbedtls_mpi_mod_residue r;
+    mbedtls_mpi_mod_residue r_copy;
     size_t n_limbs, n_bytes, a_bytes;
 
     mbedtls_mpi_mod_modulus_init( &m );
@@ -201,6 +203,7 @@ 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 */
@@ -219,11 +222,18 @@ void mpi_mod_io( char * input_N, data_t * input_A, int 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 ) );
+    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 );
 exit:
     mbedtls_mpi_mod_modulus_free( &m );
     mbedtls_free( N );
     mbedtls_free( R );
+    mbedtls_free( R_COPY );
     mbedtls_free( r_buff );
 }
 /* END_CASE */