/* BEGIN_HEADER */
#include "mbedtls/bignum.h"
#include "mbedtls/entropy.h"
#include "bignum_core.h"
#include "constant_time_internal.h"
#include "test/constant_flow.h"

/** Verifies mbedtls_mpi_core_add().
 *
 * \param[in] A       Little-endian presentation of the left operand.
 * \param[in] B       Little-endian presentation of the right operand.
 * \param limbs       Number of limbs in each MPI (\p A, \p B, \p S and \p X).
 * \param[in] S       Little-endian presentation of the expected sum.
 * \param carry       Expected carry from the addition.
 * \param[in,out] X   Temporary storage to be used for results.
 *
 * \return  1 if mbedtls_mpi_core_add() passes this test, otherwise 0.
 */
static int mpi_core_verify_add(mbedtls_mpi_uint *A,
                               mbedtls_mpi_uint *B,
                               size_t limbs,
                               mbedtls_mpi_uint *S,
                               int carry,
                               mbedtls_mpi_uint *X)
{
    int ret = 0;

    size_t bytes = limbs * sizeof(*A);

    /* The test cases have A <= B to avoid repetition, so we test A + B then,
     * if A != B, B + A. If A == B, we can test when A and B are aliased */

    /* A + B */

    /* A + B => correct result and carry */
    TEST_EQUAL(carry, mbedtls_mpi_core_add(X, A, B, limbs));
    ASSERT_COMPARE(X, bytes, S, bytes);

    /* A + B; alias output and first operand => correct result and carry */
    memcpy(X, A, bytes);
    TEST_EQUAL(carry, mbedtls_mpi_core_add(X, X, B, limbs));
    ASSERT_COMPARE(X, bytes, S, bytes);

    /* A + B; alias output and second operand => correct result and carry */
    memcpy(X, B, bytes);
    TEST_EQUAL(carry, mbedtls_mpi_core_add(X, A, X, limbs));
    ASSERT_COMPARE(X, bytes, S, bytes);

    if (memcmp(A, B, bytes) == 0) {
        /* A == B, so test where A and B are aliased */

        /* A + A => correct result and carry */
        TEST_EQUAL(carry, mbedtls_mpi_core_add(X, A, A, limbs));
        ASSERT_COMPARE(X, bytes, S, bytes);

        /* A + A, output aliased to both operands => correct result and carry */
        memcpy(X, A, bytes);
        TEST_EQUAL(carry, mbedtls_mpi_core_add(X, X, X, limbs));
        ASSERT_COMPARE(X, bytes, S, bytes);
    } else {
        /* A != B, so test B + A */

        /* B + A => correct result and carry */
        TEST_EQUAL(carry, mbedtls_mpi_core_add(X, B, A, limbs));
        ASSERT_COMPARE(X, bytes, S, bytes);

        /* B + A; alias output and first operand => correct result and carry */
        memcpy(X, B, bytes);
        TEST_EQUAL(carry, mbedtls_mpi_core_add(X, X, A, limbs));
        ASSERT_COMPARE(X, bytes, S, bytes);

        /* B + A; alias output and second operand => correct result and carry */
        memcpy(X, A, bytes);
        TEST_EQUAL(carry, mbedtls_mpi_core_add(X, B, X, limbs));
        ASSERT_COMPARE(X, bytes, S, bytes);
    }

    ret = 1;

exit:
    return ret;
}

/** Verifies mbedtls_mpi_core_add_if().
 *
 * \param[in] A       Little-endian presentation of the left operand.
 * \param[in] B       Little-endian presentation of the right operand.
 * \param limbs       Number of limbs in each MPI (\p A, \p B, \p S and \p X).
 * \param[in] S       Little-endian presentation of the expected sum.
 * \param carry       Expected carry from the addition.
 * \param[in,out] X   Temporary storage to be used for results.
 *
 * \return  1 if mbedtls_mpi_core_add_if() passes this test, otherwise 0.
 */
static int mpi_core_verify_add_if(mbedtls_mpi_uint *A,
                                  mbedtls_mpi_uint *B,
                                  size_t limbs,
                                  mbedtls_mpi_uint *S,
                                  int carry,
                                  mbedtls_mpi_uint *X)
{
    int ret = 0;

    size_t bytes = limbs * sizeof(*A);

    /* The test cases have A <= B to avoid repetition, so we test A + B then,
     * if A != B, B + A. If A == B, we can test when A and B are aliased */

    /* A + B */

    /* cond = 0 => X unchanged, no carry */
    memcpy(X, A, bytes);
    TEST_EQUAL(0, mbedtls_mpi_core_add_if(X, B, limbs, 0));
    ASSERT_COMPARE(X, bytes, A, bytes);

    /* cond = 1 => correct result and carry */
    TEST_EQUAL(carry, mbedtls_mpi_core_add_if(X, B, limbs, 1));
    ASSERT_COMPARE(X, bytes, S, bytes);

    if (memcmp(A, B, bytes) == 0) {
        /* A == B, so test where A and B are aliased */

        /* cond = 0 => X unchanged, no carry */
        memcpy(X, B, bytes);
        TEST_EQUAL(0, mbedtls_mpi_core_add_if(X, X, limbs, 0));
        ASSERT_COMPARE(X, bytes, B, bytes);

        /* cond = 1 => correct result and carry */
        TEST_EQUAL(carry, mbedtls_mpi_core_add_if(X, X, limbs, 1));
        ASSERT_COMPARE(X, bytes, S, bytes);
    } else {
        /* A != B, so test B + A */

        /* cond = 0 => d unchanged, no carry */
        memcpy(X, B, bytes);
        TEST_EQUAL(0, mbedtls_mpi_core_add_if(X, A, limbs, 0));
        ASSERT_COMPARE(X, bytes, B, bytes);

        /* cond = 1 => correct result and carry */
        TEST_EQUAL(carry, mbedtls_mpi_core_add_if(X, A, limbs, 1));
        ASSERT_COMPARE(X, bytes, S, bytes);
    }

    ret = 1;

exit:
    return ret;
}

