/*
 *  Copyright (C) 2021 - This file is part of libecc project
 *
 *  Authors:
 *      Ryad BENADJILA <ryadbenadjila@gmail.com>
 *      Arnaud EBALARD <arnaud.ebalard@ssi.gouv.fr>
 *      Jean-Pierre FLORI <jean-pierre.flori@ssi.gouv.fr>
 *
 *  Contributors:
 *      Nicolas VIVET <nicolas.vivet@ssi.gouv.fr>
 *      Karim KHALFALLAH <karim.khalfallah@ssi.gouv.fr>
 *
 *  This software is licensed under a dual BSD and GPL v2 license.
 *  See LICENSE file at the root folder of the project.
 */
#include <libecc/nn/nn_mul_redc1.h>
#include <libecc/nn/nn_div_public.h>
#include <libecc/nn/nn_logical.h>
#include <libecc/nn/nn_mod_pow.h>
#include <libecc/nn/nn_rand.h>
#include <libecc/nn/nn.h>

/*
 * NOT constant time with regard to the bitlength of exp.
 *
 * Implements a left to right Montgomery Ladder for modular exponentiation.
 * This is an internal common helper and assumes that all the initialization
 * and aliasing of inputs/outputs are handled by the callers. Depending on
 * the inputs, redcification is optionally used.
 * The base is reduced if necessary.
 *
 * Montgomery Ladder is masked using Itoh et al. anti-ADPA
 * (Address-bit DPA) countermeasure.
 * See "A Practical Countermeasure against Address-Bit Differential Power Analysis"
 * by Itoh, Izu and Takenaka for more information.

 * Returns 0 on success, -1 on error.
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_exp_monty_ladder_ltr(nn_t out, nn_src_t base, nn_src_t exp, nn_src_t mod, nn_src_t r, nn_src_t r_square, word_t mpinv)
{
        nn T[3];
        nn mask;
        bitcnt_t explen, oldexplen;
        u8 expbit, rbit;
        int ret, cmp;
        T[0].magic = T[1].magic = T[2].magic = mask.magic = WORD(0);

        /* Initialize out */
        ret = nn_init(out, 0); EG(ret, err);

        ret = nn_init(&T[0], 0); EG(ret, err);
        ret = nn_init(&T[1], 0); EG(ret, err);
        ret = nn_init(&T[2], 0); EG(ret, err);

        /* Generate our Itoh random mask */
        ret = nn_get_random_len(&mask, NN_MAX_BYTE_LEN); EG(ret, err);

        ret = nn_bitlen(exp, &explen); EG(ret, err);


        /* From now on, since we deal with Itoh's countermeasure, we must have at
         * least 2 bits in the exponent. We will deal with the particular cases of 0 and 1
         * bit exponents later.
         */
        oldexplen = explen;
        explen = (explen < 2) ? 2 : explen;

        ret = nn_getbit(&mask, (bitcnt_t)(explen - 1), &rbit); EG(ret, err);

        /* Reduce the base if necessary */
        ret = nn_cmp(base, mod, &cmp); EG(ret, err);
        if(cmp >= 0){
                /* Modular reduction */
                ret = nn_mod(&T[rbit], base, mod); EG(ret, err);
                if(r != NULL){
                        /* Redcify the base if necessary */
                        ret = nn_mul_redc1(&T[rbit], &T[rbit], r_square, mod, mpinv); EG(ret, err);
                }
        }
        else{
                if(r != NULL){
                        /* Redcify the base if necessary */
                        ret = nn_mul_redc1(&T[rbit], base, r_square, mod, mpinv); EG(ret, err);
                }
                else{
                        ret = nn_copy(&T[rbit], base); EG(ret, err);
                }
        }

        /* We implement the Montgomery ladder exponentiation with Itoh masking using three
         * registers T[0], T[1] and T[2]. The random mask is in 'mask'.
         */
        if(r != NULL){
                ret = nn_mul_redc1(&T[1-rbit], &T[rbit], &T[rbit], mod, mpinv); EG(ret, err);
        }
        else{
                ret = nn_mod_mul(&T[1-rbit], &T[rbit], &T[rbit], mod); EG(ret, err);
        }

        /* Now proceed with the Montgomery Ladder algorithm.
         */
        explen = (bitcnt_t)(explen - 1);
        while (explen > 0) {
                u8 rbit_next;
                explen = (bitcnt_t)(explen - 1);

                /* rbit is r[i+1], and rbit_next is r[i] */
                ret = nn_getbit(&mask, explen, &rbit_next); EG(ret, err);
                /* Get the exponent bit */
                ret = nn_getbit(exp, explen, &expbit); EG(ret, err);

                /* Square */
                if(r != NULL){
                        ret = nn_mul_redc1(&T[2], &T[expbit ^ rbit], &T[expbit ^ rbit], mod, mpinv); EG(ret, err);
                }
                else{
                        ret = nn_mod_mul(&T[2], &T[expbit ^ rbit], &T[expbit ^ rbit], mod); EG(ret, err);
                }
                /* Multiply */
                if(r != NULL){
                        ret = nn_mul_redc1(&T[1], &T[0], &T[1], mod, mpinv); EG(ret, err);
                }
                else{
                        ret = nn_mod_mul(&T[1], &T[0], &T[1], mod); EG(ret, err);
                }
                /* Copy */
                ret = nn_copy(&T[0], &T[2 - (expbit ^ rbit_next)]); EG(ret, err);
                ret = nn_copy(&T[1], &T[1 + (expbit ^ rbit_next)]); EG(ret, err);
                /* Update rbit */
                rbit = rbit_next;
        }
        ret = nn_one(&T[1 - rbit]);
        if(r != NULL){
                /* Unredcify in out */
                ret = nn_mul_redc1(&T[rbit], &T[rbit], &T[1 - rbit], mod, mpinv); EG(ret, err);
        }

        /* Deal with the particular cases of 0 and 1 bit exponents */
        /* Case with 0 bit exponent: T[1 - rbit] contains 1 modulo mod */
        ret = nn_mod(&T[1 - rbit], &T[1 - rbit], mod); EG(ret, err);
        /* Case with 1 bit exponent */
        ret = nn_mod(&T[2], base, mod); EG(ret, err);
        /* Proceed with the output */
        ret = nn_cnd_swap((oldexplen == 0), out, &T[1 - rbit]);
        ret = nn_cnd_swap((oldexplen == 1), out, &T[2]);
        ret = nn_cnd_swap(((oldexplen != 0) && (oldexplen != 1)), out, &T[rbit]);

