From 8b3336331507ee2ccc49e3a13c73fde1820cff5a Mon Sep 17 00:00:00 2001
From: Minos Galanakis <minos.galanakis@arm.com>
Date: Tue, 11 Oct 2022 11:28:24 +0100
Subject: [PATCH] bignum_mod: Updated modulus lifecycle with mm and rr.

This patch updates the `mbedtls_mpi_mod_modulus_setup/free()`
methods to precalculate mm and rr(Montgomery const squared) during
setup and zeroize it during free.

A static `set_mont_const_square()` is added to manage the memory allocation
and parameter checking before invoking the
`mbedtls_mpi_core_get_mont_r2_unsafe()`

Signed-off-by: Minos Galanakis <minos.galanakis@arm.com>
---
 library/bignum_mod.c | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/library/bignum_mod.c b/library/bignum_mod.c
index 7cf1b012c5..92c011cffe 100644
--- a/library/bignum_mod.c
+++ b/library/bignum_mod.c
@@ -77,6 +77,9 @@ void mbedtls_mpi_mod_modulus_free( mbedtls_mpi_mod_modulus *m )
     switch( m->int_rep )
     {
         case MBEDTLS_MPI_MOD_REP_MONTGOMERY:
+            mbedtls_platform_zeroize( (mbedtls_mpi_uint *) m->rep.mont.rr,
+                                      m->limbs );
+            mbedtls_free( (mbedtls_mpi_uint *)m->rep.mont.rr );
             m->rep.mont.rr = NULL;
             m->rep.mont.mm = 0; break;
         case MBEDTLS_MPI_MOD_REP_OPT_RED:
@@ -93,6 +96,38 @@ void mbedtls_mpi_mod_modulus_free( mbedtls_mpi_mod_modulus *m )
     m->int_rep = MBEDTLS_MPI_MOD_REP_INVALID;
 }
 
+static int set_mont_const_square( const mbedtls_mpi_uint **X,
+                                  const mbedtls_mpi_uint *A,
+                                  size_t limbs )
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    mbedtls_mpi N;
+    mbedtls_mpi RR;
+
+    mbedtls_mpi_init( &N );
+    mbedtls_mpi_init( &RR );
+
+    if ( A == NULL || limbs == 0 || limbs >= ( MBEDTLS_MPI_MAX_LIMBS / 2 ) - 2 )
+        goto cleanup;
+
+    if ( !mbedtls_mpi_grow( &N,  limbs ))
+        memcpy( N.p, A, sizeof(mbedtls_mpi_uint) *  limbs );
+    else
+        goto cleanup;
+
+    mbedtls_mpi_core_get_mont_r2_unsafe(&RR, &N);
+
+    *X = RR.p;
+    RR.p = NULL;
+    ret = 0;
+
+cleanup:
+    mbedtls_mpi_free(&N);
+    mbedtls_mpi_free(&RR);
+    ret = ( ret != 0 ) ? MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED : 0;
+    return( ret );
+}
+
 int mbedtls_mpi_mod_modulus_setup( mbedtls_mpi_mod_modulus *m,
                                    const mbedtls_mpi_uint *p,
                                    size_t p_limbs,
@@ -120,8 +155,9 @@ int mbedtls_mpi_mod_modulus_setup( mbedtls_mpi_mod_modulus *m,
     {
         case MBEDTLS_MPI_MOD_REP_MONTGOMERY:
             m->int_rep = int_rep;
-            m->rep.mont.rr = NULL;
-            m->rep.mont.mm = 0; break;
+            m->rep.mont.mm = mbedtls_mpi_core_montmul_init( m->p );
+            set_mont_const_square( &m->rep.mont.rr, m->p, m->limbs );
+            break;
         case MBEDTLS_MPI_MOD_REP_OPT_RED:
             m->int_rep = int_rep;
             m->rep.ored = NULL;