/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_BIGNUM_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE */
void mpi_core_io_null()
{
    mbedtls_mpi_uint X = 0;
    int ret;

    ret = mbedtls_mpi_core_read_be(&X, 1, NULL, 0);
    TEST_EQUAL(ret, 0);
    ret = mbedtls_mpi_core_write_be(&X, 1, NULL, 0);
    TEST_EQUAL(ret, 0);

    ret = mbedtls_mpi_core_read_be(NULL, 0, NULL, 0);
    TEST_EQUAL(ret, 0);
    ret = mbedtls_mpi_core_write_be(NULL, 0, NULL, 0);
    TEST_EQUAL(ret, 0);

    ret = mbedtls_mpi_core_read_le(&X, 1, NULL, 0);
    TEST_EQUAL(ret, 0);
    ret = mbedtls_mpi_core_write_le(&X, 1, NULL, 0);
    TEST_EQUAL(ret, 0);

    ret = mbedtls_mpi_core_read_le(NULL, 0, NULL, 0);
    TEST_EQUAL(ret, 0);
    ret = mbedtls_mpi_core_write_le(NULL, 0, NULL, 0);
    TEST_EQUAL(ret, 0);

exit:
    ;
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_io_be(data_t *input, int nb_int, int nx_32_int, int iret,
                    int oret)
{
    if (iret != 0) {
        TEST_ASSERT(oret == 0);
    }

    TEST_LE_S(0, nb_int);
    size_t nb = nb_int;

    unsigned char buf[1024];
    TEST_LE_U(nb, sizeof(buf));

    /* nx_32_int is the number of 32 bit limbs, if we have 64 bit limbs we need
     * to halve the number of limbs to have the same size. */
    size_t nx;
    TEST_LE_S(0, nx_32_int);
    if (sizeof(mbedtls_mpi_uint) == 8) {
        nx = nx_32_int / 2 + nx_32_int % 2;
    } else {
        nx = nx_32_int;
    }

    mbedtls_mpi_uint X[sizeof(buf) / sizeof(mbedtls_mpi_uint)];
    TEST_LE_U(nx, sizeof(X) / sizeof(X[0]));

    int ret = mbedtls_mpi_core_read_be(X, nx, input->x, input->len);
    TEST_EQUAL(ret, iret);

    if (iret == 0) {
        ret =  mbedtls_mpi_core_write_be(X, nx, buf, nb);
        TEST_EQUAL(ret, oret);
    }

    if ((iret == 0) && (oret == 0)) {
        if (nb > input->len) {
            size_t leading_zeroes = nb - input->len;
            TEST_ASSERT(memcmp(buf + nb - input->len, input->x, input->len) == 0);
            for (size_t i = 0; i < leading_zeroes; i++) {
                TEST_EQUAL(buf[i], 0);
            }
        } else {
            size_t leading_zeroes = input->len - nb;
            TEST_ASSERT(memcmp(input->x + input->len - nb, buf, nb) == 0);
            for (size_t i = 0; i < leading_zeroes; i++) {
                TEST_EQUAL(input->x[i], 0);
            }
        }
    }

exit:
    ;
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_io_le(data_t *input, int nb_int, int nx_32_int, int iret,
                    int oret)
{
    if (iret != 0) {
        TEST_ASSERT(oret == 0);
    }

    TEST_LE_S(0, nb_int);
    size_t nb = nb_int;

    unsigned char buf[1024];
    TEST_LE_U(nb, sizeof(buf));

    /* nx_32_int is the number of 32 bit limbs, if we have 64 bit limbs we need
     * to halve the number of limbs to have the same size. */
    size_t nx;
    TEST_LE_S(0, nx_32_int);
    if (sizeof(mbedtls_mpi_uint) == 8) {
        nx = nx_32_int / 2 + nx_32_int % 2;
    } else {
        nx = nx_32_int;
    }

    mbedtls_mpi_uint X[sizeof(buf) / sizeof(mbedtls_mpi_uint)];
    TEST_LE_U(nx, sizeof(X) / sizeof(X[0]));

    int ret =  mbedtls_mpi_core_read_le(X, nx, input->x, input->len);
    TEST_EQUAL(ret, iret);

    if (iret == 0) {
        ret =  mbedtls_mpi_core_write_le(X, nx, buf, nb);
        TEST_EQUAL(ret, oret);
    }

    if ((iret == 0) && (oret == 0)) {
        if (nb > input->len) {
            TEST_ASSERT(memcmp(buf, input->x, input->len) == 0);
            for (size_t i = input->len; i < nb; i++) {
                TEST_EQUAL(buf[i], 0);
            }
        } else {
            TEST_ASSERT(memcmp(input->x, buf, nb) == 0);
            for (size_t i = nb; i < input->len; i++) {
                TEST_EQUAL(input->x[i], 0);
            }
        }
    }

exit:
    ;
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_bitlen(char *input_X, int nr_bits)
{
    mbedtls_mpi_uint *X = NULL;
    size_t limbs;

    TEST_EQUAL(mbedtls_test_read_mpi_core(&X, &limbs, input_X), 0);
    TEST_EQUAL(mbedtls_mpi_core_bitlen(X, limbs), nr_bits);

exit:
    mbedtls_free(X);
}
/* END_CASE */


/* BEGIN_CASE */
void mpi_core_clz(int leading_zeros, int trailing_zeros)
{
    if ((size_t) (leading_zeros + trailing_zeros) >= (sizeof(mbedtls_mpi_uint) * 8)) {
        // can't fit required number of leading and trailing zeros - skip test
        goto exit;
    }

    // Construct a test input value where the count of leading zeros and
    // trailing zeros is given in the test case, and we add ones to fill
    // the gap.
    mbedtls_mpi_uint x;
    if ((leading_zeros + trailing_zeros) > 0) {
        // some zero bits
        uint32_t s = (sizeof(mbedtls_mpi_uint) * 8 - leading_zeros - trailing_zeros);
        x = ((((mbedtls_mpi_uint) 1) << s) - 1) << trailing_zeros;
    } else {
        // all bits set
        x = ~((mbedtls_mpi_uint) 0);
    }

    size_t n = mbedtls_mpi_core_clz(x);
    TEST_EQUAL(n, leading_zeros);
exit:
    ;
}
/* END_CASE */


/* BEGIN_CASE */
void mpi_core_lt_ct(char *input_X, char *input_Y, int exp_ret)
{
    mbedtls_mpi_uint *X = NULL;
    size_t X_limbs;
    mbedtls_mpi_uint *Y = NULL;
    size_t Y_limbs;
    int ret;

    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&X, &X_limbs, input_X));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&Y, &Y_limbs, input_Y));

    /* We need two same-length limb arrays */
    TEST_EQUAL(X_limbs, Y_limbs);

    TEST_CF_SECRET(X, X_limbs * sizeof(mbedtls_mpi_uint));
    TEST_CF_SECRET(Y, X_limbs * sizeof(mbedtls_mpi_uint));

    ret = mbedtls_mpi_core_lt_ct(X, Y, X_limbs);
    TEST_EQUAL(ret, exp_ret);