err:
        nn_uninit(&T[0]);
        nn_uninit(&T[1]);
        nn_uninit(&T[2]);
        nn_uninit(&mask);

        return ret;
}

/*
 * NOT constant time with regard to the bitlength of exp.
 *
 * Reduces the base modulo mod if it is not already reduced,
 * which is also a small divergence wrt constant time leaking
 * the information that base <= mod or not: please use with care
 * in the callers if this information is sensitive.
 *
 * Aliasing not supported. Expects caller to check parameters
 * have been initialized. This is an internal helper.
 *
 * Compute (base ** exp) mod (mod) using a Montgomery Ladder algorithm
 * with Montgomery redcification, hence the Montgomery coefficients as input.
 * The module "mod" is expected to be odd for redcification to be used.
 *
 * Returns 0 on success, -1 on error.
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_mod_pow_redc(nn_t out, nn_src_t base, nn_src_t exp, nn_src_t mod, nn_src_t r, nn_src_t r_square, word_t mpinv)
{
        return _nn_exp_monty_ladder_ltr(out, base, exp, mod, r, r_square, mpinv);
}

/*
 * NOT constant time with regard to the bitlength of exp.
 *
 * Reduces the base modulo mod if it is not already reduced,
 * which is also a small divergence wrt constant time leaking
 * the information that base <= mod or not: please use with care
 * in the callers if this information is sensitive.
 *
 * Aliasing is supported. Expects caller to check parameters
 * have been initialized. This is an internal helper.
 *
 * Compute (base ** exp) mod (mod) using a Montgomery Ladder algorithm.
 * This function works for all values of "mod", but is slower that the one
 * using Montgomery multiplication (which only works for odd "mod"). Hence,
 * it is only used on even "mod" by upper layers.
 *
 * Returns 0 on success, -1 on error.
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_mod_pow(nn_t out, nn_src_t base, nn_src_t exp, nn_src_t mod)
{
        int ret;

        if ((out == base) || (out == exp) || (out == mod)) {
                nn _out;
                _out.magic = WORD(0);

                ret = nn_init(&_out, 0); EG(ret, err);
                ret = _nn_exp_monty_ladder_ltr(&_out, base, exp, mod, NULL, NULL, WORD(0)); EG(ret, err);
                ret = nn_copy(out, &_out);
        }
        else{
                ret = _nn_exp_monty_ladder_ltr(out, base, exp, mod, NULL, NULL, WORD(0));
        }

err:
        return ret;
}

/*
 * Same purpose as above but handles aliasing.
 * Expects caller to check parameters
 * have been initialized. This is an internal helper.
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_mod_pow_redc_aliased(nn_t out, nn_src_t base, nn_src_t exp, nn_src_t mod, nn_src_t r, nn_src_t r_square, word_t mpinv)
{
        nn _out;
        int ret;
        _out.magic = WORD(0);

        ret = nn_init(&_out, 0); EG(ret, err);
        ret = _nn_mod_pow_redc(&_out, base, exp, mod, r, r_square, mpinv); EG(ret, err);
        ret = nn_copy(out, &_out);

err:
        nn_uninit(&_out);

        return ret;
}

/* Aliased version of previous one.
 * NOTE: our nn_mod_pow_redc primitives suppose that the modulo is odd for Montgomery multiplication
 * primitives to provide consistent results.
 *
 * Aliasing is supported.
 */
