/***********************************************************************
 * Distributed under the MIT software license, see the accompanying    *
 * file COPYING or https://www.opensource.org/licenses/mit-license.php.*
 ***********************************************************************/

#ifndef SECP256K1_MODULE_CHILLDKG_UTIL_IMPL_H
#define SECP256K1_MODULE_CHILLDKG_UTIL_IMPL_H

#include <string.h>

#include "../../../include/secp256k1.h"

#include "util.h"
#include "../../ecmult.h"
#include "../../ecmult_const.h"
#include "../../util.h"

static void secp256k1_chilldkg_sha256_tagged_params_hash(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP DKG/params_hash", sizeof("BIP DKG/params_hash") - 1);
}

static void secp256k1_chilldkg_sha256_tagged_encpedpop_seed(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP DKG/encpedpop seed", sizeof("BIP DKG/encpedpop seed") - 1);
}

static void secp256k1_chilldkg_sha256_tagged_simplpedpop_aux(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP DKG/simplpedpop aux", sizeof("BIP DKG/simplpedpop aux") - 1);
}

static void secp256k1_chilldkg_sha256_tagged_encpedpop_secnonce(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP DKG/encpedpop secnonce", sizeof("BIP DKG/encpedpop secnonce") - 1);
}

static void secp256k1_chilldkg_sha256_tagged_encpedpop_ecdh(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP DKG/encpedpop ecdh", sizeof("BIP DKG/encpedpop ecdh") - 1);
}

static void secp256k1_chilldkg_sha256_tagged_self_pad(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP DKG/encaps_multi self_pad", sizeof("BIP DKG/encaps_multi self_pad") - 1);
}

static void secp256k1_chilldkg_sha256_tagged_vss_coeffs(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"BIP DKG/vss coeffs", sizeof("BIP DKG/vss coeffs") - 1);
}

static void secp256k1_chilldkg_sha256_tagged_taptweak(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha) {
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, (const unsigned char *)"TapTweak", sizeof("TapTweak") - 1);
}

static void secp256k1_chilldkg_point_save(unsigned char *out33, const secp256k1_ge *p) {
    if (secp256k1_ge_is_infinity(p)) {
        memset(out33, 0, 33);
        return;
    }
    {
        secp256k1_ge tmp = *p;
        /* Serialization operates on public data (commitments, pubnonces, host
         * public keys), so variable-time normalization is fine. */
        secp256k1_fe_normalize_var(&tmp.x);
        secp256k1_fe_normalize_var(&tmp.y);
        out33[0] = secp256k1_fe_is_odd(&tmp.y) ? 0x03 : 0x02;
        secp256k1_fe_get_b32(&out33[1], &tmp.x);
    }
}

static int secp256k1_chilldkg_point_load(secp256k1_ge *p, const unsigned char *in33) {
    static const unsigned char zeros33[33] = { 0 };
    secp256k1_fe x;

    /* Parsed data comes from protocol messages, i.e., it is public. */
    if (secp256k1_memcmp_var(in33, zeros33, 33) == 0) {
        secp256k1_ge_set_infinity(p);
        return 1;
    }
    if (in33[0] != 0x02 && in33[0] != 0x03) {
        secp256k1_ge_set_infinity(p);
        return 0;
    }
    if (!secp256k1_fe_set_b32_limit(&x, &in33[1])) {
        secp256k1_ge_set_infinity(p);
        return 0;
    }
    if (!secp256k1_ge_set_xo_var(p, &x, in33[0] == 0x03)) {
        secp256k1_ge_set_infinity(p);
        return 0;
    }
    return 1;
}

static void secp256k1_chilldkg_xonly_save(unsigned char *out32, const secp256k1_ge *p) {
    secp256k1_ge tmp = *p;
    VERIFY_CHECK(!secp256k1_ge_is_infinity(p));
    secp256k1_fe_normalize_var(&tmp.x);
    secp256k1_fe_get_b32(out32, &tmp.x);
}

static int secp256k1_chilldkg_xonly_load(secp256k1_ge *p, const unsigned char *in32) {
    secp256k1_fe x;
    if (!secp256k1_fe_set_b32_limit(&x, in32)) {
        return 0;
    }
    return secp256k1_ge_set_xo_var(p, &x, 0);
}