exit:
    mbedtls_free(X);
    mbedtls_free(Y);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_uint_le_mpi(char *input_A)
{
    mbedtls_mpi_uint *A = NULL;
    size_t A_limbs = 0;

    TEST_EQUAL(mbedtls_test_read_mpi_core(&A, &A_limbs, input_A), 0);

    int is_large = 0; /* nonzero limbs beyond the lowest-order one? */
    for (size_t i = 1; i < A_limbs; i++) {
        if (A[i] != 0) {
            is_large = 1;
            break;
        }
    }

    TEST_CF_SECRET(A, A_limbs * sizeof(*A));

    TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi(0, A, A_limbs), 1);
    TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi(A[0], A, A_limbs), 1);

    if (is_large) {
        TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi(A[0] + 1,
                                                A, A_limbs), 1);
        TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi((mbedtls_mpi_uint) (-1) >> 1,
                                                A, A_limbs), 1);
        TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi((mbedtls_mpi_uint) (-1),
                                                A, A_limbs), 1);
    } else {
        TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi(A[0] + 1,
                                                A, A_limbs),
                   A[0] + 1 <= A[0]);
        TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi((mbedtls_mpi_uint) (-1) >> 1,
                                                A, A_limbs),
                   (mbedtls_mpi_uint) (-1) >> 1 <= A[0]);
        TEST_EQUAL(mbedtls_mpi_core_uint_le_mpi((mbedtls_mpi_uint) (-1),
                                                A, A_limbs),
                   (mbedtls_mpi_uint) (-1) <= A[0]);
    }

exit:
    mbedtls_free(A);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_cond_assign(char *input_X,
                          char *input_Y,
                          int input_bytes)
{
    mbedtls_mpi_uint *X = NULL;
    mbedtls_mpi_uint *Y = NULL;
    size_t limbs_X;
    size_t limbs_Y;

    TEST_EQUAL(mbedtls_test_read_mpi_core(&X, &limbs_X, input_X), 0);
    TEST_EQUAL(mbedtls_test_read_mpi_core(&Y, &limbs_Y, input_Y), 0);

    size_t limbs = limbs_X;
    size_t copy_limbs = CHARS_TO_LIMBS(input_bytes);
    size_t bytes = limbs * sizeof(mbedtls_mpi_uint);
    size_t copy_bytes = copy_limbs * sizeof(mbedtls_mpi_uint);

    TEST_EQUAL(limbs_X, limbs_Y);
    TEST_ASSERT(copy_limbs <= limbs);

    /* condition is false */
    TEST_CF_SECRET(X, bytes);
    TEST_CF_SECRET(Y, bytes);

    mbedtls_mpi_core_cond_assign(X, Y, copy_limbs, 0);

    TEST_CF_PUBLIC(X, bytes);
    TEST_CF_PUBLIC(Y, bytes);

    TEST_ASSERT(memcmp(X, Y, bytes) != 0);

    /* condition is true */
    TEST_CF_SECRET(X, bytes);
    TEST_CF_SECRET(Y, bytes);

    mbedtls_mpi_core_cond_assign(X, Y, copy_limbs, 1);

    TEST_CF_PUBLIC(X, bytes);
    TEST_CF_PUBLIC(Y, bytes);

    /* Check if the given length is copied even it is smaller
       than the length of the given MPIs. */
    if (copy_limbs < limbs) {
        TEST_CF_PUBLIC(X, bytes);
        TEST_CF_PUBLIC(Y, bytes);

        ASSERT_COMPARE(X, copy_bytes, Y, copy_bytes);
        TEST_ASSERT(memcmp(X, Y, bytes) != 0);
    } else {
        ASSERT_COMPARE(X, bytes, Y, bytes);
    }

exit:
    mbedtls_free(X);
    mbedtls_free(Y);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_cond_swap(char *input_X,
                        char *input_Y,
                        int input_bytes)
{
    mbedtls_mpi_uint *tmp_X = NULL;
    mbedtls_mpi_uint *tmp_Y = NULL;
    mbedtls_mpi_uint *X = NULL;
    mbedtls_mpi_uint *Y = NULL;
    size_t limbs_X;
    size_t limbs_Y;

    TEST_EQUAL(mbedtls_test_read_mpi_core(&tmp_X, &limbs_X, input_X), 0);
    TEST_EQUAL(mbedtls_test_read_mpi_core(&tmp_Y, &limbs_Y, input_Y), 0);

    size_t limbs = limbs_X;
    size_t copy_limbs = CHARS_TO_LIMBS(input_bytes);
    size_t bytes = limbs * sizeof(mbedtls_mpi_uint);
    size_t copy_bytes = copy_limbs * sizeof(mbedtls_mpi_uint);

    TEST_EQUAL(limbs_X, limbs_Y);
    TEST_ASSERT(copy_limbs <= limbs);

    ASSERT_ALLOC(X, limbs);
    memcpy(X, tmp_X, bytes);

    ASSERT_ALLOC(Y, limbs);
    memcpy(Y, tmp_Y, bytes);

    /* condition is false */
    TEST_CF_SECRET(X, bytes);
    TEST_CF_SECRET(Y, bytes);

    mbedtls_mpi_core_cond_swap(X, Y, copy_limbs, 0);

    TEST_CF_PUBLIC(X, bytes);
    TEST_CF_PUBLIC(Y, bytes);

    ASSERT_COMPARE(X, bytes, tmp_X, bytes);
    ASSERT_COMPARE(Y, bytes, tmp_Y, bytes);

    /* condition is true */
    TEST_CF_SECRET(X, bytes);
    TEST_CF_SECRET(Y, bytes);

    mbedtls_mpi_core_cond_swap(X, Y, copy_limbs, 1);

    TEST_CF_PUBLIC(X, bytes);
    TEST_CF_PUBLIC(Y, bytes);

    /* Check if the given length is copied even it is smaller
       than the length of the given MPIs. */
    if (copy_limbs < limbs) {
        ASSERT_COMPARE(X, copy_bytes, tmp_Y, copy_bytes);
        ASSERT_COMPARE(Y, copy_bytes, tmp_X, copy_bytes);
        TEST_ASSERT(memcmp(X, tmp_X, bytes) != 0);
        TEST_ASSERT(memcmp(X, tmp_Y, bytes) != 0);
        TEST_ASSERT(memcmp(Y, tmp_X, bytes) != 0);
        TEST_ASSERT(memcmp(Y, tmp_Y, bytes) != 0);
    } else {
        ASSERT_COMPARE(X, bytes, tmp_Y, bytes);
        ASSERT_COMPARE(Y, bytes, tmp_X, bytes);
    }

exit:
    mbedtls_free(tmp_X);
    mbedtls_free(tmp_Y);
    mbedtls_free(X);
    mbedtls_free(Y);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_shift_r(char *input, int count, char *result)
{
    mbedtls_mpi_uint *X = NULL;
    mbedtls_mpi_uint *Y = NULL;
    size_t limbs, n;

    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&X, &limbs, input));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&Y, &n, result));
    TEST_EQUAL(limbs, n);

    mbedtls_mpi_core_shift_r(X, limbs, count);
    ASSERT_COMPARE(X, limbs * ciL, Y, limbs * ciL);

