Merge pull request #6777 from tom-cosgrove-arm/issue-6292-mod_inv

	Bignum: Implement high level fixed width modular inversion

15 files changed, 593 insertions, 46 deletions

diff --git a/library/bignum_core.c b/library/bignum_core.c
index 064b158..bda62f2 100644
--- a/library/bignum_core.c
+++ b/library/bignum_core.c
@@ -824,6 +824,40 @@ mbedtls_mpi_uint mbedtls_mpi_core_sub_int( mbedtls_mpi_uint *X,
     return( c );
 }
 
+mbedtls_mpi_uint mbedtls_mpi_core_check_zero_ct( const mbedtls_mpi_uint *A,
+                                                 size_t limbs )
+{
+    mbedtls_mpi_uint bits = 0;
+
+    for( size_t i = 0; i < limbs; i++ )
+        bits |= A[i];
+
+    return( bits );
+}
+
+void mbedtls_mpi_core_to_mont_rep( mbedtls_mpi_uint *X,
+                                   const mbedtls_mpi_uint *A,
+                                   const mbedtls_mpi_uint *N,
+                                   size_t AN_limbs,
+                                   mbedtls_mpi_uint mm,
+                                   const mbedtls_mpi_uint *rr,
+                                   mbedtls_mpi_uint *T )
+{
+    mbedtls_mpi_core_montmul( X, A, rr, AN_limbs, N, AN_limbs, mm, T );
+}
+
+void mbedtls_mpi_core_from_mont_rep( mbedtls_mpi_uint *X,
+                                     const mbedtls_mpi_uint *A,
+                                     const mbedtls_mpi_uint *N,
+                                     size_t AN_limbs,
+                                     mbedtls_mpi_uint mm,
+                                     mbedtls_mpi_uint *T )
+{
+    const mbedtls_mpi_uint Rinv = 1;    /* 1/R in Mont. rep => 1 */
+
+    mbedtls_mpi_core_montmul( X, A, &Rinv, 1, N, AN_limbs, mm, T );
+}
+
 /* END MERGE SLOT 3 */
 
 /* BEGIN MERGE SLOT 4 */
diff --git a/library/bignum_core.h b/library/bignum_core.h
index bfc9725..771f208 100644
--- a/library/bignum_core.h
+++ b/library/bignum_core.h
@@ -555,6 +555,10 @@ int mbedtls_mpi_core_random( mbedtls_mpi_uint *X,
  * \brief          Returns the number of limbs of working memory required for
  *                 a call to `mbedtls_mpi_core_exp_mod()`.
  *
+ * \note           This will always be at least
+ *                 `mbedtls_mpi_core_montmul_working_limbs(AN_limbs)`,
+ *                 i.e. sufficient for a call to `mbedtls_mpi_core_montmul()`.
+ *
  * \param AN_limbs The number of limbs in the input `A` and the modulus `N`
  *                 (they must be the same size) that will be given to
  *                 `mbedtls_mpi_core_exp_mod()`.
@@ -625,6 +629,111 @@ mbedtls_mpi_uint mbedtls_mpi_core_sub_int( mbedtls_mpi_uint *X,
                                            mbedtls_mpi_uint b,
                                            size_t limbs );
 
+/**
+ * \brief Determine if a given MPI has the value \c 0 in constant time with
+ *        respect to the value (but not with respect to the number of limbs).
+ *
+ * \param[in] A   The MPI to test.
+ * \param limbs   Number of limbs in \p A.
+ *
+ * \return        0 if `A == 0`
+ *                non-0 (may be any value) if `A != 0`.
+ */
+mbedtls_mpi_uint mbedtls_mpi_core_check_zero_ct( const mbedtls_mpi_uint *A,
+                                                 size_t limbs );
+
+/**
+ * \brief          Returns the number of limbs of working memory required for
+ *                 a call to `mbedtls_mpi_core_montmul()`.
+ *
+ * \param AN_limbs The number of limbs in the input `A` and the modulus `N`
+ *                 (they must be the same size) that will be given to
+ *                 `mbedtls_mpi_core_montmul()` or one of the other functions
+ *                 that specifies this as the amount of working memory needed.
+ *
+ * \return         The number of limbs of working memory required by
+ *                 `mbedtls_mpi_core_montmul()` (or other similar function).
+ */
+static inline size_t mbedtls_mpi_core_montmul_working_limbs( size_t AN_limbs )
+{
+    return( 2 * AN_limbs + 1 );
+}
+
+/** Convert an MPI into Montgomery form.
+ *
+ * \p X may be aliased to \p A, but may not otherwise overlap it.
+ *
+ * \p X may not alias \p N (it is in canonical form, so must be stricly less
+ * than \p N). Nor may it alias or overlap \p rr (this is unlikely to be
+ * required in practice.)
+ *
+ * This function is a thin wrapper around `mbedtls_mpi_core_montmul()` that is
+ * an alternative to calling `mbedtls_mpi_mod_raw_to_mont_rep()` when we
+ * don't want to allocate memory.
+ *
+ * \param[out]    X         The result of the conversion.
+ *                          Must have the same number of limbs as \p A.
+ * \param[in]     A         The MPI to convert into Montgomery form.
+ *                          Must have the same number of limbs as the modulus.
+ * \param[in]     N         The address of the modulus, which gives the size of
+ *                          the base `R` = 2^(biL*N->limbs).
+ * \param[in]     AN_limbs  The number of limbs in \p X, \p A, \p N and \p rr.
+ * \param         mm        The Montgomery constant for \p N: -N^-1 mod 2^biL.
+ *                          This can be determined by calling
+ *                          `mbedtls_mpi_core_montmul_init()`.
+ * \param[in]     rr        The residue for `2^{2*n*biL} mod N`.
+ * \param[in,out] T         Temporary storage of size at least
+ *                          `mbedtls_mpi_core_montmul_working_limbs(AN_limbs)`
+ *                          limbs.
+ *                          Its initial content is unused and
+ *                          its final content is indeterminate.
+ *                          It must not alias or otherwise overlap any of the
+ *                          other parameters.
+ */
+void mbedtls_mpi_core_to_mont_rep( mbedtls_mpi_uint *X,
+                                   const mbedtls_mpi_uint *A,
+                                   const mbedtls_mpi_uint *N,
+                                   size_t AN_limbs,
+                                   mbedtls_mpi_uint mm,
+                                   const mbedtls_mpi_uint *rr,
+                                   mbedtls_mpi_uint *T );
+
+/** Convert an MPI from Montgomery form.
+ *
+ * \p X may be aliased to \p A, but may not otherwise overlap it.
+ *
+ * \p X may not alias \p N (it is in canonical form, so must be stricly less
+ * than \p N).
+ *
+ * This function is a thin wrapper around `mbedtls_mpi_core_montmul()` that is
+ * an alternative to calling `mbedtls_mpi_mod_raw_from_mont_rep()` when we
+ * don't want to allocate memory.
+ *
+ * \param[out]    X         The result of the conversion.
+ *                          Must have the same number of limbs as \p A.
+ * \param[in]     A         The MPI to convert from Montgomery form.
+ *                          Must have the same number of limbs as the modulus.
+ * \param[in]     N         The address of the modulus, which gives the size of
+ *                          the base `R` = 2^(biL*N->limbs).
+ * \param[in]     AN_limbs  The number of limbs in \p X, \p A and \p N.
+ * \param         mm        The Montgomery constant for \p N: -N^-1 mod 2^biL.
+ *                          This can be determined by calling
+ *                          `mbedtls_mpi_core_montmul_init()`.
+ * \param[in,out] T         Temporary storage of size at least
+ *                          `mbedtls_mpi_core_montmul_working_limbs(AN_limbs)`
+ *                          limbs.
+ *                          Its initial content is unused and
+ *                          its final content is indeterminate.
+ *                          It must not alias or otherwise overlap any of the
+ *                          other parameters.
+ */
+void mbedtls_mpi_core_from_mont_rep( mbedtls_mpi_uint *X,
+                                     const mbedtls_mpi_uint *A,
+                                     const mbedtls_mpi_uint *N,
+                                     size_t AN_limbs,
+                                     mbedtls_mpi_uint mm,
+                                     mbedtls_mpi_uint *T );
+
 /* END MERGE SLOT 3 */
 
 /* BEGIN MERGE SLOT 4 */
diff --git a/library/bignum_mod.c b/library/bignum_mod.c
index 0057eba..31e18e7 100644
--- a/library/bignum_mod.c
+++ b/library/bignum_mod.c
@@ -79,7 +79,7 @@ void mbedtls_mpi_mod_modulus_free( mbedtls_mpi_mod_modulus *m )
             if (m->rep.mont.rr != NULL)
             {
                 mbedtls_platform_zeroize( (mbedtls_mpi_uint *) m->rep.mont.rr,
-                                           m->limbs );
+                                           m->limbs * sizeof(mbedtls_mpi_uint) );
                 mbedtls_free( (mbedtls_mpi_uint *)m->rep.mont.rr );
                 m->rep.mont.rr = NULL;
             }