static void secp256k1_chilldkg_pad33(unsigned char *out33, const char *str) {
    size_t len = strlen(str);
    /* Every call site passes a string literal of the module, so this cannot
     * trigger. The clamp must not sit inside VERIFY_CHECK, which is compiled
     * out in noverify builds: an over-long tag would overflow out33 and make
     * the memset length below underflow to a huge value. */
    if (len > 33) {
        VERIFY_CHECK(0);
        len = 33;
    }
    memcpy(out33, str, len);
    memset(out33 + len, 0, 33 - len);
}

/* Initializes sha with the BIP 340 tagged hash tag tag_prefix || subtag,
 * e.g., "BIP DKG/pop message" || "/nonce". Tag strings are public constants. */
static void secp256k1_chilldkg_schnorrsig_sha256_tagged(const secp256k1_hash_ctx *hash_ctx, secp256k1_sha256 *sha, const char *tag_prefix, const char *subtag) {
    unsigned char tag[64];
    size_t prefix_len = strlen(tag_prefix);
    size_t subtag_len = strlen(subtag);

    /* The longest tag the module builds is "BIP DKG/pop message" ||
     * "/challenge", 29 bytes. As in secp256k1_chilldkg_pad33, the bound is
     * enforced outside VERIFY_CHECK so that a future over-long tag cannot
     * overflow tag[] in a noverify build. Clamping rather than returning
     * early keeps sha initialized for the caller; a truncated tag changes
     * every hash the module computes, so the test vectors fail loudly. */
    if (prefix_len > sizeof(tag)) {
        VERIFY_CHECK(0);
        prefix_len = sizeof(tag);
    }
    if (subtag_len > sizeof(tag) - prefix_len) {
        VERIFY_CHECK(0);
        subtag_len = sizeof(tag) - prefix_len;
    }
    memcpy(tag, tag_prefix, prefix_len);
    memcpy(tag + prefix_len, subtag, subtag_len);
    secp256k1_sha256_initialize_tagged(hash_ctx, sha, tag, prefix_len + subtag_len);
}

/* BIP 340 nonce derivation with parameterized tag prefix, mirroring
 * schnorr_sign in secp256k1lab/bip340.py:
 *   t = seckey32 XOR TH(tag_prefix || "/aux", aux_rand32)
 *   nonce32 = TH(tag_prefix || "/nonce", t || xonly_pk32 || msg) */
static void secp256k1_chilldkg_schnorrsig_nonce(const secp256k1_hash_ctx *hash_ctx, unsigned char *nonce32, const unsigned char *msg, size_t msglen, const unsigned char *seckey32, const unsigned char *xonly_pk32, const unsigned char *aux_rand32, const char *tag_prefix) {
    secp256k1_sha256 sha;
    unsigned char masked_key[32];
    unsigned char rand[32];
    int i;

    secp256k1_chilldkg_schnorrsig_sha256_tagged(hash_ctx, &sha, tag_prefix, "/aux");
    secp256k1_sha256_write(hash_ctx, &sha, aux_rand32, 32);
    secp256k1_sha256_finalize(hash_ctx, &sha, rand);
    for (i = 0; i < 32; i++) {
        masked_key[i] = seckey32[i] ^ rand[i];
    }
    secp256k1_chilldkg_schnorrsig_sha256_tagged(hash_ctx, &sha, tag_prefix, "/nonce");
    secp256k1_sha256_write(hash_ctx, &sha, masked_key, 32);
    secp256k1_sha256_write(hash_ctx, &sha, xonly_pk32, 32);
    secp256k1_sha256_write(hash_ctx, &sha, msg, msglen);
    secp256k1_sha256_finalize(hash_ctx, &sha, nonce32);
    secp256k1_sha256_clear(&sha);
    secp256k1_memclear_explicit(masked_key, sizeof(masked_key));
    secp256k1_memclear_explicit(rand, sizeof(rand));
}

/* BIP 340 challenge hash with parameterized tag prefix:
 *   e = TH(tag_prefix || "/challenge", r32 || pubkey32 || msg) mod n
 * The reduction modulo the group order matches the reference, which reduces
 * the hash output with int_from_bytes(...) % GE.ORDER. */