exit:
    mbedtls_free(X);
    mbedtls_free(Y);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_add_and_add_if(char *input_A, char *input_B,
                             char *input_S, int carry)
{
    mbedtls_mpi_uint *A = NULL; /* first value to add */
    mbedtls_mpi_uint *B = NULL; /* second value to add */
    mbedtls_mpi_uint *S = NULL; /* expected result */
    mbedtls_mpi_uint *X = NULL; /* destination - the in/out first operand */
    size_t A_limbs, B_limbs, S_limbs;

    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&A, &A_limbs, input_A));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&B, &B_limbs, input_B));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&S, &S_limbs, input_S));

    /* add and add_if expect all operands to be the same length */
    TEST_EQUAL(A_limbs, B_limbs);
    TEST_EQUAL(A_limbs, S_limbs);

    size_t limbs = A_limbs;
    ASSERT_ALLOC(X, limbs);

    TEST_ASSERT(mpi_core_verify_add(A, B, limbs, S, carry, X));
    TEST_ASSERT(mpi_core_verify_add_if(A, B, limbs, S, carry, X));

exit:
    mbedtls_free(A);
    mbedtls_free(B);
    mbedtls_free(S);
    mbedtls_free(X);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_sub(char *input_A, char *input_B,
                  char *input_X, int carry)
{
    mbedtls_mpi A, B, X;
    mbedtls_mpi_uint *a = NULL;
    mbedtls_mpi_uint *b = NULL;
    mbedtls_mpi_uint *x = NULL; /* expected */
    mbedtls_mpi_uint *r = NULL; /* result */

    mbedtls_mpi_init(&A);
    mbedtls_mpi_init(&B);
    mbedtls_mpi_init(&X);

    TEST_EQUAL(0, mbedtls_test_read_mpi(&A, input_A));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&B, input_B));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&X, input_X));

    /* All of the inputs are +ve (or zero) */
    TEST_EQUAL(1, A.s);
    TEST_EQUAL(1, B.s);
    TEST_EQUAL(1, X.s);

    /* Get the number of limbs we will need */
    size_t limbs = MAX(A.n, B.n);
    size_t bytes = limbs * sizeof(mbedtls_mpi_uint);

    /* The result shouldn't have more limbs than the longest input */
    TEST_LE_U(X.n, limbs);

    /* Now let's get arrays of mbedtls_mpi_uints, rather than MPI structures */

    /* ASSERT_ALLOC() uses calloc() under the hood, so these do get zeroed */
    ASSERT_ALLOC(a, bytes);
    ASSERT_ALLOC(b, bytes);
    ASSERT_ALLOC(x, bytes);
    ASSERT_ALLOC(r, bytes);

    /* Populate the arrays. As the mbedtls_mpi_uint[]s in mbedtls_mpis (and as
     * processed by mbedtls_mpi_core_sub()) are little endian, we can just
     * copy what we have as long as MSBs are 0 (which they are from ASSERT_ALLOC())
     */
    memcpy(a, A.p, A.n * sizeof(mbedtls_mpi_uint));
    memcpy(b, B.p, B.n * sizeof(mbedtls_mpi_uint));
    memcpy(x, X.p, X.n * sizeof(mbedtls_mpi_uint));

    /* 1a) r = a - b => we should get the correct carry */
    TEST_EQUAL(carry, mbedtls_mpi_core_sub(r, a, b, limbs));

    /* 1b) r = a - b => we should get the correct result */
    ASSERT_COMPARE(r, bytes, x, bytes);

    /* 2 and 3 test "r may be aliased to a or b" */
    /* 2a) r = a; r -= b => we should get the correct carry (use r to avoid clobbering a) */
    memcpy(r, a, bytes);
    TEST_EQUAL(carry, mbedtls_mpi_core_sub(r, r, b, limbs));

    /* 2b) r -= b => we should get the correct result */
    ASSERT_COMPARE(r, bytes, x, bytes);

    /* 3a) r = b; r = a - r => we should get the correct carry (use r to avoid clobbering b) */
    memcpy(r, b, bytes);
    TEST_EQUAL(carry, mbedtls_mpi_core_sub(r, a, r, limbs));

    /* 3b) r = a - b => we should get the correct result */
    ASSERT_COMPARE(r, bytes, x, bytes);

    /* 4 tests "r may be aliased to [...] both" */
    if (A.n == B.n && memcmp(A.p, B.p, bytes) == 0) {
        memcpy(r, b, bytes);
        TEST_EQUAL(carry, mbedtls_mpi_core_sub(r, r, r, limbs));
        ASSERT_COMPARE(r, bytes, x, bytes);
    }

