/* Copyright (c) (2011,2012,2015,2016,2017,2018,2019,2020) Apple Inc. All rights reserved.
 *
 * corecrypto is licensed under Apple Inc.’s Internal Use License Agreement (which
 * is contained in the License.txt file distributed with corecrypto) and only to
 * people who accept that license. IMPORTANT:  Any license rights granted to you by
 * Apple Inc. (if any) are limited to internal use within your organization only on
 * devices and computers you own or control, for the sole purpose of verifying the
 * security characteristics and correct functioning of the Apple Software.  You may
 * not, directly or indirectly, redistribute the Apple Software or any portions thereof.
 */

#include "ccprime_internal.h"
#include "ccn_internal.h"

/*
 * Table of small known primes, generated by makePrimes tool.
 */
static const uint16_t ccn_prime_table[] = {
    3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37,
    41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89,
    97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151,
    157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223,
    227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281,
    283, 293, 307, 311, 313, 317, 331, 337, 347, 349, 353, 359,
    367, 373, 379, 383, 389, 397, 401, 409, 419, 421, 431, 433,
    439, 443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503,
    509, 521, 523, 541, 547, 557, 563, 569, 571, 577, 587, 593,
    599, 601, 607, 613, 617, 619, 631, 641, 643, 647, 653, 659,
    661, 673, 677, 683, 691, 701, 709, 719, 727, 733, 739, 743,
    751, 757, 761, 769, 773, 787, 797, 809, 811, 821, 823, 827,
    829, 839, 853, 857, 859, 863, 877, 881, 883, 887, 907, 911,
    919, 929, 937, 941, 947, 953, 967, 971, 977, 983, 991, 997,
    1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069,
    1087, 1091, 1093, 1097, 1103, 1109, 1117, 1123, 1129, 1151, 1153, 1163,
    1171, 1181, 1187, 1193, 1201, 1213, 1217, 1223, 1229, 1231, 1237, 1249,
    1259, 1277, 1279, 1283, 1289, 1291, 1297, 1301, 1303, 1307, 1319, 1321,
    1327, 1361, 1367, 1373, 1381, 1399, 1409, 1423, 1427, 1429, 1433, 1439,
    1447, 1451, 1453, 1459, 1471, 1481, 1483, 1487, 1489, 1493, 1499, 1511,
    1523, 1531, 1543, 1549, 1553, 1559, 1567, 1571, 1579, 1583, 1597, 1601,
    1607, 1609, 1613, 1619
};

const size_t ccn_prime_table_n = CC_ARRAY_LEN(ccn_prime_table);

// Check for obvious composites through trial division
// by two and all small known primes listed above.
CC_NONNULL_ALL
static cc_unit ccprime_find_small_prime_factor(cc_size n, const cc_unit *p)
{
    // Check for factor two.
    if ((p[0] & 1) == 0) {
        return 2;
    }

    for (unsigned i = 0; i < ccn_prime_table_n; i++) {
        if (ccn_divides1(n, p, ccn_prime_table[i])) {
            return ccn_prime_table[i];
        }
    }

    return 1;
}

/*
  Pick a pseudorandom base in [2, p).

  @see FIPS 186-4, C.3.1 Miller-Rabin Probabilistic Primality Test
*/
#define CCPRIME_PICK_RANDOM_BASE_EXTRA_BITS 64
#define CCPRIME_PICK_RANDOM_BASE_MAX_RETRIES 100
CC_NONNULL_ALL
static int ccprime_pick_random_base_ws(cc_ws_t ws, ccprime_mr_t mr, cc_unit *base, struct ccrng_state *rng)
{
    int rv = CCERR_INTERNAL;

    cczp_t zp = ccprime_mr_zp(mr);
    cc_size n = cczp_n(zp);

    CC_DECL_BP_WS(ws, bp)
    cc_unit *e = CC_ALLOC_WS(ws, 2 * n);
    ccn_clear(2 * n, e);

    cc_size e_nbits = cczp_bitlen(zp) + CCPRIME_PICK_RANDOM_BASE_EXTRA_BITS;

    // The REDC algorithm requires 0 <= e < R*p.
    // Adjust e_nbits for primes of length <= 64 bits.
    // Only 65+ bit primes will use the full extra entropy.
    if (CC_UNLIKELY(cczp_bitlen(zp) <= CCPRIME_PICK_RANDOM_BASE_EXTRA_BITS)) {
        e_nbits = cczp_bitlen(zp) + ccn_bitsof_n(n) - 1;
    }

    for (int i = 0; i < CCPRIME_PICK_RANDOM_BASE_MAX_RETRIES; i++) {
        rv = ccn_random_bits(e_nbits, e, rng);
        if (rv) {
            goto done;
        }

        // Compute base = e/R mod p.
        cczp_mod_ws(ws, zp, base, e);

        // Return if base < 2.
        if (!ccn_is_zero_or_one(n, base)) {
            goto done;
        }
    }

    // Looks like we couldn't find a good base?
    cc_try_abort("Failed to pick a random Miller-Rabin base");
    rv = CCERR_INTERNAL;

done:
    CC_FREE_BP_WS(ws, bp);
    return rv;
}

