/*

 * Copyright (c) 2022, stelar7 <dudedbz@gmail.com>

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#include <AK/Random.h>

#include <LibCrypto/Curves/Curve25519.h>

#include <LibCrypto/Curves/Ed25519.h>

#include <LibCrypto/Hash/SHA2.h>

namespace Crypto::Curves {

namespace {

struct Ed25519Point {

    u32 x[8] {};

    u32 y[8] {};

    u32 z[8] {};

    u32 t[8] {};

};

constexpr Ed25519Point BASE_POINT = {

    { 0x8F25D51A, 0xC9562D60, 0x9525A7B2, 0x692CC760, 0xFDD6DC5C, 0xC0A4E231, 0xCD6E53FE, 0x216936D3 },

    { 0x66666658, 0x66666666, 0x66666666, 0x66666666, 0x66666666, 0x66666666, 0x66666666, 0x66666666 },

    { 0x00000001, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000 },

    { 0xA5B7DDA3, 0x6DDE8AB3, 0x775152F5, 0x20F09F80, 0x64ABE37D, 0x66EA4E8E, 0xD78B7665, 0x67875F0F }

};

void point_double(Ed25519Point* result, Ed25519Point const* point)

{

    u32 a[8] {};

    u32 b[8] {};

    u32 c[8] {};

    u32 e[8] {};

    u32 f[8] {};

    u32 g[8] {};

    u32 h[8] {};

    Curve25519::modular_square(a, point->x);

    Curve25519::modular_square(b, point->y);

    Curve25519::modular_square(c, point->z);

    Curve25519::modular_add(c, c, c);

    Curve25519::modular_add(e, a, b);

    Curve25519::modular_add(f, point->x, point->y);

    Curve25519::modular_square(f, f);

    Curve25519::modular_subtract(f, e, f);

    Curve25519::modular_subtract(g, a, b);

    Curve25519::modular_add(h, c, g);

    Curve25519::modular_multiply(result->x, f, h);

    Curve25519::modular_multiply(result->y, e, g);

    Curve25519::modular_multiply(result->z, g, h);

    Curve25519::modular_multiply(result->t, e, f);

}

void point_add(Ed25519Point* result, Ed25519Point const* p, Ed25519Point const* q)

{

    // Compute R = P + Q

    u32 a[8] {};

    u32 b[8] {};

    u32 c[8] {};

    u32 d[8] {};

    u32 e[8] {};

    u32 f[8] {};

    u32 g[8] {};

    u32 h[8] {};

    Curve25519::modular_add(c, p->y, p->x);

    Curve25519::modular_add(d, q->y, q->x);

    Curve25519::modular_multiply(a, c, d);

    Curve25519::modular_subtract(c, p->y, p->x);

    Curve25519::modular_subtract(d, q->y, q->x);

    Curve25519::modular_multiply(b, c, d);

    Curve25519::modular_multiply(c, p->z, q->z);

    Curve25519::modular_add(c, c, c);

    Curve25519::modular_multiply(d, p->t, q->t);

    Curve25519::modular_multiply(d, d, Curve25519::CURVE_D_2);

    Curve25519::modular_add(e, a, b);

    Curve25519::modular_subtract(f, a, b);

    Curve25519::modular_add(g, c, d);

    Curve25519::modular_subtract(h, c, d);

    Curve25519::modular_multiply(result->x, f, h);

    Curve25519::modular_multiply(result->y, e, g);

    Curve25519::modular_multiply(result->z, g, h);

    Curve25519::modular_multiply(result->t, e, f);

}

void point_multiply_scalar(Ed25519Point* result, u8 const* scalar, Ed25519Point const* point)

{

    // Set U to the neutral element (0, 1, 1, 0)

    Ed25519Point u {};

    Ed25519Point v {};

    Curve25519::set(u.x, 0);

    Curve25519::set(u.y, 1);

    Curve25519::set(u.z, 1);

    Curve25519::set(u.t, 0);

    for (i32 i = Curve25519::BITS - 1; i >= 0; i--) {

        u8 b = (scalar[i / 8] >> (i % 8)) & 1;

        // Compute U = 2 * U

        point_double(&u, &u);

        // Compute V = U + P

        point_add(&v, &u, point);

        // If b is set, then U = V

        Curve25519::select(u.x, u.x, v.x, b);

        Curve25519::select(u.y, u.y, v.y, b);

        Curve25519::select(u.z, u.z, v.z, b);

        Curve25519::select(u.t, u.t, v.t, b);

    }

    Curve25519::copy(result->x, u.x);

    Curve25519::copy(result->y, u.y);

    Curve25519::copy(result->z, u.z);

    Curve25519::copy(result->t, u.t);

}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.2

void encode_point(Ed25519Point* point, u8* data)

{

    // Retrieve affine representation

    Curve25519::modular_multiply_inverse(point->z, point->z);

    Curve25519::modular_multiply(point->x, point->x, point->z);

    Curve25519::modular_multiply(point->y, point->y, point->z);

    Curve25519::set(point->z, 1);

    Curve25519::modular_multiply(point->t, point->x, point->y);

    // First, encode the y-coordinate (in the range 0 <= y < p) as a little-endian string of 32 octets.

    // The most significant bit of the final octet is always zero.

    Curve25519::export_state(point->y, data);

    // To form the encoding of the point [s]B,

    // copy the least significant bit of the x coordinate to the most significant bit of the final octet.

    data[31] |= (point->x[0] & 1) << 7;

}

void multiply(u8* result_low, u8* result_high, u8 const* a, u8 const* b, u8 n)

{

    // Comba's algorithm

    u32 temp = 0;

    for (u32 i = 0; i < n; i++) {

        for (u32 j = 0; j <= i; j++) {

            temp += (uint16_t)a[j] * b[i - j];

        }

        if (result_low != NULL) {

            result_low[i] = temp & 0xFF;

        }

        temp >>= 8;

    }

    if (result_high != NULL) {

        for (u32 i = n; i < (2 * n); i++) {

            for (u32 j = i + 1 - n; j < n; j++) {

                temp += (uint16_t)a[j] * b[i - j];

            }

            result_high[i - n] = temp & 0xFF;

            temp >>= 8;

        }

    }

}

u8 subtract(u8* result, u8 const* a, u8 const* b, u8 n)

{

    i16 temp = 0;

    // Compute R = A - B

    for (i8 i = 0; i < n; i++) {

        temp += a[i];

        temp -= b[i];

        result[i] = temp & 0xFF;

        temp >>= 8;

    }

    // Return 1 if the result of the subtraction is negative

    return temp & 1;

}

void copy(u8* a, u8 const* b, u32 n)

{

    for (u32 i = 0; i < n; i++) {

        a[i] = b[i];

    }

}

void select(u8* r, u8 const* a, u8 const* b, u8 c, u8 n)

{

    u8 mask = c - 1;

    for (u8 i = 0; i < n; i++) {

        r[i] = (a[i] & mask) | (b[i] & ~mask);

    }

}

void barrett_reduce(u8* result, u8 const* input)

{

    // Barrett reduction b = 2^8 && k = 32

    u8 is_negative;

    u8 u[33];

    u8 v[33];

    multiply(NULL, u, input + 31, Curve25519::BARRETT_REDUCTION_QUOTIENT, 33);

    multiply(v, NULL, u, Curve25519::BASE_POINT_L_ORDER, 33);

    subtract(u, input, v, 33);

    is_negative = subtract(v, u, Curve25519::BASE_POINT_L_ORDER, 33);

    select(u, v, u, is_negative, 33);

    is_negative = subtract(v, u, Curve25519::BASE_POINT_L_ORDER, 33);

    select(u, v, u, is_negative, 33);

    copy(result, u, 32);

}

void add(u8* result, u8 const* a, u8 const* b, u8 n)

{

    // Compute R = A + B

    u16 temp = 0;

    for (u8 i = 0; i < n; i++) {

        temp += a[i];

        temp += b[i];

        result[i] = temp & 0xFF;

        temp >>= 8;

    }

}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.3

u32 decode_point(Ed25519Point* point, u8 const* data)

{

    u32 u[8];

    u32 v[8];

    u32 ret;

    u64 temp = 19;

    // 1. First, interpret the string as an integer in little-endian representation.

    // Bit 255 of this number is the least significant bit of the x-coordinate and denote this value x_0.

    u8 x0 = data[31] >> 7;

    // The y-coordinate is recovered simply by clearing this bit.

    Curve25519::import_state(point->y, data);

    point->y[7] &= 0x7FFFFFFF;

    // Compute U = Y + 19

    for (u32 i = 0; i < 8; i++) {

        temp += point->y[i];

        u[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    // If the resulting value is >= p, decoding fails.

    ret = (u[7] >> 31) & 1;

    // 2. To recover the x-coordinate, the curve equation implies x^2 = (y^2 - 1) / (d y^2 + 1) (mod p).

    // The denominator is always non-zero mod p.

    // Let u = y^2 - 1 and v = d * y^2 + 1

    Curve25519::modular_square(v, point->y);

    Curve25519::modular_subtract_single(u, v, 1);

    Curve25519::modular_multiply(v, v, Curve25519::CURVE_D);

    Curve25519::modular_add_single(v, v, 1);

    // 3. Compute u = sqrt(u / v)

    ret |= Curve25519::modular_square_root(u, u, v);

    // If x = 0, and x_0 = 1, decoding fails.

    ret |= (Curve25519::compare(u, Curve25519::ZERO) ^ 1) & x0;

    // 4.  Finally, use the x_0 bit to select the right square root.

    Curve25519::modular_subtract(v, Curve25519::ZERO, u);

    Curve25519::select(point->x, u, v, (x0 ^ u[0]) & 1);

    Curve25519::set(point->z, 1);

    Curve25519::modular_multiply(point->t, point->x, point->y);

    // Return 0 if the point has been successfully decoded, else 1

    return ret;

}

u8 compare(u8 const* a, u8 const* b, u8 n)

{

    u8 mask = 0;

    for (u32 i = 0; i < n; i++) {

        mask |= a[i] ^ b[i];

    }

    // Return 0 if A = B, else 1

    return ((u8)(mask | (~mask + 1))) >> 7;

}

}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.5

ErrorOr<ByteBuffer> Ed25519::generate_private_key()

{

    // The private key is 32 octets (256 bits, corresponding to b) of

    // cryptographically secure random data.  See [RFC4086] for a discussion

    // about randomness.

    auto buffer = TRY(ByteBuffer::create_uninitialized(key_size()));

    fill_with_random(buffer);

    return buffer;

}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.5

ErrorOr<ByteBuffer> Ed25519::generate_public_key(ReadonlyBytes private_key)

{

    // The 32-byte public key is generated by the following steps.

    // 1. Hash the 32-byte private key using SHA-512, storing the digest in a 64-octet large buffer, denoted h.

    // Only the lower 32 bytes are used for generating the public key.

    auto digest = Crypto::Hash::SHA512::hash(private_key);

    // NOTE: we do these steps in the opposite order (since its easier to modify s)

    // 3. Interpret the buffer as the little-endian integer, forming a secret scalar s.

    u8 s[32] {};

    memcpy(s, digest.data, 32);

    // 2. Prune the buffer:

    // The lowest three bits of the first octet are cleared,

    s[0] &= 0xF8;

    // the highest bit of the last octet is cleared,

    s[31] &= 0x7F;

    // and the second highest bit of the last octet is set.

    s[31] |= 0x40;

    // Perform a fixed-base scalar multiplication [s]B.

    Ed25519Point sb {};

    point_multiply_scalar(&sb, s, &BASE_POINT);

    // 4.  The public key A is the encoding of the point [s]B.

    // First, encode the y-coordinate (in the range 0 <= y < p) as a little-endian string of 32 octets.

    // The most significant bit of the final octet is always zero.

    // To form the encoding of the point [s]B, copy the least significant bit of the x coordinate

    // to the most significant bit of the final octet.  The result is the public key.

    auto public_key = TRY(ByteBuffer::create_uninitialized(key_size()));

    encode_point(&sb, public_key.data());

    return public_key;

}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.6

ErrorOr<ByteBuffer> Ed25519::sign(ReadonlyBytes public_key, ReadonlyBytes private_key, ReadonlyBytes message)

{

    // 1. Hash the private key, 32 octets, using SHA-512.

    // Let h denote the resulting digest.

    auto h = Crypto::Hash::SHA512::hash(private_key);

    // Construct the secret scalar s from the first half of the digest,

    u8 s[32] {};

    memcpy(s, h.data, 32);

    // NOTE: This is done later in step 4, but we can also do it here.

    s[0] &= 0xF8;

    s[31] &= 0x7F;

    s[31] |= 0x40;

    // and the corresponding public key A, as described in the previous section.

    // NOTE: The public key A is taken as input to this function.

    // Let prefix denote the second half of the hash digest, h[32],...,h[63].

    u8 prefix[32] {};

    memcpy(prefix, h.data + 32, 32);

    // 2. Compute SHA-512(dom2(F, C) || p || PH(M)), where M is the message to be signed.

    Crypto::Hash::SHA512 hash;

    // NOTE: dom2(F, C) is a blank octet string when signing or verifying Ed25519

    hash.update(prefix, 32);

    // NOTE: PH(M) = M

    hash.update(message.data(), message.size());

    // Interpret the 64-octet digest as a little-endian integer r.

    // For efficiency, do this by first reducing r modulo L, the group order of B.

    auto digest = hash.digest();

    u8 r[32] {};

    barrett_reduce(r, digest.data);

    // 3.  Compute the point [r]B.

    Ed25519Point rb {};

    point_multiply_scalar(&rb, r, &BASE_POINT);

    auto R = TRY(ByteBuffer::create_uninitialized(32));

    // Let the string R be the encoding of this point

    encode_point(&rb, R.data());

    // 4. Compute SHA512(dom2(F, C) || R || A || PH(M)),

    // NOTE: We can reuse hash here, since digest() calls reset()

    // NOTE: dom2(F, C) is a blank octet string when signing or verifying Ed25519

    hash.update(R.data(), R.size());

    // NOTE: A == public_key

    hash.update(public_key.data(), public_key.size());

    // NOTE: PH(M) = M

    hash.update(message.data(), message.size());

    digest = hash.digest();

    // and interpret the 64-octet digest as a little-endian integer k.

    u8 k[64] {};

    memcpy(k, digest.data, 64);

    // 5. Compute S = (r + k * s) mod L. For efficiency, again reduce k modulo L first.

    barrett_reduce(prefix, k);

    multiply(k, k + 32, prefix, s, 32);

    barrett_reduce(prefix, k);

    add(s, prefix, r, 32);

    // modular reduction

    auto reduced_s = TRY(ByteBuffer::create_uninitialized(32));

    auto is_negative = subtract(prefix, s, Curve25519::BASE_POINT_L_ORDER, 32);

    select(reduced_s.data(), prefix, s, is_negative, 32);

    // 6.  Form the signature of the concatenation of R (32 octets) and the little-endian encoding of S

    // (32 octets; the three most significant bits of the final octet are always zero).

    auto signature = TRY(ByteBuffer::copy(R));

    signature.append(reduced_s);

    return signature;

}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.7

bool Ed25519::verify(ReadonlyBytes public_key, ReadonlyBytes signature, ReadonlyBytes message)

{

    auto not_valid = false;

    // 1. To verify a signature on a message M using public key A,

    // with F being 0 for Ed25519ctx, 1 for Ed25519ph, and if Ed25519ctx or Ed25519ph is being used, C being the context

    // first split the signature into two 32-octet halves.

    // If any of the decodings fail (including S being out of range), the signature is invalid.

    // NOTE: We dont care about F, since we dont implement Ed25519ctx or Ed25519PH

    // NOTE: C is the internal state, so its not a parameter

    if (public_key.size() != key_size() || signature.size() != signature_size())

        return false;

    auto half_signature_size = signature_size() / 2;

    // Decode the first half as a point R

    u8 r[32] {};

    memcpy(r, signature.data(), half_signature_size);

    // and the second half as an integer S, in the range  0 <= s < L.

    u8 s[32] {};

    memcpy(s, signature.data() + half_signature_size, half_signature_size);

    // NOTE: Ed25519 and Ed448 signatures are not malleable due to the verification check that decoded S is smaller than l.

    // Without this check, one can add a multiple of l into a scalar part and still pass signature verification,

    // resulting in malleable signatures.

    u8 p[32] {};

    auto is_negative = subtract(p, s, Curve25519::BASE_POINT_L_ORDER, half_signature_size);

    not_valid |= 1 ^ is_negative;

    // Decode the public key A as point A'.

    Ed25519Point ka {};

    not_valid |= decode_point(&ka, public_key.data());

    // 2. Compute SHA512(dom2(F, C) || R || A || PH(M)), and interpret the 64-octet digest as a little-endian integer k.

    Crypto::Hash::SHA512 hash;

    // NOTE: dom2(F, C) is a blank octet string when signing or verifying Ed25519

    hash.update(r, half_signature_size);

    // NOTE: A == public_key

    hash.update(public_key.data(), key_size());

    // NOTE: PH(M) = M

    hash.update(message.data(), message.size());

    auto digest = hash.digest();

    auto k = digest.data;

    // 3.  Check the group equation [8][S]B = [8]R + [8][k]A'.

    // It's sufficient, but not required, to instead check [S]B = R + [k]A'.

    // NOTE: For efficiency, do this by first reducing k modulo L.

    barrett_reduce(k, k);

    // NOTE: We check [S]B - [k]A' == R

    Curve25519::modular_subtract(ka.x, Curve25519::ZERO, ka.x);

    Curve25519::modular_subtract(ka.t, Curve25519::ZERO, ka.t);

    Ed25519Point sb {};

    point_multiply_scalar(&sb, s, &BASE_POINT);

    point_multiply_scalar(&ka, k, &ka);

    point_add(&ka, &sb, &ka);

    encode_point(&ka, p);

    not_valid |= compare(p, r, half_signature_size);

    return !not_valid;

}

}