exit:
    mbedtls_free(a);
    mbedtls_free(b);
    mbedtls_free(x);
    mbedtls_free(r);

    mbedtls_mpi_free(&A);
    mbedtls_mpi_free(&B);
    mbedtls_mpi_free(&X);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_mla(char *input_A, char *input_B, char *input_S,
                  char *input_X4, char *input_cy4,
                  char *input_X8, char *input_cy8)
{
    /* We are testing A += B * s; A, B are MPIs, s is a scalar.
     *
     * However, we encode s as an MPI in the .data file as the test framework
     * currently only supports `int`-typed scalars, and that doesn't cover the
     * full range of `mbedtls_mpi_uint`.
     *
     * We also have the different results for sizeof(mbedtls_mpi_uint) == 4 or 8.
     */
    mbedtls_mpi A, B, S, X4, X8, cy4, cy8;
    mbedtls_mpi_uint *a = NULL;
    mbedtls_mpi_uint *x = NULL;

    mbedtls_mpi_init(&A);
    mbedtls_mpi_init(&B);
    mbedtls_mpi_init(&S);
    mbedtls_mpi_init(&X4);
    mbedtls_mpi_init(&X8);
    mbedtls_mpi_init(&cy4);
    mbedtls_mpi_init(&cy8);

    TEST_EQUAL(0, mbedtls_test_read_mpi(&A, input_A));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&B, input_B));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&S, input_S));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&X4, input_X4));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&cy4, input_cy4));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&X8, input_X8));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&cy8, input_cy8));

    /* The MPI encoding of scalar s must be only 1 limb */
    TEST_EQUAL(1, S.n);

    /* We only need to work with X4 or X8, and cy4 or cy8, depending on sizeof(mbedtls_mpi_uint) */
    mbedtls_mpi *X = (sizeof(mbedtls_mpi_uint) == 4) ? &X4 : &X8;
    mbedtls_mpi *cy = (sizeof(mbedtls_mpi_uint) == 4) ? &cy4 : &cy8;

    /* The carry should only have one limb */
    TEST_EQUAL(1, cy->n);

    /* All of the inputs are +ve (or zero) */
    TEST_EQUAL(1, A.s);
    TEST_EQUAL(1, B.s);
    TEST_EQUAL(1, S.s);
    TEST_EQUAL(1, X->s);
    TEST_EQUAL(1, cy->s);

    /* Get the (max) number of limbs we will need */
    size_t limbs = MAX(A.n, B.n);
    size_t bytes = limbs * sizeof(mbedtls_mpi_uint);

    /* The result shouldn't have more limbs than the longest input */
    TEST_LE_U(X->n, limbs);

    /* Now let's get arrays of mbedtls_mpi_uints, rather than MPI structures */

    /* ASSERT_ALLOC() uses calloc() under the hood, so these do get zeroed */
    ASSERT_ALLOC(a, bytes);
    ASSERT_ALLOC(x, bytes);

    /* Populate the arrays. As the mbedtls_mpi_uint[]s in mbedtls_mpis (and as
     * processed by mbedtls_mpi_core_mla()) are little endian, we can just
     * copy what we have as long as MSBs are 0 (which they are from ASSERT_ALLOC()).
     */
    memcpy(a, A.p, A.n * sizeof(mbedtls_mpi_uint));
    memcpy(x, X->p, X->n * sizeof(mbedtls_mpi_uint));

    /* 1a) A += B * s => we should get the correct carry */
    TEST_EQUAL(mbedtls_mpi_core_mla(a, limbs, B.p, B.n, *S.p), *cy->p);

    /* 1b) A += B * s => we should get the correct result */
    ASSERT_COMPARE(a, bytes, x, bytes);

    if (A.n == B.n && memcmp(A.p, B.p, bytes) == 0) {
        /* Check when A and B are aliased */
        memcpy(a, A.p, A.n * sizeof(mbedtls_mpi_uint));
        TEST_EQUAL(mbedtls_mpi_core_mla(a, limbs, a, limbs, *S.p), *cy->p);
        ASSERT_COMPARE(a, bytes, x, bytes);
    }

exit:
    mbedtls_free(a);
    mbedtls_free(x);

    mbedtls_mpi_free(&A);
    mbedtls_mpi_free(&B);
    mbedtls_mpi_free(&S);
    mbedtls_mpi_free(&X4);
    mbedtls_mpi_free(&X8);
    mbedtls_mpi_free(&cy4);
    mbedtls_mpi_free(&cy8);
}
/* END_CASE */


/* BEGIN_CASE */
void mpi_montg_init(char *input_N, char *input_mm)
{
    mbedtls_mpi N, mm;

    mbedtls_mpi_init(&N);
    mbedtls_mpi_init(&mm);

    TEST_EQUAL(0, mbedtls_test_read_mpi(&N, input_N));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&mm, input_mm));

    /* The MPI encoding of mm should be 1 limb (sizeof(mbedtls_mpi_uint) == 8) or
     * 2 limbs (sizeof(mbedtls_mpi_uint) == 4).
     *
     * The data file contains the expected result for sizeof(mbedtls_mpi_uint) == 8;
     * for sizeof(mbedtls_mpi_uint) == 4 it's just the LSW of this.
     */
    TEST_ASSERT(mm.n == 1  || mm.n == 2);

    /* All of the inputs are +ve (or zero) */
    TEST_EQUAL(1, N.s);
    TEST_EQUAL(1, mm.s);

    /* mbedtls_mpi_core_montmul_init() only returns a result, no error possible */
    mbedtls_mpi_uint result = mbedtls_mpi_core_montmul_init(N.p);

    /* Check we got the correct result */
    TEST_EQUAL(result, mm.p[0]);