int nn_mod_pow_redc(nn_t out, nn_src_t base, nn_src_t exp, nn_src_t mod, nn_src_t r, nn_src_t r_square, word_t mpinv)
{
        int ret, isodd;

        ret = nn_check_initialized(base); EG(ret, err);
        ret = nn_check_initialized(exp); EG(ret, err);
        ret = nn_check_initialized(mod); EG(ret, err);
        ret = nn_check_initialized(r); EG(ret, err);
        ret = nn_check_initialized(r_square); EG(ret, err);

        /* Check that we have an odd number for our modulo */
        ret = nn_isodd(mod, &isodd); EG(ret, err);
        MUST_HAVE(isodd, ret, err);

        /* Handle the case where our prime is less than two words size.
         * We need it to be >= 2 words size for the Montgomery multiplication to be
         * usable.
         */
        if(mod->wlen < 2){
                /* Local copy our modulo */
                nn _mod;
                _mod.magic = WORD(0);

                /* And set its length accordingly */
                ret = nn_copy(&_mod, mod); EG(ret, err1);
                ret = nn_set_wlen(&_mod, 2); EG(ret, err1);
                /* Handle output aliasing */
                if ((out == base) || (out == exp) || (out == mod) || (out == r) || (out == r_square)) {
                        ret = _nn_mod_pow_redc_aliased(out, base, exp, &_mod, r, r_square, mpinv); EG(ret, err1);
                } else {
                        ret = _nn_mod_pow_redc(out, base, exp, &_mod, r, r_square, mpinv); EG(ret, err1);
                }
err1:
                nn_uninit(&_mod);
                EG(ret, err);
        }
        else{
                /* Handle output aliasing */
                if ((out == base) || (out == exp) || (out == mod) || (out == r) || (out == r_square)) {
                        ret = _nn_mod_pow_redc_aliased(out, base, exp, mod, r, r_square, mpinv);
                } else {
                        ret = _nn_mod_pow_redc(out, base, exp, mod, r, r_square, mpinv);
                }
        }

err:
        return ret;
}


/*
 * NOT constant time with regard to the bitlength of exp.
 * Aliasing is supported.
 *
 * Compute (base ** exp) mod (mod) using a Montgomery Ladder algorithm.
 * Internally, this computes Montgomery coefficients and uses the redc
 * function.
 *
 * Returns 0 on success, -1 on error.
 */
int nn_mod_pow(nn_t out, nn_src_t base, nn_src_t exp, nn_src_t mod)
{
        nn r, r_square;
        word_t mpinv;
        int ret, isodd;
        r.magic = r_square.magic = WORD(0);

        /* Handle the case where our modulo is even: in this case, we cannot
         * use our Montgomery multiplication primitives as they are only suitable
         * for odd numbers. We fallback to less efficient regular modular exponentiation.
         */
        ret = nn_isodd(mod, &isodd); EG(ret, err);
        if(!isodd){
                /* mod is even: use the regular unoptimized modular exponentiation */
                ret = _nn_mod_pow(out, base, exp, mod);
        }
        else{
                /* mod is odd */
                /* Compute the Montgomery coefficients */
                ret = nn_compute_redc1_coefs(&r, &r_square, mod, &mpinv); EG(ret, err);

                /* Now use the redc version */
                ret = nn_mod_pow_redc(out, base, exp, mod, &r, &r_square, mpinv);
        }

err:
        nn_uninit(&r);
        nn_uninit(&r_square);

        return ret;
}