@@ -191,6 +191,95 @@ int mbedtls_mpi_mod_sub( mbedtls_mpi_mod_residue *X,
 
     return( 0 );
 }
+
+static int mbedtls_mpi_mod_inv_mont( mbedtls_mpi_mod_residue *X,
+                                     const mbedtls_mpi_mod_residue *A,
+                                     const mbedtls_mpi_mod_modulus *N,
+                                     mbedtls_mpi_uint *working_memory )
+{
+    /* Input already in Montgomery form, so there's little to do */
+    mbedtls_mpi_mod_raw_inv_prime( X->p, A->p,
+                                   N->p, N->limbs,
+                                   N->rep.mont.rr,
+                                   working_memory );
+    return( 0 );
+}
+
+static int mbedtls_mpi_mod_inv_non_mont( mbedtls_mpi_mod_residue *X,
+                                         const mbedtls_mpi_mod_residue *A,
+                                         const mbedtls_mpi_mod_modulus *N,
+                                         mbedtls_mpi_uint *working_memory )
+{
+    /* Need to convert input into Montgomery form */
+
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    mbedtls_mpi_mod_modulus Nmont;
+    mbedtls_mpi_mod_modulus_init( &Nmont );
+
+    MBEDTLS_MPI_CHK( mbedtls_mpi_mod_modulus_setup( &Nmont, N->p, N->limbs,
+                                         MBEDTLS_MPI_MOD_REP_MONTGOMERY ) );
+
+    /* We'll use X->p to hold the Montgomery form of the input A->p */
+    mbedtls_mpi_core_to_mont_rep( X->p, A->p, Nmont.p, Nmont.limbs,
+                                  Nmont.rep.mont.mm, Nmont.rep.mont.rr,
+                                  working_memory );
+
+    mbedtls_mpi_mod_raw_inv_prime( X->p, X->p,
+                                   Nmont.p, Nmont.limbs,
+                                   Nmont.rep.mont.rr,
+                                   working_memory );
+
+    /* And convert back from Montgomery form */
+
+    mbedtls_mpi_core_from_mont_rep( X->p, X->p, Nmont.p, Nmont.limbs,
+                                    Nmont.rep.mont.mm, working_memory );
+
+cleanup:
+    mbedtls_mpi_mod_modulus_free( &Nmont );
+    return( ret );
+}
+
+int mbedtls_mpi_mod_inv( mbedtls_mpi_mod_residue *X,
+                         const mbedtls_mpi_mod_residue *A,
+                         const mbedtls_mpi_mod_modulus *N )
+{
+    if( X->limbs != N->limbs || A->limbs != N->limbs )
+        return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
+
+    /* Zero has the same value regardless of Montgomery form or not */
+    if( mbedtls_mpi_core_check_zero_ct( A->p, A->limbs ) == 0 )
+        return( MBEDTLS_ERR_MPI_BAD_INPUT_DATA );
+
+    size_t working_limbs =
+                    mbedtls_mpi_mod_raw_inv_prime_working_limbs( N->limbs );
+
+    mbedtls_mpi_uint *working_memory = mbedtls_calloc( working_limbs,
+                                                     sizeof(mbedtls_mpi_uint) );
+    if( working_memory == NULL )
+        return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
+
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    switch( N->int_rep )
+    {
+        case MBEDTLS_MPI_MOD_REP_MONTGOMERY:
+            ret = mbedtls_mpi_mod_inv_mont( X, A, N, working_memory );
+            break;
+        case MBEDTLS_MPI_MOD_REP_OPT_RED:
+            ret = mbedtls_mpi_mod_inv_non_mont( X, A, N, working_memory );
+            break;
+        default:
+            ret = MBEDTLS_ERR_MPI_BAD_INPUT_DATA;
+            break;
+    }
+
+    mbedtls_platform_zeroize( working_memory,
+                              working_limbs * sizeof(mbedtls_mpi_uint) );
+    mbedtls_free( working_memory );
+
+    return ret;
+}
 /* END MERGE SLOT 3 */
 
 /* BEGIN MERGE SLOT 4 */