exit:
    mbedtls_mpi_free(&N);
    mbedtls_mpi_free(&mm);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_montmul(int limbs_AN4, int limbs_B4,
                      int limbs_AN8, int limbs_B8,
                      char *input_A,
                      char *input_B,
                      char *input_N,
                      char *input_X4,
                      char *input_X8)
{
    mbedtls_mpi A, B, N, X4, X8, T, R;

    mbedtls_mpi_init(&A);
    mbedtls_mpi_init(&B);
    mbedtls_mpi_init(&N);
    mbedtls_mpi_init(&X4);      /* expected result, sizeof(mbedtls_mpi_uint) == 4 */
    mbedtls_mpi_init(&X8);      /* expected result, sizeof(mbedtls_mpi_uint) == 8 */
    mbedtls_mpi_init(&T);
    mbedtls_mpi_init(&R);       /* for the result */

    TEST_EQUAL(0, mbedtls_test_read_mpi(&A, input_A));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&B, input_B));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&N, input_N));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&X4, input_X4));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&X8, input_X8));

    mbedtls_mpi *X = (sizeof(mbedtls_mpi_uint) == 4) ? &X4 : &X8;

    int limbs_AN = (sizeof(mbedtls_mpi_uint) == 4) ? limbs_AN4 : limbs_AN8;
    int limbs_B = (sizeof(mbedtls_mpi_uint) == 4) ? limbs_B4 : limbs_B8;

    TEST_LE_U(A.n, (size_t) limbs_AN);
    TEST_LE_U(X->n, (size_t) limbs_AN);
    TEST_LE_U(B.n, (size_t) limbs_B);
    TEST_LE_U(limbs_B, limbs_AN);

    /* All of the inputs are +ve (or zero) */
    TEST_EQUAL(1, A.s);
    TEST_EQUAL(1, B.s);
    TEST_EQUAL(1, N.s);
    TEST_EQUAL(1, X->s);

    TEST_EQUAL(0, mbedtls_mpi_grow(&A, limbs_AN));
    TEST_EQUAL(0, mbedtls_mpi_grow(&N, limbs_AN));
    TEST_EQUAL(0, mbedtls_mpi_grow(X, limbs_AN));
    TEST_EQUAL(0, mbedtls_mpi_grow(&B, limbs_B));

    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);

    TEST_EQUAL(0, mbedtls_mpi_grow(&R, limbs_AN));     /* ensure it's got the right number of limbs */

    mbedtls_mpi_core_montmul(R.p, A.p, B.p, B.n, N.p, N.n, mm, T.p);
    size_t bytes = N.n * sizeof(mbedtls_mpi_uint);
    ASSERT_COMPARE(R.p, bytes, X->p, bytes);

    /* The output (R, above) may be aliased to A - use R to save the value of A */

    memcpy(R.p, A.p, bytes);

    mbedtls_mpi_core_montmul(A.p, A.p, B.p, B.n, N.p, N.n, mm, T.p);
    ASSERT_COMPARE(A.p, bytes, X->p, bytes);

    memcpy(A.p, R.p, bytes);    /* restore A */

    /* The output may be aliased to N - use R to save the value of N */

    memcpy(R.p, N.p, bytes);

    mbedtls_mpi_core_montmul(N.p, A.p, B.p, B.n, N.p, N.n, mm, T.p);
    ASSERT_COMPARE(N.p, bytes, X->p, bytes);

    memcpy(N.p, R.p, bytes);

    if (limbs_AN == limbs_B) {
        /* Test when A aliased to B (requires A == B on input values) */
        if (memcmp(A.p, B.p, bytes) == 0) {
            /* Test with A aliased to B and output, since this is permitted -
             * don't bother with yet another test with only A and B aliased */

            mbedtls_mpi_core_montmul(B.p, B.p, B.p, B.n, N.p, N.n, mm, T.p);
            ASSERT_COMPARE(B.p, bytes, X->p, bytes);

            memcpy(B.p, A.p, bytes);    /* restore B from equal value A */
        }

        /* The output may be aliased to B - last test, so we don't save B */

        mbedtls_mpi_core_montmul(B.p, A.p, B.p, B.n, N.p, N.n, mm, T.p);
        ASSERT_COMPARE(B.p, bytes, X->p, bytes);
    }