static void secp256k1_chilldkg_schnorrsig_challenge(const secp256k1_hash_ctx *hash_ctx, secp256k1_scalar *e, const unsigned char *r32, const unsigned char *msg, size_t msglen, const unsigned char *pubkey32, const char *tag_prefix) {
    unsigned char buf[32];
    secp256k1_sha256 sha;

    secp256k1_chilldkg_schnorrsig_sha256_tagged(hash_ctx, &sha, tag_prefix, "/challenge");
    secp256k1_sha256_write(hash_ctx, &sha, r32, 32);
    secp256k1_sha256_write(hash_ctx, &sha, pubkey32, 32);
    secp256k1_sha256_write(hash_ctx, &sha, msg, msglen);
    secp256k1_sha256_finalize(hash_ctx, &sha, buf);
    secp256k1_scalar_set_b32(e, buf, NULL);
}

/* Mirrors schnorr_sign in secp256k1lab/bip340.py. The structure follows
 * secp256k1_schnorrsig_sign_internal, which cannot be reused directly because
 * it hardcodes the "BIP0340" tags. */
static int secp256k1_chilldkg_schnorrsig_sign(const secp256k1_context *ctx, unsigned char *sig64, const unsigned char *msg, size_t msglen, const unsigned char *seckey32, const unsigned char *aux_rand32, const char *tag_prefix) {
    const secp256k1_hash_ctx *hash_ctx;
    secp256k1_scalar sk;
    secp256k1_scalar e;
    secp256k1_scalar k;
    secp256k1_ge pk;
    secp256k1_ge r;
    unsigned char nonce32[32] = { 0 };
    unsigned char pk32[32];
    unsigned char seckey[32];
    int overflow;
    int ret = 1;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(secp256k1_ecmult_gen_context_is_built(&ctx->ecmult_gen_ctx));
    ARG_CHECK(sig64 != NULL);
    ARG_CHECK(msg != NULL || msglen == 0);
    ARG_CHECK(seckey32 != NULL);
    ARG_CHECK(aux_rand32 != NULL);
    ARG_CHECK(tag_prefix != NULL);
    hash_ctx = secp256k1_get_hash_context(ctx);

    secp256k1_scalar_set_b32(&sk, seckey32, &overflow);
    overflow |= secp256k1_scalar_is_zero(&sk);
    /* Branching on the validity of the secret key is fine: whether the
     * caller's secret key is in range 1..n-1 is not secret. */
    secp256k1_declassify(ctx, &overflow, sizeof(overflow));
    if (overflow) {
        memset(sig64, 0, 64);
        secp256k1_scalar_clear(&sk);
        return 0;
    }

    secp256k1_ecmult_gen_ge(&ctx->ecmult_gen_ctx, &pk, &sk);
    /* The public key is not secret, so variable-time normalization and
     * branching on its y parity are fine. */
    secp256k1_declassify(ctx, &pk, sizeof(pk));
    secp256k1_fe_normalize_var(&pk.x);
    secp256k1_fe_normalize_var(&pk.y);
    if (secp256k1_fe_is_odd(&pk.y)) {
        secp256k1_scalar_negate(&sk, &sk);
    }
    secp256k1_scalar_get_b32(seckey, &sk);
    secp256k1_fe_get_b32(pk32, &pk.x);

    secp256k1_chilldkg_schnorrsig_nonce(hash_ctx, nonce32, msg, msglen, seckey, pk32, aux_rand32, tag_prefix);
    /* The reference reduces the nonce hash modulo the group order. */
    secp256k1_scalar_set_b32(&k, nonce32, NULL);
    ret &= !secp256k1_scalar_is_zero(&k);
    secp256k1_scalar_cmov(&k, &secp256k1_scalar_one, !ret);

    secp256k1_ecmult_gen_ge(&ctx->ecmult_gen_ctx, &r, &k);
    /* We declassify r to allow using it as a branch point. This is fine
     * because r is not a secret. */
    secp256k1_declassify(ctx, &r, sizeof(r));
    secp256k1_fe_normalize_var(&r.y);
    if (secp256k1_fe_is_odd(&r.y)) {
        secp256k1_scalar_negate(&k, &k);
    }
    secp256k1_fe_normalize_var(&r.x);
    secp256k1_fe_get_b32(&sig64[0], &r.x);

    secp256k1_chilldkg_schnorrsig_challenge(hash_ctx, &e, &sig64[0], msg, msglen, pk32, tag_prefix);
    secp256k1_scalar_mul(&e, &e, &sk);
    secp256k1_scalar_add(&e, &e, &k);
    secp256k1_scalar_get_b32(&sig64[32], &e);

    secp256k1_memczero(sig64, 64, !ret);
    secp256k1_scalar_clear(&k);
    secp256k1_scalar_clear(&sk);
    secp256k1_scalar_clear(&e);
    secp256k1_ge_clear(&pk);
    secp256k1_ge_clear(&r);
    secp256k1_memclear_explicit(seckey, sizeof(seckey));
    secp256k1_memclear_explicit(nonce32, sizeof(nonce32));

    /* Branching on the return value only leaks whether the derived nonce is
     * zero, which happens with negligible probability (the case of an invalid
     * secret key is declassified above). */
    secp256k1_declassify(ctx, &ret, sizeof(ret));
    return ret;
}

