Add conditional swap and assign function for MPI core

Signed-off-by: Gabor Mezei <gabor.mezei@arm.com>
diff --git a/library/bignum_core.c b/library/bignum_core.c
index c47292e..6aa1e00 100644
--- a/library/bignum_core.c
+++ b/library/bignum_core.c
@@ -25,6 +25,7 @@
 
 #include "mbedtls/error.h"
 #include "mbedtls/platform_util.h"
+#include "constant_time_internal.h"
 
 #if defined(MBEDTLS_PLATFORM_C)
 #include "mbedtls/platform.h"
@@ -161,6 +162,61 @@
     }
 }
 
+int mbedtls_mpi_core_cond_assign( mbedtls_mpi_uint *X,
+                                  size_t X_limbs,
+                                  const mbedtls_mpi_uint *Y,
+                                  size_t Y_limbs,
+                                  unsigned char assign )
+{
+    if( X_limbs < Y_limbs )
+        return( MBEDTLS_ERR_MPI_BUFFER_TOO_SMALL );
+
+    if( X != NULL && Y != NULL )
+    {
+        /* all-bits 1 if assign is 1, all-bits 0 if assign is 0 */
+        mbedtls_mpi_uint limb_mask = mbedtls_ct_mpi_uint_mask( assign );
+
+        mbedtls_ct_mpi_uint_cond_assign( X_limbs, X, Y, assign );
+
+        for( size_t i = Y_limbs; i < X_limbs; i++ )
+            X[i] &= ~limb_mask;
+
+        return( 0 );
+    }
+
+    return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
+}
+
+int mbedtls_mpi_core_cond_swap( mbedtls_mpi_uint *X,
+                                size_t X_limbs,
+                                mbedtls_mpi_uint *Y,
+                                size_t Y_limbs,
+                                unsigned char swap )
+{
+    if( X == Y )
+        return( 0 );
+
+    if( X_limbs != Y_limbs )
+        return( MBEDTLS_ERR_MPI_BUFFER_TOO_SMALL );
+
+    if( X != NULL && Y != NULL )
+    {
+        /* all-bits 1 if swap is 1, all-bits 0 if swap is 0 */
+        mbedtls_mpi_uint limb_mask = mbedtls_ct_mpi_uint_mask( swap );
+
+        for( size_t i = 0; i < X_limbs; i++ )
+        {
+            mbedtls_mpi_uint tmp = X[i];
+            X[i] = ( X[i] & ~limb_mask ) | ( Y[i] & limb_mask );
+            Y[i] = ( Y[i] & ~limb_mask ) | (  tmp & limb_mask );
+        }
+
+        return( 0 );
+    }
+
+    return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
+}
+
 int mbedtls_mpi_core_read_le( mbedtls_mpi_uint *X,
                               size_t X_limbs,
                               const unsigned char *input,
diff --git a/library/bignum_core.h b/library/bignum_core.h
index 8e227f8..24650fe 100644
--- a/library/bignum_core.h
+++ b/library/bignum_core.h
@@ -74,6 +74,70 @@
 void mbedtls_mpi_core_bigendian_to_host( mbedtls_mpi_uint *A,
                                          size_t A_limbs );
 
+/**
+ * \brief   Perform a safe conditional copy of MPI which doesn't reveal whether
+ *          the condition was true or not.
+ *
+ * \param[OUT] X        The address of the first MPI. This must be initialized.
+ * \param      X_limbs  The number of limbs of \p X.
+ * \param[IN]  Y        The address of the second MPI. This must be initialized.
+ * \param      Y_limbs  The number of limbs of \p Y.
+ * \param      assign   The condition deciding whether to perform the
+ *                      assignment or not. Must be either 0 or 1:
+ *                      * \c 1: Perform the assignment `X = Y`.
+ *                      * \c 0: Keep the original value of \p X.
+ *
+ * \note           This function avoids leaking any information about whether
+ *                 the assignment was done or not.
+ *
+ * \warning        If \p assign is neither 0 nor 1, the result of this function
+ *                 is indeterminate, and the resulting value in \p X might be
+ *                 neither its original value nor the value in \p Y.
+ *
+ * \return         \c 0 if successful.
+ * \return         #MBEDTLS_ERR_MPI_BUFFER_TOO_SMALL if \p X isn't
+ *                 large enough to hold the value in \p Y.
+ * \return         #MBEDTLS_ERR_MPI_BAD_INPUT_DATA if \p X or \p Y is invalid.
+ */
+int mbedtls_mpi_core_cond_assign( mbedtls_mpi_uint *X,
+                                  size_t X_limbs,
+                                  const mbedtls_mpi_uint *Y,
+                                  size_t Y_limbs,
+                                  unsigned char assign );
+
+/**
+ * \brief   Perform a safe conditional copy of MPI which doesn't reveal whether
+ *          the condition was true or not.
+ *
+ * \param[IN,OUT] X         The address of the first MPI.
+ *                          This must be initialized.
+ * \param         X_limbs   The number of limbs of \p X.
+ * \param[IN,OUT] Y         The address of the second MPI.
+ *                          This must be initialized.
+ * \param         Y_limbs   The number of limbs of \p Y.
+ * \param         swap      The condition deciding whether to perform
+ *                          the swap or not. Must be either 0 or 1:
+ *                          * \c 1: Swap the values of \p X and \p Y.
+ *                          * \c 0: Keep the original values of \p X and \p Y.
+ *
+ * \note           This function avoids leaking any information about whether
+ *                 the swap was done or not.
+ *
+ * \warning        If \p swap is neither 0 nor 1, the result of this function
+ *                 is indeterminate, and both \p X and \p Y might end up with
+ *                 values different to either of the original ones.
+ *
+ * \return         \c 0 if successful.
+ * \return         #MBEDTLS_ERR_MPI_BUFFER_TOO_SMALL if the size of
+ *                 \p X and \p Y is differ.
+ * \return         #MBEDTLS_ERR_MPI_BAD_INPUT_DATA if \p X or \p Y is invalid.
+ */
+int mbedtls_mpi_core_cond_swap( mbedtls_mpi_uint *X,
+                                size_t X_limbs,
+                                mbedtls_mpi_uint *Y,
+                                size_t Y_limbs,
+                                unsigned char swap );
+
 /** Import X from unsigned binary data, little-endian.
  *
  * The MPI needs to have enough limbs to store the full value (including any