exit:
    mbedtls_mpi_free(&A);
    mbedtls_mpi_free(&B);
    mbedtls_mpi_free(&N);
    mbedtls_mpi_free(&X4);
    mbedtls_mpi_free(&X8);
    mbedtls_mpi_free(&T);
    mbedtls_mpi_free(&R);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_get_mont_r2_unsafe_neg()
{
    mbedtls_mpi N, RR;
    mbedtls_mpi_init(&N);
    mbedtls_mpi_init(&RR);
    const char *n = "7ffffffffffffff1";

    /* Test for zero divisor */
    TEST_EQUAL(MBEDTLS_ERR_MPI_DIVISION_BY_ZERO,
               mbedtls_mpi_core_get_mont_r2_unsafe(&RR, &N));

    /* Test for negative input */
    TEST_EQUAL(0, mbedtls_test_read_mpi(&N, n));
    N.s = -1;
    TEST_EQUAL(MBEDTLS_ERR_MPI_NEGATIVE_VALUE,
               mbedtls_mpi_core_get_mont_r2_unsafe(&RR, &N));
    N.s = 1;

exit:
    mbedtls_mpi_free(&N);
    mbedtls_mpi_free(&RR);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_get_mont_r2_unsafe(char *input_N,
                                 char *input_RR_X4,
                                 char *input_RR_X8)
{
    mbedtls_mpi N, RR, RR_REF;

    /* Select the appropriate output */
    char *input_rr = (sizeof(mbedtls_mpi_uint) == 4) ? input_RR_X4 : input_RR_X8;

    mbedtls_mpi_init(&N);
    mbedtls_mpi_init(&RR);
    mbedtls_mpi_init(&RR_REF);

    /* Read inputs */
    TEST_EQUAL(0, mbedtls_test_read_mpi(&N, input_N));
    TEST_EQUAL(0, mbedtls_test_read_mpi(&RR_REF, input_rr));

    /* All of the inputs are +ve (or zero) */
    TEST_EQUAL(1, N.s);
    TEST_EQUAL(1, RR_REF.s);

    /* Test valid input */
    TEST_EQUAL(0, mbedtls_mpi_core_get_mont_r2_unsafe(&RR, &N));

    /* Test that the moduli is odd */
    TEST_EQUAL(N.p[0] ^ 1, N.p[0] - 1);

    /* Output is +ve (or zero) */
    TEST_EQUAL(1, RR_REF.s);

    /* rr is updated to a valid pointer */
    TEST_ASSERT(RR.p != NULL);

    /* Calculated rr matches expected value */
    TEST_ASSERT(mbedtls_mpi_cmp_mpi(&RR, &RR_REF) == 0);

exit:
    mbedtls_mpi_free(&N);
    mbedtls_mpi_free(&RR);
    mbedtls_mpi_free(&RR_REF);
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_TEST_HOOKS */
void mpi_core_ct_uint_table_lookup(int bitlen, int window_size)
{
    size_t limbs = BITS_TO_LIMBS(bitlen);
    size_t count = ((size_t) 1) << window_size;

    mbedtls_mpi_uint *table = NULL;
    mbedtls_mpi_uint *dest = NULL;

    ASSERT_ALLOC(table, limbs * count);
    ASSERT_ALLOC(dest, limbs);

    /*
     * Fill the table with a unique counter so that differences are easily
     * detected. (And have their relationship to the index relatively non-trivial just
     * to be sure.)
     */
    for (size_t i = 0; i < count * limbs; i++) {
        table[i] = ~i - 1;
    }

    for (size_t i = 0; i < count; i++) {
        mbedtls_mpi_uint *current = table + i * limbs;
        memset(dest, 0x00, limbs * sizeof(*dest));

        /*
         * We shouldn't leak anything through timing.
         * We need to set these in every loop as we need to make the loop
         * variable public for the loop head and the buffers for comparison.
         */
        TEST_CF_SECRET(&i, sizeof(i));
        TEST_CF_SECRET(dest, limbs * sizeof(*dest));
        TEST_CF_SECRET(table, count * limbs * sizeof(*table));

        mbedtls_mpi_core_ct_uint_table_lookup(dest, table, limbs, count, i);

        TEST_CF_PUBLIC(dest, limbs * sizeof(*dest));
        TEST_CF_PUBLIC(table, count * limbs * sizeof(*table));
        ASSERT_COMPARE(dest, limbs * sizeof(*dest),
                       current, limbs * sizeof(*current));
        TEST_CF_PUBLIC(&i, sizeof(i));
    }

exit:
    mbedtls_free(table);
    mbedtls_free(dest);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_fill_random(int wanted_bytes_arg, int extra_rng_bytes,
                          int extra_limbs, int before, int expected_ret)
{
    size_t wanted_bytes = wanted_bytes_arg;
    mbedtls_mpi_uint *X = NULL;
    size_t X_limbs = CHARS_TO_LIMBS(wanted_bytes) + extra_limbs;
    size_t rng_bytes = wanted_bytes + extra_rng_bytes;
    unsigned char *rnd_data = NULL;
    mbedtls_test_rnd_buf_info rnd_info = { NULL, rng_bytes, NULL, NULL };
    int ret;

    /* Prepare an RNG with known output, limited to rng_bytes. */
    ASSERT_ALLOC(rnd_data, rng_bytes);
    TEST_EQUAL(0, mbedtls_test_rnd_std_rand(NULL, rnd_data, rng_bytes));
    rnd_info.buf = rnd_data;

    /* Allocate an MPI with room for wanted_bytes plus extra_limbs.
     * extra_limbs may be negative but the total limb count must be positive.
     * Fill the MPI with the byte value in before. */
    TEST_LE_U(1, X_limbs);
    ASSERT_ALLOC(X, X_limbs);
    memset(X, before, X_limbs * sizeof(*X));

    ret = mbedtls_mpi_core_fill_random(X, X_limbs, wanted_bytes,
                                       mbedtls_test_rnd_buffer_rand,
                                       &rnd_info);
    TEST_EQUAL(expected_ret, ret);

    if (expected_ret == 0) {
        /* mbedtls_mpi_core_fill_random is documented to use bytes from the
         * RNG as a big-endian representation of the number. We used an RNG
         * with known output, so check that the output contains the
         * expected value. Bytes above wanted_bytes must be zero. */
        for (size_t i = 0; i < wanted_bytes; i++) {
            mbedtls_test_set_step(i);
            TEST_EQUAL(GET_BYTE(X, i), rnd_data[wanted_bytes - 1 - i]);
        }
        for (size_t i = wanted_bytes; i < X_limbs * ciL; i++) {
            mbedtls_test_set_step(i);
            TEST_EQUAL(GET_BYTE(X, i), 0);
        }
    }

exit:
    mbedtls_free(rnd_data);
    mbedtls_free(X);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_mul(char *input_A,
                  char *input_B,
                  char *result)
{
    mbedtls_mpi_uint *A      = NULL;
    mbedtls_mpi_uint *A_orig = NULL;
    mbedtls_mpi_uint *B      = NULL;
    mbedtls_mpi_uint *B_orig = NULL;
    mbedtls_mpi_uint *R      = NULL;
    mbedtls_mpi_uint *X      = NULL;
    size_t A_limbs, B_limbs, R_limbs;

    TEST_EQUAL(mbedtls_test_read_mpi_core(&A, &A_limbs, input_A), 0);
    TEST_EQUAL(mbedtls_test_read_mpi_core(&B, &B_limbs, input_B), 0);
    TEST_EQUAL(mbedtls_test_read_mpi_core(&R, &R_limbs, result), 0);

    TEST_EQUAL(R_limbs, A_limbs + B_limbs);

    const size_t X_limbs = A_limbs + B_limbs;
    const size_t X_bytes = X_limbs * sizeof(mbedtls_mpi_uint);
    ASSERT_ALLOC(X, X_limbs);

    const size_t A_bytes = A_limbs * sizeof(mbedtls_mpi_uint);
    ASSERT_ALLOC(A_orig, A_limbs);
    memcpy(A_orig, A, A_bytes);

    const size_t B_bytes = B_limbs * sizeof(mbedtls_mpi_uint);
    ASSERT_ALLOC(B_orig, B_limbs);
    memcpy(B_orig, B, B_bytes);

    /* Set result to something that is unlikely to be correct */
    memset(X, '!', X_bytes);

    /* 1. X = A * B - result should be correct, A and B unchanged */
    mbedtls_mpi_core_mul(X, A, A_limbs, B, B_limbs);
    ASSERT_COMPARE(X, X_bytes, R, X_bytes);
    ASSERT_COMPARE(A, A_bytes, A_orig, A_bytes);
    ASSERT_COMPARE(B, B_bytes, B_orig, B_bytes);

    /* 2. A == B: alias A and B - result should be correct, A and B unchanged */
    if (A_bytes == B_bytes && memcmp(A, B, A_bytes) == 0) {
        memset(X, '!', X_bytes);
        mbedtls_mpi_core_mul(X, A, A_limbs, A, A_limbs);
        ASSERT_COMPARE(X, X_bytes, R, X_bytes);
        ASSERT_COMPARE(A, A_bytes, A_orig, A_bytes);
    }
    /* 3. X = B * A - result should be correct, A and B unchanged */
    else {
        memset(X, '!', X_bytes);
        mbedtls_mpi_core_mul(X, B, B_limbs, A, A_limbs);
        ASSERT_COMPARE(X, X_bytes, R, X_bytes);
        ASSERT_COMPARE(A, A_bytes, A_orig, A_bytes);
        ASSERT_COMPARE(B, B_bytes, B_orig, B_bytes);
    }

exit:
    mbedtls_free(A);
    mbedtls_free(A_orig);
    mbedtls_free(B);
    mbedtls_free(B_orig);
    mbedtls_free(R);
    mbedtls_free(X);
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_exp_mod(char *input_N, char *input_A,
                      char *input_E, char *input_X)
{
    mbedtls_mpi_uint *A = NULL;
    mbedtls_mpi_uint *E = NULL;
    mbedtls_mpi_uint *N = NULL;
    mbedtls_mpi_uint *X = NULL;
    size_t A_limbs, E_limbs, N_limbs, X_limbs;
    const mbedtls_mpi_uint *R2 = NULL;
    mbedtls_mpi_uint *Y = NULL;
    mbedtls_mpi_uint *T = NULL;
    /* Legacy MPIs for computing R2 */
    mbedtls_mpi N_mpi;
    mbedtls_mpi_init(&N_mpi);
    mbedtls_mpi R2_mpi;
    mbedtls_mpi_init(&R2_mpi);

    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&A, &A_limbs, input_A));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&E, &E_limbs, input_E));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&N, &N_limbs, input_N));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&X, &X_limbs, input_X));
    ASSERT_ALLOC(Y, N_limbs);

    TEST_EQUAL(A_limbs, N_limbs);
    TEST_EQUAL(X_limbs, N_limbs);

    TEST_EQUAL(0, mbedtls_mpi_grow(&N_mpi, N_limbs));
    memcpy(N_mpi.p, N, N_limbs * sizeof(*N));
    N_mpi.n = N_limbs;
    TEST_EQUAL(0,
               mbedtls_mpi_core_get_mont_r2_unsafe(&R2_mpi, &N_mpi));
    TEST_EQUAL(0, mbedtls_mpi_grow(&R2_mpi, N_limbs));
    R2 = R2_mpi.p;

    size_t working_limbs = mbedtls_mpi_core_exp_mod_working_limbs(N_limbs,
                                                                  E_limbs);

    /* No point exactly duplicating the code in mbedtls_mpi_core_exp_mod_working_limbs()
     * to see if the output is correct, but we can check that it's in a
     * reasonable range.  The current calculation works out as
     * `1 + N_limbs * (welem + 3)`, where welem is the number of elements in
     * the window (1 << 1 up to 1 << 6).
     */
    size_t min_expected_working_limbs = 1 + N_limbs * 4;
    size_t max_expected_working_limbs = 1 + N_limbs * 67;

    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);

    TEST_EQUAL(0, memcmp(X, Y, N_limbs * sizeof(mbedtls_mpi_uint)));

    /* Check when output aliased to input */

    mbedtls_mpi_core_exp_mod(A, A, N, N_limbs, E, E_limbs, R2, T);

    TEST_EQUAL(0, memcmp(X, A, N_limbs * sizeof(mbedtls_mpi_uint)));