void ccprime_rabin_miller_init(cc_ws_t ws, ccprime_mr_t mr, cc_size n, const cc_unit *p)
{
    CC_DECL_BP_WS(ws, bp);

    // Use Montgomery multiplication to speed up exponentiation.
    cczp_t zpmm = ccprime_mr_zp(mr);
    cczp_mm_init_ws(ws, zpmm, n, p);

    cc_unit *pm1 = ccprime_mr_pm1(mr);
    cc_unit *s = ccprime_mr_s(mr);
    cc_unit *d = ccprime_mr_d(mr);

    // Compute p - 1.
    ccn_set(n, pm1, p);
    pm1[0] &= ~CC_UNIT_C(1);

    // Compute s and d, where p - 1 = 2^s * d.
    *s = (cc_unit)ccn_trailing_zeros(n, pm1);
    ccn_shift_right_multi(n, d, pm1, (size_t)*s);

    // Compute (p - 1) * R (mod p).
    // R * p = 0 (mod p) ==> R * (p - 1) = - R (mod p)
    ccn_sub(n, pm1, p, cczp_mm_r1(zpmm));

    CC_FREE_BP_WS(ws, bp);
}

// With s defined by p - 1 = 2^s * d, we need to iterate the main loop below in
// a way that doesn't reveal the value of s. To speed things up, we iterate up
// to 64 times to account for possible values of s in [1, 64). Anything above
// that will be rejected and is likely (2^-64) not a random prime candidate.
#define MAX_POWER_OF_TWO 64

int ccprime_rabin_miller_iteration(cc_ws_t ws, ccprime_mr_t mr, const cc_unit *base)
{
    const cc_unit *pm1 = ccprime_mr_pm1(mr);
    const cc_unit *d = ccprime_mr_d(mr);
    cc_unit s = *ccprime_mr_s(mr);

    // Reject excessive values of s for performance reasons.
    // p is very likely not a random prime candidate.
    if (CC_UNLIKELY(s >= MAX_POWER_OF_TWO)) {
        return 0;
    }

    cczp_const_t zp = ccprime_mr_zp(mr);
    cc_size n = cczp_n(zp);

    CC_DECL_BP_WS(ws, bp)
    cc_unit *x = CC_ALLOC_WS(ws, n);
    cc_unit *b = x;

    cczp_to_ws(ws, zp, b, base);

    // base (b) is a composite witness for p iff the following are true:
    //  - b^d != 1 (mod p)
    //  - b^(d*2^j) != -1 (mod p), for 0 <= j < s

    // Hide d's actual bit length.
    size_t d_bitlen = ccn_bitsof_n(n);

    int rv = cczp_powern_ws(ws, zp, x, b, d_bitlen, d);
    if (rv) {
        goto cleanup;
    }

    // If x = 1 or x = p-1, p is probably prime.
    cc_unit is_pm1 = ((cc_unit)ccn_cmp(n, pm1, x) & 1) ^ 1;
    cc_unit probably_prime = cczp_is_one_ws(ws, zp, x) | is_pm1;

    // To avoid leaking s, we iterate until r = MAX_POWER_OF_TWO-1.
    // We can exit early if we find a composite witness.
    for (cc_unit r = 1; r < MAX_POWER_OF_TWO; r++) {
        cc_unit r_neq_s;
        CC_HEAVISIDE_STEP(r_neq_s, r ^ s);

        // If r == s and we haven't seen x = 1 or x = p-1
        // then p is composite and we can "leak" that.
        if ((r_neq_s | probably_prime) ^ 1) {
            break;
        }

        cczp_sqr_ws(ws, zp, x, x);

        // If x = p-1, p is probably prime.
        probably_prime |= ((cc_unit)ccn_cmp(n, pm1, x) & 1) ^ 1;

        // Early out if it's clear we never reach p-1 via squaring.
        // p is composite, so we can "leak" this information.
        if (cczp_is_one_ws(ws, zp, x) & (probably_prime ^ 1)) {
            break;
        }
    }

    rv = (int)probably_prime;

cleanup:
    CC_FREE_BP_WS(ws, bp);
    return rv;
}

int ccprime_rabin_miller_ws(cc_ws_t ws, cc_size n, const cc_unit *p, size_t mr_depth, struct ccrng_state *rng)
{
    int rv = 0;

    if (mr_depth == 0) {
        return CCERR_PARAMETER;
    }

    cc_size np = ccn_n(n, p);

    // Zero and one aren't primes.
    if (np == 0 || (np == 1 && p[0] == 1)) {
        return 0;
    }

    // Try to find a small prime factor for p.
    cc_unit f = ccprime_find_small_prime_factor(n, p);

    // If we found a prime factor f>1 then either p=f or p is composite.
    if (f > 1) {
        if (np == 1 && p[0] == f) {
            return 1;
        }

        return 0;
    }

    CC_DECL_BP_WS(ws, bp);

    ccprime_mr_t mr = (ccprime_mr_t)CC_ALLOC_WS(ws, ccprime_mr_nof_n(n));
    ccprime_rabin_miller_init(ws, mr, n, p);
    cc_unit *base = CC_ALLOC_WS(ws, n);

    for (size_t j = 0; j < mr_depth; ++j) {
        // Pick a random base in [1, p).
        // Note: The base returned will be interpreted as being in Montgomery
        // space. So b will effectively be b/R, as we don't (need to) convert it.
        rv = ccprime_pick_random_base_ws(ws, mr, base, rng);
        if (rv) {
            break;
        }

        // One MR iteration with a random base.
        // Continue as long as p is _probably_ a prime.
        rv = ccprime_rabin_miller_iteration(ws, mr, base);
        if (rv != 1) {
            break;
        }
    }

    CC_FREE_BP_WS(ws, bp);
    return rv;
}

int ccprime_rabin_miller(cc_size n, const cc_unit *p, size_t mr_depth, struct ccrng_state *rng)
{
    CC_DECL_WORKSPACE_OR_FAIL(ws, CCPRIME_RABIN_MILLER_WORKSPACE_N(n));
    int rv = ccprime_rabin_miller_ws(ws, n, p, mr_depth, rng);
    CC_CLEAR_AND_FREE_WORKSPACE(ws);
    return rv;
}