diff --git a/library/bignum_mod.h b/library/bignum_mod.h
index f089f65..95aaacc 100644
--- a/library/bignum_mod.h
+++ b/library/bignum_mod.h
@@ -249,6 +249,35 @@ int mbedtls_mpi_mod_sub( mbedtls_mpi_mod_residue *X,
                          const mbedtls_mpi_mod_residue *A,
                          const mbedtls_mpi_mod_residue *B,
                          const mbedtls_mpi_mod_modulus *N );
+
+/**
+ * \brief Perform modular inversion of an MPI with respect to a modulus \p N.
+ *
+ * \p A and \p X must be associated with the modulus \p N and will therefore
+ * have the same number of limbs as \p N.
+ *
+ * \p X may be aliased to \p A.
+ *
+ * \warning  Currently only supports prime moduli, but does not check for them.
+ *
+ * \param[out] X   The modular inverse of \p A with respect to \p N.
+ * \param[in] A    The number to calculate the modular inverse of.
+ *                 Must not be 0.
+ * \param[in] N    The modulus to use.
+ *
+ * \return         \c 0 if successful.
+ * \return         #MBEDTLS_ERR_MPI_BAD_INPUT_DATA if \p A and \p N do not
+ *                 have the same number of limbs.
+ * \return         #MBEDTLS_ERR_MPI_BAD_INPUT_DATA if \p A is zero.
+ * \return         #MBEDTLS_ERR_MPI_ALLOC_FAILED if couldn't allocate enough
+ *                 memory (needed for conversion to and from Mongtomery form
+ *                 when not in Montgomery form already, and for temporary use
+ *                 by the inversion calculation itself).
+ */
+
+int mbedtls_mpi_mod_inv( mbedtls_mpi_mod_residue *X,
+                         const mbedtls_mpi_mod_residue *A,
+                         const mbedtls_mpi_mod_modulus *N );
 /* END MERGE SLOT 3 */
 
 /* BEGIN MERGE SLOT 4 */