exit:
    mbedtls_free(T);
    mbedtls_free(A);
    mbedtls_free(E);
    mbedtls_free(N);
    mbedtls_free(X);
    mbedtls_free(Y);
    mbedtls_mpi_free(&N_mpi);
    mbedtls_mpi_free(&R2_mpi);
    // R2 doesn't need to be freed as it is only aliasing R2_mpi
}
/* END_CASE */

/* BEGIN_CASE */
void mpi_core_sub_int(char *input_A, char *input_B,
                      char *input_X, int borrow)
{
    /* We are testing A - b, where A is an MPI and b is a scalar, expecting
     * result X with borrow borrow.  However, for ease of handling we encode b
     * as a 1-limb MPI (B) in the .data file. */

    mbedtls_mpi_uint *A = NULL;
    mbedtls_mpi_uint *B = NULL;
    mbedtls_mpi_uint *X = NULL;
    mbedtls_mpi_uint *R = NULL;
    size_t A_limbs, B_limbs, X_limbs;

    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&A, &A_limbs, input_A));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&B, &B_limbs, input_B));
    TEST_EQUAL(0, mbedtls_test_read_mpi_core(&X, &X_limbs, input_X));

    /* The MPI encoding of scalar b must be only 1 limb */
    TEST_EQUAL(B_limbs, 1);

    /* The subtraction is fixed-width, so A and X must have the same number of limbs */
    TEST_EQUAL(A_limbs, X_limbs);
    size_t limbs = A_limbs;

    ASSERT_ALLOC(R, limbs);

#define TEST_COMPARE_CORE_MPIS(A, B, limbs) \
    ASSERT_COMPARE(A, (limbs) * sizeof(mbedtls_mpi_uint), B, (limbs) * sizeof(mbedtls_mpi_uint))

    /* 1. R = A - b. Result and borrow should be correct */
    TEST_EQUAL(mbedtls_mpi_core_sub_int(R, A, B[0], limbs), borrow);
    TEST_COMPARE_CORE_MPIS(R, X, limbs);

    /* 2. A = A - b. Result and borrow should be correct */
    TEST_EQUAL(mbedtls_mpi_core_sub_int(A, A, B[0], limbs), borrow);
    TEST_COMPARE_CORE_MPIS(A, X, limbs);

exit:
    mbedtls_free(A);
    mbedtls_free(B);
    mbedtls_free(X);
    mbedtls_free(R);
}
/* 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 */