/*
 * LITLS -- X25519 key exchange
 * Extracted from TweetNaCl (public domain, DJB / Bernstein, Schwabe, et al.)
 * Single-letter variable names expanded for readability.
 */

#include "litls.h"
#include <string.h>

typedef int64_t fe[16]; /* Field element in GF(2^255-19), 16 limbs of ~16 bits */

static const fe fe_zero            = {0};
static const fe fe_one             = {1};
static const fe fe_121665          = {0xDB41, 1};
static const uint8_t basepoint[32] = {9};

static void fe_copy(fe out, const fe a) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = a[i];
}

static void fe_carry(fe o) {
    int i;
    int64_t c;
    for (i = 0; i < 16; i++) {
        o[i] += (1LL << 16);
        c = o[i] >> 16;
        o[(i + 1) * (i < 15)] += c - 1 + 37 * (c - 1) * (i == 15);
        o[i] -= c << 16;
    }
}

static void fe_cswap(fe p, fe q, int b) {
    int64_t t, c = ~(b - 1);
    int i;
    for (i = 0; i < 16; i++) {
        t = c & (p[i] ^ q[i]);
        p[i] ^= t;
        q[i] ^= t;
    }
}

static void fe_pack(uint8_t out[32], const fe n) {
    int i, j, b;
    fe m, t;
    for (i = 0; i < 16; i++)
        t[i] = n[i];
    fe_carry(t);
    fe_carry(t);
    fe_carry(t);
    for (j = 0; j < 2; j++) {
        m[0] = t[0] - 0xffed;
        for (i = 1; i < 15; i++) {
            m[i] = t[i] - 0xffff - ((m[i - 1] >> 16) & 1);
            m[i - 1] &= 0xffff;
        }
        m[15] = t[15] - 0x7fff - ((m[14] >> 16) & 1);
        b     = (m[15] >> 16) & 1;
        m[14] &= 0xffff;
        fe_cswap(t, m, 1 - b);
    }
    for (i = 0; i < 16; i++) {
        out[2 * i]     = t[i] & 0xff;
        out[2 * i + 1] = t[i] >> 8;
    }
}

static void fe_unpack(fe out, const uint8_t n[32]) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = n[2 * i] + ((int64_t)n[2 * i + 1] << 8);
    out[15] &= 0x7fff;
}

static void fe_add(fe out, const fe a, const fe b) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = a[i] + b[i];
}

static void fe_sub(fe out, const fe a, const fe b) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = a[i] - b[i];
}

static void fe_mul(fe out, const fe a, const fe b) {
    int64_t i, j, t[31];
    for (i = 0; i < 31; i++)
        t[i] = 0;
    for (i = 0; i < 16; i++)
        for (j = 0; j < 16; j++)
            t[i + j] += a[i] * b[j];
    for (i = 0; i < 15; i++)
        t[i] += 38 * t[i + 16];
    for (i = 0; i < 16; i++)
        out[i] = t[i];
    fe_carry(out);
    fe_carry(out);
}

static void fe_sq(fe out, const fe a) {
    fe_mul(out, a, a);
}

static void fe_inv(fe out, const fe in) {
    fe c;
    int a;
    for (a = 0; a < 16; a++)
        c[a] = in[a];
    /* Compute in^(2^255 - 21) = in^(p-2) via repeated squaring */
    for (a = 253; a >= 0; a--) {
        fe_sq(c, c);
        if (a != 2 && a != 4)
            fe_mul(c, c, in);
    }
    for (a = 0; a < 16; a++)
        out[a] = c[a];
}

/*
 * X25519 scalar multiplication using the Montgomery ladder.
 * Clamps the scalar per RFC 7748: clear bits 0,1,2 and 255; set bit 254.
 */
int litls_x25519(uint8_t shared[32], const uint8_t private_key[32], const uint8_t peer_public[32]) {
    uint8_t scalar[32];
    int64_t u[80];
    int64_t r;
    int i;
    fe a, b, c, d, e, f;

    for (i = 0; i < 32; i++)
        scalar[i] = private_key[i];
    scalar[31] = (scalar[31] & 127) | 64;
    scalar[0] &= 248;

    /* u[0..15] = unpacked peer u-coordinate */
    fe_unpack(u, peer_public);
    for (i = 0; i < 16; i++) {
        b[i] = u[i];
        d[i] = a[i] = c[i] = 0;
    }
    a[0] = d[0] = 1;

    /* Montgomery ladder */
    for (i = 254; i >= 0; --i) {
        r = (scalar[i >> 3] >> (i & 7)) & 1;
        fe_cswap(a, b, r);
        fe_cswap(c, d, r);
        fe_add(e, a, c);
        fe_sub(a, a, c);
        fe_add(c, b, d);
        fe_sub(b, b, d);
        fe_sq(d, e);
        fe_sq(f, a);
        fe_mul(a, c, a);
        fe_mul(c, b, e);
        fe_add(e, a, c);
        fe_sub(a, a, c);
        fe_sq(b, a);
        fe_sub(c, d, f);
        fe_mul(a, c, fe_121665);
        fe_add(a, a, d);
        fe_mul(c, c, a);
        fe_mul(a, d, f);
        fe_mul(d, b, u);
        fe_sq(b, e);
        fe_cswap(a, b, r);
        fe_cswap(c, d, r);
    }

    /* Final: result = a/c (projective to affine) */
    for (i = 0; i < 16; i++) {
        u[i + 16] = a[i];
        u[i + 32] = c[i];
    }
    fe_inv(u + 32, u + 32);
    fe_mul(u + 16, u + 16, u + 32);
    fe_pack(shared, u + 16);

    /* Reject all-zero output (low-order input point) */
    {
        uint8_t zeros[32] = {0};
        if (litls_secure_memcmp(shared, zeros, 32) == 0)
            return -1;
    }
    return 0;
}

void litls_x25519_base(uint8_t public_key[32], const uint8_t private_key[32]) {
    litls_x25519(public_key, private_key, basepoint);
}