/* Mirrors schnorr_verify in secp256k1lab/bip340.py. The structure follows
 * secp256k1_schnorrsig_verify, which cannot be reused directly because it
 * hardcodes the "BIP0340/challenge" tag. */
static int secp256k1_chilldkg_schnorrsig_verify(const secp256k1_context *ctx, const unsigned char *sig64, const unsigned char *msg, size_t msglen, const unsigned char *pubkey32, const char *tag_prefix) {
    secp256k1_scalar s;
    secp256k1_scalar e;
    secp256k1_gej rj;
    secp256k1_ge pk;
    secp256k1_gej pkj;
    secp256k1_fe rx;
    secp256k1_ge r;
    unsigned char buf[32];
    int overflow;

    VERIFY_CHECK(ctx != NULL);
    ARG_CHECK(sig64 != NULL);
    ARG_CHECK(msg != NULL || msglen == 0);
    ARG_CHECK(pubkey32 != NULL);
    ARG_CHECK(tag_prefix != NULL);

    if (!secp256k1_fe_set_b32_limit(&rx, &sig64[0])) {
        return 0;
    }

    secp256k1_scalar_set_b32(&s, &sig64[32], &overflow);
    if (overflow) {
        return 0;
    }

    if (!secp256k1_chilldkg_xonly_load(&pk, pubkey32)) {
        return 0;
    }

    /* Compute e. */
    secp256k1_fe_get_b32(buf, &pk.x);
    secp256k1_chilldkg_schnorrsig_challenge(secp256k1_get_hash_context(ctx), &e, &sig64[0], msg, msglen, buf, tag_prefix);

    /* Compute rj = s*G + (-e)*pkj */
    secp256k1_scalar_negate(&e, &e);
    secp256k1_gej_set_ge(&pkj, &pk);
    secp256k1_ecmult(&rj, &pkj, &e, &s);

    secp256k1_ge_set_gej_var(&r, &rj);
    if (secp256k1_ge_is_infinity(&r)) {
        return 0;
    }

    secp256k1_fe_normalize_var(&r.y);
    return !secp256k1_fe_is_odd(&r.y) &&
           secp256k1_fe_equal(&rx, &r.x);
}