diff --git a/library/bignum_mod_raw.c b/library/bignum_mod_raw.c
index c0877bd..5950ff6 100644
--- a/library/bignum_mod_raw.c
+++ b/library/bignum_mod_raw.c
@@ -193,13 +193,13 @@ int mbedtls_mpi_mod_raw_to_mont_rep( mbedtls_mpi_uint *X,
                                      const mbedtls_mpi_mod_modulus *m )
 {
     mbedtls_mpi_uint *T;
-    const size_t t_limbs = m->limbs * 2 + 1;
+    const size_t t_limbs = mbedtls_mpi_core_montmul_working_limbs( m->limbs );
 
     if( ( T = (mbedtls_mpi_uint *) mbedtls_calloc( t_limbs, ciL ) ) == NULL )
         return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
 
-    mbedtls_mpi_core_montmul( X, X, m->rep.mont.rr, m->limbs, m->p, m->limbs,
-                              m->rep.mont.mm, T );
+    mbedtls_mpi_core_to_mont_rep( X, X, m->p, m->limbs,
+                                  m->rep.mont.mm, m->rep.mont.rr, T );
 
     mbedtls_platform_zeroize( T, t_limbs * ciL );
     mbedtls_free( T );
@@ -209,15 +209,13 @@ int mbedtls_mpi_mod_raw_to_mont_rep( mbedtls_mpi_uint *X,
 int mbedtls_mpi_mod_raw_from_mont_rep( mbedtls_mpi_uint *X,
                                        const mbedtls_mpi_mod_modulus *m )
 {
-    const mbedtls_mpi_uint one = 1;
-    const size_t t_limbs = m->limbs * 2 + 1;
+    const size_t t_limbs = mbedtls_mpi_core_montmul_working_limbs( m->limbs );
     mbedtls_mpi_uint *T;
 
     if( ( T = (mbedtls_mpi_uint *) mbedtls_calloc( t_limbs, ciL ) ) == NULL )
         return( MBEDTLS_ERR_MPI_ALLOC_FAILED );
 
-    mbedtls_mpi_core_montmul( X, X, &one, 1, m->p, m->limbs,
-                              m->rep.mont.mm, T );
+    mbedtls_mpi_core_from_mont_rep( X, X, m->p, m->limbs, m->rep.mont.mm, T );
 
     mbedtls_platform_zeroize( T, t_limbs * ciL );
     mbedtls_free( T );
diff --git a/library/bignum_mod_raw.h b/library/bignum_mod_raw.h
index 380f30b..0fac6f8 100644
--- a/library/bignum_mod_raw.h
+++ b/library/bignum_mod_raw.h
@@ -258,6 +258,10 @@ void mbedtls_mpi_mod_raw_mul( mbedtls_mpi_uint *X,
  * \brief          Returns the number of limbs of working memory required for
  *                 a call to `mbedtls_mpi_mod_raw_inv_prime()`.
  *
+ * \note           This will always be at least
+ *                 `mbedtls_mpi_core_montmul_working_limbs(AN_limbs)`,
+ *                 i.e. sufficient for a call to `mbedtls_mpi_core_montmul()`.
+ *
  * \param AN_limbs The number of limbs in the input `A` and the modulus `N`
  *                 (they must be the same size) that will be given to
  *                 `mbedtls_mpi_mod_raw_inv_prime()`.
diff --git a/scripts/mbedtls_dev/bignum_common.py b/scripts/mbedtls_dev/bignum_common.py
index 0339b1a..c4efabf 100644
--- a/scripts/mbedtls_dev/bignum_common.py
+++ b/scripts/mbedtls_dev/bignum_common.py
@@ -39,6 +39,11 @@ def invmod(a: int, n: int) -> int:
         return b
     raise ValueError("Not invertible")
 
+def invmod_positive(a: int, n: int) -> int:
+    """Return a non-negative inverse of a to modulo n."""
+    inv = invmod(a, n)
+    return inv if inv >= 0 else inv + n
+
 def hex_to_int(val: str) -> int:
     """Implement the syntax accepted by mbedtls_test_read_mpi().
 
@@ -244,6 +249,8 @@ class ModOperationCommon(OperationCommon):
     #pylint: disable=abstract-method
     """Target for bignum mod_raw test case generation."""
     moduli = MODULI_DEFAULT # type: List[str]
+    montgomery_form_a = False
+    disallow_zero_a = False
 
     def __init__(self, val_n: str, val_a: str, val_b: str = "0",
                  bits_in_limb: int = 64) -> None:
@@ -264,6 +271,14 @@ class ModOperationCommon(OperationCommon):
         return self.int_n
 
     @property
+    def arg_a(self) -> str:
+        if self.montgomery_form_a:
+            value_a = self.to_montgomery(self.int_a)
+        else:
+            value_a = self.int_a
+        return self.format_arg('{:x}'.format(value_a))
+
+    @property
     def arg_n(self) -> str:
         return self.format_arg(self.val_n)
 
@@ -287,6 +302,8 @@ class ModOperationCommon(OperationCommon):
     def is_valid(self) -> bool:
         if self.int_a >= self.int_n:
             return False
+        if self.disallow_zero_a and self.int_a == 0:
+            return False
         if self.arity == 2 and self.int_b >= self.int_n:
             return False
         return True
diff --git a/scripts/mbedtls_dev/bignum_core.py b/scripts/mbedtls_dev/bignum_core.py
index 158ada9..24d37cb 100644
--- a/scripts/mbedtls_dev/bignum_core.py
+++ b/scripts/mbedtls_dev/bignum_core.py
@@ -757,15 +757,7 @@ class BignumCoreExpMod(BignumCoreTarget, bignum_common.ModOperationCommon):
     test_function = "mpi_core_exp_mod"
     test_name = "Core modular exponentiation (Mongtomery form only)"
     input_style = "fixed"
-
-    def arguments(self) -> List[str]:
-        # Input 'a' has to be given in Montgomery form
-        mont_a = self.to_montgomery(self.int_a)
-        arg_mont_a = self.format_arg('{:x}'.format(mont_a))
-        return [bignum_common.quote_str(n) for n in [self.arg_n,
-                                                     arg_mont_a,
-                                                     self.arg_b]
-               ] + self.result()
+    montgomery_form_a = True
 
     def result(self) -> List[str]:
         # Result has to be given in Montgomery form too
@@ -818,6 +810,20 @@ class BignumCoreSubInt(BignumCoreTarget, bignum_common.OperationCommon):
             str(-borrow)
         ]
 
+class BignumCoreZeroCheckCT(BignumCoreTarget, bignum_common.OperationCommon):
+    """Test cases for bignum core zero check (constant flow)."""
+    count = 0
+    symbol = "== 0"
+    test_function = "mpi_core_check_zero_ct"
+    test_name = "mpi_core_check_zero_ct"
+    input_style = "variable"
+    arity = 1
+    suffix = True
+
+    def result(self) -> List[str]:
+        result = 1 if self.int_a == 0 else 0
+        return [str(result)]
+
 # END MERGE SLOT 3
 
 # BEGIN MERGE SLOT 4
diff --git a/scripts/mbedtls_dev/bignum_data.py b/scripts/mbedtls_dev/bignum_data.py
index 9658933..0a48e53 100644
--- a/scripts/mbedtls_dev/bignum_data.py
+++ b/scripts/mbedtls_dev/bignum_data.py
@@ -121,6 +121,9 @@ MODULI_DEFAULT = [
 ONLY_PRIME_MODULI = [
         "53", # safe prime
         "8ac72304057392b5",     # 9999999997777777333 (longer, not safe, prime)
+        # The next prime has a different R in Montgomery form depending on
+        # whether 32- or 64-bit MPIs are used.
+        "152d02c7e14af67fe0bf", # 99999999999999999991999
         SAFE_PRIME_192_BIT_SEED_1,  # safe prime
         SAFE_PRIME_1024_BIT_SEED_3, # safe prime
         ]
diff --git a/scripts/mbedtls_dev/bignum_mod.py b/scripts/mbedtls_dev/bignum_mod.py
index a16699a..25afe30 100644
--- a/scripts/mbedtls_dev/bignum_mod.py
+++ b/scripts/mbedtls_dev/bignum_mod.py
@@ -18,6 +18,7 @@ from typing import Dict, List
 
 from . import test_data_generation
 from . import bignum_common
+from .bignum_data import ONLY_PRIME_MODULI
 
 class BignumModTarget(test_data_generation.BaseTarget):
     #pylint: disable=abstract-method, too-few-public-methods
@@ -48,6 +49,42 @@ class BignumModSub(bignum_common.ModOperationCommon, BignumModTarget):
         # generated cases
         return [self.format_result(result), "0"]
 
+class BignumModInvNonMont(bignum_common.ModOperationCommon, BignumModTarget):
+    """Test cases for bignum mpi_mod_inv() - not in Montgomery form."""
+    moduli = ONLY_PRIME_MODULI  # for now only prime moduli supported
+    symbol = "^ -1"
+    test_function = "mpi_mod_inv_non_mont"
+    test_name = "mbedtls_mpi_mod_inv non-Mont. form"
+    input_style = "fixed"
+    arity = 1
+    suffix = True
+    disallow_zero_a = True
+
+    def result(self) -> List[str]:
+        result = bignum_common.invmod_positive(self.int_a, self.int_n)
+        # To make negative tests easier, append 0 for success to the
+        # generated cases
+        return [self.format_result(result), "0"]
+
+class BignumModInvMont(bignum_common.ModOperationCommon, BignumModTarget):
+    """Test cases for bignum mpi_mod_inv() - Montgomery form."""
+    moduli = ONLY_PRIME_MODULI  # for now only prime moduli supported
+    symbol = "^ -1"
+    test_function = "mpi_mod_inv_mont"
+    test_name = "mbedtls_mpi_mod_inv Mont. form"
+    input_style = "arch_split"  # Mont. form requires arch_split
+    arity = 1
+    suffix = True
+    disallow_zero_a = True
+    montgomery_form_a = True
+
+    def result(self) -> List[str]:
+        result = bignum_common.invmod_positive(self.int_a, self.int_n)
+        mont_result = self.to_montgomery(result)
+        # To make negative tests easier, append 0 for success to the
+        # generated cases
+        return [self.format_result(mont_result), "0"]
+
 # END MERGE SLOT 3
 
 # BEGIN MERGE SLOT 4
diff --git a/scripts/mbedtls_dev/bignum_mod_raw.py b/scripts/mbedtls_dev/bignum_mod_raw.py
index 296e2d2..09bbbee 100644
--- a/scripts/mbedtls_dev/bignum_mod_raw.py
+++ b/scripts/mbedtls_dev/bignum_mod_raw.py
@@ -80,24 +80,14 @@ class BignumModRawInvPrime(bignum_common.ModOperationCommon,
     symbol = "^ -1"
     test_function = "mpi_mod_raw_inv_prime"
     test_name = "mbedtls_mpi_mod_raw_inv_prime (Montgomery form only)"
-    input_style = "fixed"
+    input_style = "arch_split"
     arity = 1
     suffix = True
-
-    @property
-    def is_valid(self) -> bool:
-        return self.int_a > 0 and self.int_a < self.int_n
-
-    @property
-    def arg_a(self) -> str:
-        # Input has to be given in Montgomery form
-        mont_a = self.to_montgomery(self.int_a)
-        return self.format_arg('{:x}'.format(mont_a))
+    montgomery_form_a = True
+    disallow_zero_a = True
 
     def result(self) -> List[str]:
-        result = bignum_common.invmod(self.int_a, self.int_n)
-        if result < 0:
-            result += self.int_n
+        result = bignum_common.invmod_positive(self.int_a, self.int_n)
         mont_result = self.to_montgomery(result)
         return [self.format_result(mont_result)]
 
diff --git a/tests/suites/test_suite_bignum_core.function b/tests/suites/test_suite_bignum_core.function
index 9cb314b..47f9013 100644
--- a/tests/suites/test_suite_bignum_core.function
+++ b/tests/suites/test_suite_bignum_core.function
@@ -848,7 +848,9 @@ void mpi_core_montmul( int limbs_AN4, int limbs_B4,
     TEST_EQUAL( 0, mbedtls_mpi_grow( X, limbs_AN ) );
     TEST_EQUAL( 0, mbedtls_mpi_grow( &B, limbs_B ) );
 
-    TEST_EQUAL( 0, mbedtls_mpi_grow( &T, limbs_AN * 2 + 1 ) );
+    size_t working_limbs = mbedtls_mpi_core_montmul_working_limbs( limbs_AN );
+    TEST_EQUAL( working_limbs, limbs_AN * 2 + 1 );
+    TEST_EQUAL( 0, mbedtls_mpi_grow( &T, working_limbs ) );
 
     /* Calculate the Montgomery constant (this is unit tested separately) */
     mbedtls_mpi_uint mm = mbedtls_mpi_core_montmul_init( N.p );
@@ -1133,6 +1135,10 @@ void mpi_core_exp_mod( char * input_N, char * input_A,
     TEST_LE_U( min_expected_working_limbs, working_limbs );
     TEST_LE_U( working_limbs, max_expected_working_limbs );
 
+    /* Should also be at least mbedtls_mpi_core_montmul_working_limbs() */
+    TEST_LE_U( mbedtls_mpi_core_montmul_working_limbs( N_limbs ),
+               working_limbs );
+
     ASSERT_ALLOC( T, working_limbs );
 
     mbedtls_mpi_core_exp_mod( Y, A, N, N_limbs, E, E_limbs, R2, T );
@@ -1212,6 +1218,25 @@ exit:
 }
 /* END_CASE */
 
+/* BEGIN_CASE */
+void mpi_core_check_zero_ct( char *input_X, int expected_is_zero )
+{
+    mbedtls_mpi_uint *X = NULL;
+    size_t X_limbs;
+
+    TEST_EQUAL( 0, mbedtls_test_read_mpi_core( &X, &X_limbs, input_X ) );
+
+    TEST_CF_SECRET( X, X_limbs * sizeof( mbedtls_mpi_uint ) );
+
+    mbedtls_mpi_uint check = mbedtls_mpi_core_check_zero_ct( X, X_limbs );
+    int is_zero = (check == 0);
+    TEST_EQUAL( is_zero, expected_is_zero );
+
+exit:
+    mbedtls_free( X );
+}
+/* END_CASE */
+
 /* END MERGE SLOT 3 */
 
 /* BEGIN MERGE SLOT 4 */
diff --git a/tests/suites/test_suite_bignum_mod.function b/tests/suites/test_suite_bignum_mod.function
index 1e64255..79f5134 100644
--- a/tests/suites/test_suite_bignum_mod.function
+++ b/tests/suites/test_suite_bignum_mod.function
@@ -208,6 +208,106 @@ exit:
     mbedtls_free( X_raw );
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void mpi_mod_inv_mont( char * input_N,
+                       char * input_A, char * input_I,
+                       int expected_ret )
+{
+    mbedtls_mpi_mod_residue a = { NULL, 0 };    /* argument */
+    mbedtls_mpi_mod_residue i = { NULL, 0 };    /* expected inverse wrt N */
+    mbedtls_mpi_mod_residue x = { NULL, 0 };    /* output */
+    mbedtls_mpi_uint *X_raw = NULL;
+
+    mbedtls_mpi_mod_modulus N;
+    mbedtls_mpi_mod_modulus_init( &N );
+
+    TEST_EQUAL( 0,
+        test_read_modulus( &N, MBEDTLS_MPI_MOD_REP_MONTGOMERY, input_N ) );
+
+    /* test_read_residue() normally checks that inputs have the same number of
+     * limbs as the modulus. For negative testing we can ask it to skip this
+     * with a non-zero final parameter. */
+    TEST_EQUAL( 0, test_read_residue( &a, &N, input_A, expected_ret != 0 ) );
+    TEST_EQUAL( 0, test_read_residue( &i, &N, input_I, expected_ret != 0 ) );
+
+    size_t limbs = N.limbs;
+    size_t bytes = limbs * sizeof( *X_raw );
+
+    ASSERT_ALLOC( X_raw, limbs );
+
+    TEST_EQUAL( 0, mbedtls_mpi_mod_residue_setup( &x, &N, X_raw, limbs ) );
+
+    TEST_EQUAL( expected_ret, mbedtls_mpi_mod_inv( &x, &a, &N ) );
+    if( expected_ret == 0 )
+    {
+        TEST_COMPARE_MPI_RESIDUES( x, i );
+
+        /* a^-1: alias x to a => Correct result */
+        memcpy( x.p, a.p, bytes );
+        TEST_EQUAL( 0, mbedtls_mpi_mod_inv( &x, &x, &N ) );
+        TEST_COMPARE_MPI_RESIDUES( x, i );
+    }
+
+exit:
+    mbedtls_free( (void *)N.p ); /* mbedtls_mpi_mod_modulus_free() sets N.p = NULL */
+    mbedtls_mpi_mod_modulus_free( &N );
+
+    mbedtls_free( a.p );
+    mbedtls_free( i.p );
+    mbedtls_free( X_raw );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void mpi_mod_inv_non_mont( char * input_N,
+                           char * input_A, char * input_I,
+                           int expected_ret )
+{
+    mbedtls_mpi_mod_residue a = { NULL, 0 };    /* argument */
+    mbedtls_mpi_mod_residue i = { NULL, 0 };    /* expected inverse wrt N */
+    mbedtls_mpi_mod_residue x = { NULL, 0 };    /* output */
+    mbedtls_mpi_uint *X_raw = NULL;
+
+    mbedtls_mpi_mod_modulus N;
+    mbedtls_mpi_mod_modulus_init( &N );
+
+    TEST_EQUAL( 0,
+        test_read_modulus( &N, MBEDTLS_MPI_MOD_REP_OPT_RED, input_N ) );
+
+    /* test_read_residue() normally checks that inputs have the same number of
+     * limbs as the modulus. For negative testing we can ask it to skip this
+     * with a non-zero final parameter. */
+    TEST_EQUAL( 0, test_read_residue( &a, &N, input_A, expected_ret != 0 ) );
+    TEST_EQUAL( 0, test_read_residue( &i, &N, input_I, expected_ret != 0 ) );
+
+    size_t limbs = N.limbs;
+    size_t bytes = limbs * sizeof( *X_raw );
+
+    ASSERT_ALLOC( X_raw, limbs );
+
+    TEST_EQUAL( 0, mbedtls_mpi_mod_residue_setup( &x, &N, X_raw, limbs ) );
+
+    TEST_EQUAL( expected_ret, mbedtls_mpi_mod_inv( &x, &a, &N ) );
+    if( expected_ret == 0 )
+    {
+        TEST_COMPARE_MPI_RESIDUES( x, i );
+
+        /* a^-1: alias x to a => Correct result */
+        memcpy( x.p, a.p, bytes );
+        TEST_EQUAL( 0, mbedtls_mpi_mod_inv( &x, &x, &N ) );
+        TEST_COMPARE_MPI_RESIDUES( x, i );
+    }
+
+exit:
+    mbedtls_free( (void *)N.p ); /* mbedtls_mpi_mod_modulus_free() sets N.p = NULL */
+    mbedtls_mpi_mod_modulus_free( &N );
+
+    mbedtls_free( a.p );
+    mbedtls_free( i.p );
+    mbedtls_free( X_raw );
+}
+/* END_CASE */
 /* END MERGE SLOT 3 */
 
 /* BEGIN MERGE SLOT 4 */
diff --git a/tests/suites/test_suite_bignum_mod.misc.data b/tests/suites/test_suite_bignum_mod.misc.data
index 7b1c85f..6240e21 100644
--- a/tests/suites/test_suite_bignum_mod.misc.data
+++ b/tests/suites/test_suite_bignum_mod.misc.data
@@ -38,6 +38,50 @@ mpi_mod_sub:"014320a022ccb75bdf470ddf25":"a99c71c7":"00033b2e3c9fd0803ce8000f93"
 mpi_mod_sub with second input too short
 mpi_mod_sub:"014320a022ccb75bdf470ddf25":"000000025a55a46e5da99c71c7":"e8000f93":"00":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
 
+mbedtls_mpi_mod_inv non-Mont. form - base case for negative testing (N, A, A^-1)
+mpi_mod_inv_non_mont:"000000000000152d02c7e14af67fe0bf":"00000000000000000000000000000038":"000000000000097418193b6f2e0b5fc3":0
+
+mbedtls_mpi_mod_inv non-Mont. form - A == 0
+mpi_mod_inv_non_mont:"000000000000152d02c7e14af67fe0bf":"00000000000000000000000000000000":"00":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv non-Mont. form - A too long
+mpi_mod_inv_non_mont:"000000000000152d02c7e14af67fe0bf":"0000000000000000000000000000000000000038":"00":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv non-Mont. form - A too short
+mpi_mod_inv_non_mont:"000000000000152d02c7e14af67fe0bf":"0000000000000038":"000000000000097418193b6f2e0b5fc3":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv 32-bit Mont. form - base case for negative testing, A = 1 (N, mont(A), mont(A^-1))
+depends_on:MBEDTLS_HAVE_INT32
+mpi_mod_inv_mont:"000000000000152d02c7e14af67fe0bf":"0000000000000d71d51539b9c02b7b28":"0000000000000d71d51539b9c02b7b28":0
+
+mbedtls_mpi_mod_inv 32-bit Mont. form - A == 0
+depends_on:MBEDTLS_HAVE_INT32
+mpi_mod_inv_mont:"000000000000152d02c7e14af67fe0bf":"00000000000000000000000000000000":"00":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv 32-bit Mont. form - A too long
+depends_on:MBEDTLS_HAVE_INT32
+mpi_mod_inv_mont:"000000000000152d02c7e14af67fe0bf":"000000000000000000000d71d51539b9c02b7b28":"00":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv 32-bit Mont. form - A too short
+depends_on:MBEDTLS_HAVE_INT32
+mpi_mod_inv_mont:"000000000000152d02c7e14af67fe0bf":"00000d71d51539b9c02b7b28":"0000000000000d71d51539b9c02b7b28":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv 64-bit Mont. form - base case for negative testing, A = 1 (N, mont(A), mont(A^-1))
+depends_on:MBEDTLS_HAVE_INT64
+mpi_mod_inv_mont:"0000000000000000000000000000152d02c7e14af67fe0bf":"000000000000000000000000000009545642424381c611fb":"000000000000000000000000000009545642424381c611fb":0
+
+mbedtls_mpi_mod_inv 64-bit Mont. form - A == 0
+depends_on:MBEDTLS_HAVE_INT64
+mpi_mod_inv_mont:"0000000000000000000000000000152d02c7e14af67fe0bf":"000000000000000000000000000000000000000000000000":"000000000000000000000000000009545642424381c611fb":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv 64-bit Mont. form - A too long
+depends_on:MBEDTLS_HAVE_INT64
+mpi_mod_inv_mont:"0000000000000000000000000000152d02c7e14af67fe0bf":"0000000000000000000000000000000000000000000009545642424381c611fb":"000000000000000000000000000009545642424381c611fb":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
+mbedtls_mpi_mod_inv 64-bit Mont. form - A too short
+depends_on:MBEDTLS_HAVE_INT64
+mpi_mod_inv_mont:"0000000000000000000000000000152d02c7e14af67fe0bf":"00000000000009545642424381c611fb":"000000000000000000000000000009545642424381c611fb":MBEDTLS_ERR_MPI_BAD_INPUT_DATA
+
 # END MERGE SLOT 3
 
 # BEGIN MERGE SLOT 4
diff --git a/tests/suites/test_suite_bignum_mod_raw.function b/tests/suites/test_suite_bignum_mod_raw.function
index b1d1f77..461a18e 100644
--- a/tests/suites/test_suite_bignum_mod_raw.function
+++ b/tests/suites/test_suite_bignum_mod_raw.function
@@ -489,6 +489,10 @@ void mpi_mod_raw_inv_prime( char * input_N, char * input_A, char * input_X )
     TEST_LE_U( min_expected_working_limbs, working_limbs );
     TEST_LE_U( working_limbs, max_expected_working_limbs );
 
+    /* Should also be at least mbedtls_mpi_core_montmul_working_limbs() */
+    TEST_LE_U( mbedtls_mpi_core_montmul_working_limbs( N_limbs ),
+               working_limbs );
+
     ASSERT_ALLOC( T, working_limbs );
 
     mbedtls_mpi_mod_raw_inv_prime( Y, A, N, N_limbs, R2, T );
@@ -624,8 +628,10 @@ void mpi_mod_raw_to_mont_rep( char * input_N, char * input_A, char * input_X )
 {
     mbedtls_mpi_uint *N = NULL;
     mbedtls_mpi_uint *A = NULL;
+    mbedtls_mpi_uint *R = NULL; /* for result of low-level conversion */
     mbedtls_mpi_uint *X = NULL;
-    size_t n_limbs, a_limbs, x_limbs, x_bytes;
+    mbedtls_mpi_uint *T = NULL;
+    size_t n_limbs, a_limbs, x_limbs;
 
     mbedtls_mpi_mod_modulus m;
     mbedtls_mpi_mod_modulus_init( &m );
@@ -634,23 +640,50 @@ void mpi_mod_raw_to_mont_rep( char * input_N, char * input_A, char * input_X )
     TEST_EQUAL( 0, mbedtls_test_read_mpi_core( &N, &n_limbs, input_N ) );
     TEST_EQUAL( 0, mbedtls_test_read_mpi_core( &A, &a_limbs, input_A ) );
     TEST_EQUAL( 0, mbedtls_test_read_mpi_core( &X, &x_limbs, input_X ) );
-    x_bytes = x_limbs * sizeof(mbedtls_mpi_uint);
 
-    /* Test that input does not require more limbs than modulo */
-    TEST_LE_U(a_limbs, n_limbs);
+    /* Number to convert must have same number of limbs as modulus */
+    TEST_EQUAL(a_limbs, n_limbs);
+
+    /* Higher-level conversion is in-place, so expected result must have the
+     * same number of limbs too */
+    TEST_EQUAL(x_limbs, n_limbs);
+
+    size_t limbs = n_limbs;
+    size_t bytes = limbs * sizeof(mbedtls_mpi_uint);
 
     TEST_EQUAL( 0, mbedtls_mpi_mod_modulus_setup( &m, N, n_limbs,
-                MBEDTLS_MPI_MOD_REP_MONTGOMERY ) );
+                                          MBEDTLS_MPI_MOD_REP_MONTGOMERY ) );
+
+    /* 1. Test low-level function first */
+
+    /* It has separate output, and requires temporary working storage */
+    size_t temp_limbs = mbedtls_mpi_core_montmul_working_limbs( limbs );
+    ASSERT_ALLOC( T, temp_limbs );
+    ASSERT_ALLOC( R, limbs );
+    mbedtls_mpi_core_to_mont_rep( R, A, N, n_limbs,
+                                  m.rep.mont.mm, m.rep.mont.rr, T );
+    /* Test that the low-level function gives the required value */
+    ASSERT_COMPARE( R, bytes, X, bytes );
+
+    /* Test when output is aliased to input */
+    memcpy( R, A, bytes );
+    mbedtls_mpi_core_to_mont_rep( R, R, N, n_limbs,
+                                  m.rep.mont.mm, m.rep.mont.rr, T );
+    ASSERT_COMPARE( R, bytes, X, bytes );
+
+    /* 2. Test higher-level cannonical to Montgomery conversion */
 
-    /* Convert from cannonical into Montgomery representation */
     TEST_EQUAL(0, mbedtls_mpi_mod_raw_to_mont_rep( A, &m ) );
 
     /* The result matches expected value */
-    ASSERT_COMPARE( A, x_bytes, X, x_bytes );
+    ASSERT_COMPARE( A, bytes, X, bytes );
+
 exit:
     mbedtls_mpi_mod_modulus_free( &m );
+    mbedtls_free( T );
     mbedtls_free( N );
     mbedtls_free( A );
+    mbedtls_free( R );
     mbedtls_free( X );
 }
 /* END_CASE */
@@ -660,8 +693,10 @@ void mpi_mod_raw_from_mont_rep( char * input_N, char * input_A, char * input_X )
 {
     mbedtls_mpi_uint *N = NULL;
     mbedtls_mpi_uint *A = NULL;
+    mbedtls_mpi_uint *R = NULL; /* for result of low-level conversion */
     mbedtls_mpi_uint *X = NULL;
-    size_t n_limbs, a_limbs, x_limbs, x_bytes;
+    mbedtls_mpi_uint *T = NULL;
+    size_t n_limbs, a_limbs, x_limbs;
 
     mbedtls_mpi_mod_modulus m;
     mbedtls_mpi_mod_modulus_init( &m );
@@ -670,23 +705,50 @@ void mpi_mod_raw_from_mont_rep( char * input_N, char * input_A, char * input_X )
     TEST_EQUAL( 0, mbedtls_test_read_mpi_core( &N, &n_limbs, input_N ) );
     TEST_EQUAL( 0, mbedtls_test_read_mpi_core( &A, &a_limbs, input_A ) );
     TEST_EQUAL( 0, mbedtls_test_read_mpi_core( &X, &x_limbs, input_X ) );
-    x_bytes = x_limbs * sizeof(mbedtls_mpi_uint);
 
-    /* Test that input does not require more limbs than modulo */
-    TEST_LE_U(a_limbs, n_limbs);
+    /* Number to convert must have same number of limbs as modulus */
+    TEST_EQUAL(a_limbs, n_limbs);
+
+    /* Higher-level conversion is in-place, so expected result must have the
+     * same number of limbs too */
+    TEST_EQUAL(x_limbs, n_limbs);
+
+    size_t limbs = n_limbs;
+    size_t bytes = limbs * sizeof(mbedtls_mpi_uint);
 
     TEST_EQUAL( 0, mbedtls_mpi_mod_modulus_setup( &m, N, n_limbs,
-                MBEDTLS_MPI_MOD_REP_MONTGOMERY ) );
+                                          MBEDTLS_MPI_MOD_REP_MONTGOMERY ) );
+
+    /* 1. Test low-level function first */
+
+    /* It has separate output, and requires temporary working storage */
+    size_t temp_limbs = mbedtls_mpi_core_montmul_working_limbs( limbs );
+    ASSERT_ALLOC( T, temp_limbs );
+    ASSERT_ALLOC( R, limbs );
+    mbedtls_mpi_core_from_mont_rep( R, A, N, n_limbs,
+                                    m.rep.mont.mm, T );
+    /* Test that the low-level function gives the required value */
+    ASSERT_COMPARE( R, bytes, X, bytes );
+
+    /* Test when output is aliased to input */
+    memcpy( R, A, bytes );
+    mbedtls_mpi_core_from_mont_rep( R, R, N, n_limbs,
+                                    m.rep.mont.mm, T );
+    ASSERT_COMPARE( R, bytes, X, bytes );
+
+    /* 2. Test higher-level Montgomery to cannonical conversion */
 
-    /* Convert from Montgomery into cannonical representation */
     TEST_EQUAL(0, mbedtls_mpi_mod_raw_from_mont_rep( A, &m ) );
 
     /* The result matches expected value */
-    ASSERT_COMPARE( A, x_bytes, X, x_bytes );
+    ASSERT_COMPARE( A, bytes, X, bytes );
+
 exit:
     mbedtls_mpi_mod_modulus_free( &m );
+    mbedtls_free( T );
     mbedtls_free( N );
     mbedtls_free( A );
+    mbedtls_free( R );
     mbedtls_free( X );
 }
 /* END_CASE */