static void secp256k1_chilldkg_encpedpop_ecdh(const secp256k1_context *ctx, secp256k1_scalar *out, const secp256k1_scalar *seckey, const secp256k1_ge *their_point, const unsigned char *my_pubkey33, const unsigned char *their_pubkey33, const unsigned char *context, size_t context_len, int sending) {
    const secp256k1_hash_ctx *hash_ctx = secp256k1_get_hash_context(ctx);
    secp256k1_gej resj;
    secp256k1_ge res;
    unsigned char x[32];
    unsigned char y[32];
    unsigned char shared[32];
    unsigned char hash32[32];
    secp256k1_sha256 sha;

    VERIFY_CHECK(ctx != NULL);
    VERIFY_CHECK(!secp256k1_scalar_is_zero(seckey));
    VERIFY_CHECK(!secp256k1_ge_is_infinity(their_point));

    /* libsecp256k1-style ECDH: SHA256 of the compressed shared point. */
    secp256k1_ecmult_const(&resj, their_point, seckey);
    secp256k1_ge_set_gej(&res, &resj);
    /* The result cannot be the point at infinity: their_point is not
     * infinity, seckey is nonzero, and the group has prime order. */
    VERIFY_CHECK(!secp256k1_ge_is_infinity(&res));
    secp256k1_fe_normalize(&res.x);
    secp256k1_fe_normalize(&res.y);
    secp256k1_fe_get_b32(x, &res.x);
    secp256k1_fe_get_b32(y, &res.y);
    /* This hash function always succeeds; call it unconditionally (it must
     * not sit inside VERIFY_CHECK, which is compiled out in noverify
     * builds). */
    if (!ecdh_hash_function_sha256_impl(hash_ctx, shared, x, y, NULL)) {
        VERIFY_CHECK(0);
        memset(shared, 0, sizeof(shared));
    }

    secp256k1_chilldkg_sha256_tagged_encpedpop_ecdh(hash_ctx, &sha);
    secp256k1_sha256_write(hash_ctx, &sha, shared, 32);
    /* The sender's pubnonce always comes first in the hash input. */
    if (sending) {
        secp256k1_sha256_write(hash_ctx, &sha, my_pubkey33, 33);
        secp256k1_sha256_write(hash_ctx, &sha, their_pubkey33, 33);
    } else {
        secp256k1_sha256_write(hash_ctx, &sha, their_pubkey33, 33);
        secp256k1_sha256_write(hash_ctx, &sha, my_pubkey33, 33);
    }
    secp256k1_sha256_write(hash_ctx, &sha, context, context_len);
    secp256k1_sha256_finalize(hash_ctx, &sha, hash32);
    secp256k1_sha256_clear(&sha);

    /* Pads are reduced modulo the group order (from_bytes_wrapping in the
     * reference). */
    secp256k1_scalar_set_b32(out, hash32, NULL);

    secp256k1_memclear_explicit(x, sizeof(x));
    secp256k1_memclear_explicit(y, sizeof(y));
    secp256k1_memclear_explicit(shared, sizeof(shared));
    secp256k1_memclear_explicit(hash32, sizeof(hash32));
    secp256k1_ge_clear(&res);
    secp256k1_gej_clear(&resj);
}

static void secp256k1_chilldkg_encpedpop_self_pad(const secp256k1_hash_ctx *hash_ctx, secp256k1_scalar *out, const unsigned char *symkey32, const unsigned char *nonce33, const unsigned char *context, size_t context_len) {
    unsigned char hash32[32];
    secp256k1_sha256 sha;

    secp256k1_chilldkg_sha256_tagged_self_pad(hash_ctx, &sha);
    secp256k1_sha256_write(hash_ctx, &sha, symkey32, 32);
    secp256k1_sha256_write(hash_ctx, &sha, nonce33, 33);
    secp256k1_sha256_write(hash_ctx, &sha, context, context_len);
    secp256k1_sha256_finalize(hash_ctx, &sha, hash32);
    secp256k1_sha256_clear(&sha);

    /* from_bytes_wrapping in the reference. */
    secp256k1_scalar_set_b32(out, hash32, NULL);
    secp256k1_memclear_explicit(hash32, sizeof(hash32));
}

static void secp256k1_chilldkg_params_hash_internal(const secp256k1_hash_ctx *hash_ctx, unsigned char *out32, const unsigned char *hostpubkeys33, size_t n, uint32_t t) {
    unsigned char buf[4];
    secp256k1_sha256 sha;

    secp256k1_chilldkg_sha256_tagged_params_hash(hash_ctx, &sha);
    secp256k1_write_be32(buf, t);
    secp256k1_sha256_write(hash_ctx, &sha, buf, sizeof(buf));
    secp256k1_sha256_write(hash_ctx, &sha, hostpubkeys33, 33 * n);
    secp256k1_sha256_finalize(hash_ctx, &sha, out32);
    secp256k1_sha256_clear(&sha);
}

#endif