/*************************************************************************
 * To the extent possible under law, the author(s) have dedicated all    *
 * copyright and related and neighboring rights to the software in this  *
 * file to the public domain worldwide. This software is distributed     *
 * without any warranty. For the CC0 Public Domain Dedication, see       *
 * EXAMPLES_COPYING or https://creativecommons.org/publicdomain/zero/1.0 *
 *************************************************************************/

/** This file demonstrates how to use the FROST enrollment module: first to
 *  grow a 2-of-3 group into a 2-of-4 one without re-running key generation,
 *  then to repair a participant's lost share. Additionally, see the
 *  documentation in include/secp256k1_frost_enrollment.h and
 *  src/modules/frost_enrollment/frost_enrollment.md.
 *
 *  The example runs all roles (trusted dealer, helpers, the party receiving
 *  the share, coordinator) in a single process. In a real deployment these
 *  roles are performed by different parties, and the enrollment shares and
 *  the aggregated values they produce MUST travel over confidential and
 *  authenticated channels: they are additive shares of real secret shares.
 *
 *  Two things this example demonstrates that are easy to get wrong:
 *
 *  1. The verification flow is only non-circular if the threshold public key
 *     comes from somewhere the helpers do not control. Here the enrollee gets
 *     it straight from the dealer step, which stands in for real
 *     authentication, and validates the helpers' public shares against it
 *     before deriving anything from them.
 *  2. Nothing in the protocol decides WHO may be enrolled. That is a
 *     caller-side precondition; see the note before enroll() below.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <secp256k1.h>
#include <secp256k1_extrakeys.h>
#include <secp256k1_frost.h>
#include <secp256k1_frost_enrollment.h>
#include <secp256k1_schnorrsig.h>

#include "examples_util.h"

/* Number of participants n before the enrollment */
#define N_PARTICIPANTS 3
/* Number of participants after the enrollment */
#define N_PARTICIPANTS_AFTER (N_PARTICIPANTS + 1)
/* Threshold t */
#define THRESHOLD 2
/* Number of helpers u taking part in the enrollment run. Any u with
 * THRESHOLD <= u <= N_PARTICIPANTS works; the result does not depend on the
 * choice. */
#define N_HELPERS 2
/* The identifier the new participant will hold. In enrollment mode this must
 * be exactly the current participant count. */
#define NEW_ID N_PARTICIPANTS
/* The signers of the demonstration signing session: participant 2 and the
 * newly enrolled participant. */
#define N_SIGNERS 2

/* What one helper holds during a run. */
struct helper {
    uint32_t id;
    /* The helper's own secret share of the group key. Never leaves the
     * helper. */
    unsigned char secshare[32];
    /* Round 1.1 output, aligned with the helper id array: entry j goes to the
     * helper at position j. The entry at this helper's own position is kept
     * locally. These are secret. */
    unsigned char shares[N_HELPERS * 32];
    /* The parameters hash this helper computed, sent alongside every share it
     * distributes and to the enrollee. Public. */
    unsigned char params_hash[32];
    /* Round 1.2 output, sent to the enrollee. Secret. */
    unsigned char sigma[32];
};

/* Run the trusted dealer key generation for the initial 2-of-3 group.
 *
 * WARNING: The trusted dealer knows the threshold secret key and all secret
 * shares, and must erase them after distributing the shares. A distributed key
 * generation protocol avoids a trusted dealer entirely; see the chilldkg
 * module. */
static int trusted_dealer_keygen(const secp256k1_context* ctx, unsigned char *threshold_seckey, unsigned char *secshares, secp256k1_pubkey *thresh_pk, secp256k1_pubkey *pubshares) {
    if (!fill_random(threshold_seckey, 32)) {
        printf("Failed to generate randomness\n");
        return 0;
    }
    if (!secp256k1_frost_trusted_dealer_keygen(ctx, secshares, thresh_pk, pubshares, N_PARTICIPANTS, THRESHOLD, threshold_seckey)) {
        return 0;
    }
    return secp256k1_frost_threshold_info_validate(ctx, thresh_pk, pubshares, N_PARTICIPANTS, THRESHOLD);
}

/* Run one enrollment or repair, all three rounds.
 *
 * `helper_pubshares` and `helpers` are aligned with `ids`. On success the new
 * secret share is written to `new_secshare` and its public counterpart to
 * `new_pubshare`.
 *
 * PRECONDITION THE LIBRARY CANNOT ENFORCE: the helpers must already have
 * agreed, out of band, that this party is entitled to a share at `new_id`.
 * The protocol has no authorization step of its own: anyone who convinces t
 * helpers to run it walks away with a valid share, and in repair mode that is
 * an existing participant's actual share. Authenticated channels establish who
 * is speaking, not that the group approved the request. */
static int enroll(const secp256k1_context* ctx, unsigned char *new_secshare, secp256k1_pubkey *new_pubshare, struct helper *helpers, const uint32_t *ids, const secp256k1_pubkey *helper_pubshares, const secp256k1_pubkey *thresh_pk, uint32_t new_id) {
    unsigned char sigmas[N_HELPERS * 32];
    unsigned char enrollee_params_hash[32];
    int i, j;

    /* --- Round 1.1 ------------------------------------------------------
     * Every helper splits its Lagrange-weighted share into one additive share
     * per helper, keeps its own and sends the rest out. */
    for (i = 0; i < N_HELPERS; i++) {
        unsigned char session_secrand[32];

        /* Fresh randomness for every run. Reusing it across runs leaks share
         * information. shares_gen wipes it before returning. */
        if (!fill_random(session_secrand, sizeof(session_secrand))) {
            printf("Failed to generate randomness\n");
            return 0;
        }
        if (!secp256k1_frost_enrollment_shares_gen(ctx, helpers[i].shares, helpers[i].params_hash, session_secrand, helpers[i].secshare, thresh_pk, ids, N_HELPERS, helpers[i].id, new_id, N_PARTICIPANTS, THRESHOLD)) {
            return 0;
        }
        secure_erase(session_secrand, sizeof(session_secrand));
    }

    /* --- Round 1.2 ------------------------------------------------------
     * Helper j receives entry j of every other helper's buffer, over a
     * confidential and authenticated channel, together with that helper's
     * parameters hash. It recomputes the hash itself and compares. */
    for (j = 0; j < N_HELPERS; j++) {
        unsigned char all_shares[N_HELPERS * 32];
        unsigned char received_hashes[N_HELPERS * 32];
        uint32_t mismatch_id;

        /* The two buffers take deliberately opposite conventions at this
         * helper's own position: the share kept locally IS read, while the
         * hash slot is never read and stays zero. The own hash is
         * recomputed, never taken from a buffer -- which is what makes this
         * a recomputation check rather than a comparison between two strings
         * the caller supplied. */
        memset(received_hashes, 0, sizeof(received_hashes));
        for (i = 0; i < N_HELPERS; i++) {
            memcpy(&all_shares[32 * i], &helpers[i].shares[32 * j], 32);
            if (i != j) {
                memcpy(&received_hashes[32 * i], helpers[i].params_hash, 32);
            }
        }
        if (!secp256k1_frost_enrollment_share_agg(ctx, helpers[j].sigma, &mismatch_id, all_shares, received_hashes, thresh_pk, ids, N_HELPERS, helpers[j].id, new_id, N_PARTICIPANTS, THRESHOLD)) {
            if (mismatch_id != UINT32_MAX) {
                printf("\nHelper %u disagrees about the enrollment parameters\n", mismatch_id);
            }
            return 0;
        }
        secure_erase(all_shares, sizeof(all_shares));
    }

    /* --- Round 2 --------------------------------------------------------
     * The enrollee receives one value from each helper, again confidentially
     * and authenticated, plus the parameters hash.
     *
     * The expected public share is derived from the helpers' public shares --
     * which the caller has already validated against an independently
     * authenticated threshold public key (see main). Without that step both
     * checks below would be circular: t colluding helpers could present a
     * consistent but fabricated polynomial and every check would pass on a
     * worthless share. */
    for (i = 0; i < N_HELPERS; i++) {
        memcpy(&sigmas[32 * i], helpers[i].sigma, 32);
    }
    memcpy(enrollee_params_hash, helpers[0].params_hash, 32);

    if (!secp256k1_frost_enrollment_pubshare_derive(ctx, new_pubshare, helper_pubshares, ids, N_HELPERS, new_id, N_PARTICIPANTS, THRESHOLD)) {
        return 0;
    }
    if (!secp256k1_frost_enrollment_secshare_gen(ctx, new_secshare, sigmas, thresh_pk, ids, N_HELPERS, new_id, N_PARTICIPANTS, THRESHOLD, enrollee_params_hash, new_pubshare)) {
        return 0;
    }
    secure_erase(sigmas, sizeof(sigmas));
    for (i = 0; i < N_HELPERS; i++) {
        secure_erase(helpers[i].shares, sizeof(helpers[i].shares));
        secure_erase(helpers[i].sigma, sizeof(helpers[i].sigma));
    }
    return 1;
}

/* Produce a BIP340 signature with the given signer set and verify it against
 * the (untweaked) threshold public key. `secshares` and `pubshares` are
 * aligned with `ids`. `n_participants` is the group size AFTER the enrollment,
 * which every participant must have updated consistently. */
static int sign_and_verify(const secp256k1_context* ctx, const uint32_t *ids, unsigned char (*secshares)[32], const secp256k1_pubkey *pubshares, const secp256k1_pubkey *thresh_pk, size_t n_participants, const unsigned char *msg, size_t msglen) {
    secp256k1_frost_tweak_cache cache;
    secp256k1_frost_secnonce secnonces[N_SIGNERS];
    secp256k1_frost_pubnonce pubnonces[N_SIGNERS];
    const secp256k1_frost_pubnonce *pubnonce_ptrs[N_SIGNERS];
    secp256k1_frost_partial_sig partial_sigs[N_SIGNERS];
    const secp256k1_frost_partial_sig *partial_sig_ptrs[N_SIGNERS];
    secp256k1_frost_aggnonce aggnonce;
    secp256k1_frost_session session;
    secp256k1_xonly_pubkey tweaked_pk;
    unsigned char tweaked_pk32[32];
    unsigned char sig[64];
    int i;
    int ret = 0;

    if (!secp256k1_frost_tweak_cache_init(ctx, &cache, thresh_pk)) {
        return 0;
    }
    /* No tweaks are applied here, so the "tweaked" key is the threshold public
     * key itself in its x-only encoding. */
    if (!secp256k1_frost_tweaked_pubkey_get(ctx, &tweaked_pk, &cache)
            || !secp256k1_xonly_pubkey_serialize(ctx, tweaked_pk32, &tweaked_pk)) {
        return 0;
    }

    for (i = 0; i < N_SIGNERS; i++) {
        unsigned char session_secrand[32];
        if (!fill_random(session_secrand, sizeof(session_secrand))) {
            printf("Failed to generate randomness\n");
            return 0;
        }
        if (!secp256k1_frost_nonce_gen(ctx, &secnonces[i], &pubnonces[i], session_secrand, secshares[i], &pubshares[i], tweaked_pk32, msg, msglen, NULL, 0)) {
            return 0;
        }
        secure_erase(session_secrand, sizeof(session_secrand));
        pubnonce_ptrs[i] = &pubnonces[i];
        partial_sig_ptrs[i] = &partial_sigs[i];
    }
    if (!secp256k1_frost_nonce_agg(ctx, &aggnonce, NULL, pubnonce_ptrs, N_SIGNERS)) {
        return 0;
    }
    if (!secp256k1_frost_session_init(ctx, &session, &aggnonce, ids, pubshares, N_SIGNERS, n_participants, THRESHOLD, &cache, msg, msglen)) {
        return 0;
    }
    for (i = 0; i < N_SIGNERS; i++) {
        if (!secp256k1_frost_sign(ctx, &partial_sigs[i], &secnonces[i], secshares[i], &session, ids, pubshares, N_SIGNERS, ids[i])) {
            goto cleanup;
        }
        if (!secp256k1_frost_partial_sig_verify(ctx, &partial_sigs[i], &pubnonces[i], &pubshares[i], &session, ids, N_SIGNERS, (size_t)i)) {
            goto cleanup;
        }
    }
    if (!secp256k1_frost_partial_sig_agg(ctx, sig, NULL, &session, partial_sig_ptrs, N_SIGNERS)) {
        goto cleanup;
    }
    ret = secp256k1_schnorrsig_verify(ctx, sig, msg, msglen, &tweaked_pk);

cleanup:
    for (i = 0; i < N_SIGNERS; i++) {
        secure_erase(&secnonces[i], sizeof(secnonces[i]));
    }
    return ret;
}

int main(void) {
    secp256k1_context* ctx;
    /* Key material of the initial 2-of-3 group. The dealer erases all of it
     * at the end of this function. */
    unsigned char threshold_seckey[32];
    unsigned char secshares[N_PARTICIPANTS * 32];
    secp256k1_pubkey thresh_pk;
    secp256k1_pubkey pubshares[N_PARTICIPANTS_AFTER];

    /* The enrollment run for the new participant 3. */
    struct helper helpers[N_HELPERS];
    uint32_t helper_ids[N_HELPERS] = { 0, 1 };
    secp256k1_pubkey helper_pubshares[N_HELPERS];
    unsigned char new_secshare[32];
    secp256k1_pubkey new_pubshare;

    /* The repair run for participant 1. */
    struct helper repair_helpers[N_HELPERS];
    uint32_t repair_ids[N_HELPERS] = { 0, 2 };
    secp256k1_pubkey repair_pubshares[N_HELPERS];
    unsigned char repaired_secshare[32];
    secp256k1_pubkey repaired_pubshare;

    /* The signing session that proves the enrolled participant works. */
    uint32_t signer_ids[N_SIGNERS] = { 2, NEW_ID };
    unsigned char signer_secshares[N_SIGNERS][32];
    secp256k1_pubkey signer_pubshares[N_SIGNERS];
    unsigned char msg[32] = "this_could_be_the_hash_of_a_msg";

    unsigned char buf[33];
    size_t outputlen;
    int i;

    ctx = secp256k1_context_create(SECP256K1_CONTEXT_NONE);

    printf("Generating threshold key material for a %d-of-%d group...", THRESHOLD, N_PARTICIPANTS);
    fflush(stdout);
    if (!trusted_dealer_keygen(ctx, threshold_seckey, secshares, &thresh_pk, pubshares)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");
    outputlen = sizeof(buf);
    if (!secp256k1_ec_pubkey_serialize(ctx, buf, &outputlen, &thresh_pk, SECP256K1_EC_COMPRESSED)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("Threshold public key:   ");
    print_hex(buf, outputlen);
    fflush(stdout);

    /* The new participant obtains the threshold public key from a source it
     * authenticates INDEPENDENTLY OF THE HELPERS. In this single-process
     * example that is the dealer step above; in a real deployment it might be
     * a signed group descriptor, an on-chain commitment, or whatever the group
     * already trusts to say what it is.
     *
     * It then validates the helpers' public shares against that key. This is
     * the step that makes the round 2 checks worth anything: without it, the
     * enrollee would be checking the helpers' arithmetic against numbers the
     * helpers also chose, and t colluding helpers could hand it a consistent
     * fake polynomial. */
    printf("Validating the group's public shares against the authenticated key...");
    fflush(stdout);
    if (!secp256k1_frost_threshold_info_validate(ctx, &thresh_pk, pubshares, N_PARTICIPANTS, THRESHOLD)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");

    /* Enrollment: helpers 0 and 1 give participant 3 a share. */
    for (i = 0; i < N_HELPERS; i++) {
        helpers[i].id = helper_ids[i];
        memcpy(helpers[i].secshare, &secshares[32 * helper_ids[i]], 32);
        helper_pubshares[i] = pubshares[helper_ids[i]];
    }
    printf("Enrolling participant %d with %d helpers...", NEW_ID, N_HELPERS);
    fflush(stdout);
    if (!enroll(ctx, new_secshare, &new_pubshare, helpers, helper_ids, helper_pubshares, &thresh_pk, NEW_ID)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");

    /* n -> n+1 bookkeeping. Every participant must extend its table of public
     * shares and update its record of n; participants that disagree about n
     * compute different Lagrange coefficients and cannot sign together. */
    pubshares[NEW_ID] = new_pubshare;
    printf("Validating the extended %d-participant key material...", N_PARTICIPANTS_AFTER);
    fflush(stdout);
    if (!secp256k1_frost_threshold_info_validate(ctx, &thresh_pk, pubshares, N_PARTICIPANTS_AFTER, THRESHOLD)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");
    outputlen = sizeof(buf);
    if (!secp256k1_ec_pubkey_serialize(ctx, buf, &outputlen, &new_pubshare, SECP256K1_EC_COMPRESSED)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("New participant's public share: ");
    print_hex(buf, outputlen);
    fflush(stdout);

    /* Sign with participant 2 and the newly enrolled participant 3. The
     * signature verifies against the group's original threshold public key:
     * enrollment does not change the group key or any existing share. */
    memcpy(signer_secshares[0], &secshares[32 * signer_ids[0]], 32);
    memcpy(signer_secshares[1], new_secshare, 32);
    signer_pubshares[0] = pubshares[signer_ids[0]];
    signer_pubshares[1] = pubshares[signer_ids[1]];
    printf("Signing with participants %u and %u...", signer_ids[0], signer_ids[1]);
    fflush(stdout);
    if (!sign_and_verify(ctx, signer_ids, signer_secshares, signer_pubshares, &thresh_pk, N_PARTICIPANTS_AFTER, msg, sizeof(msg))) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");

    /* Repair: participant 1 lost its share. The same three rounds at
     * new_id = 1 reproduce it exactly -- the share is a fixed value, f(x_1),
     * not a fresh random one. Note that n does NOT change here, and that this
     * is the mode where the missing authorization step bites hardest: whoever
     * convinces the helpers to run it receives participant 1's actual share.
     */
    for (i = 0; i < N_HELPERS; i++) {
        repair_helpers[i].id = repair_ids[i];
        memcpy(repair_helpers[i].secshare, &secshares[32 * repair_ids[i]], 32);
        repair_pubshares[i] = pubshares[repair_ids[i]];
    }
    printf("Repairing participant 1's lost share...");
    fflush(stdout);
    if (!enroll(ctx, repaired_secshare, &repaired_pubshare, repair_helpers, repair_ids, repair_pubshares, &thresh_pk, 1)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    if (memcmp(repaired_secshare, &secshares[32], 32) != 0) {
        printf("FAILED (the repaired share differs from the original)\n");
        return EXIT_FAILURE;
    }
    if (memcmp(&repaired_pubshare, &pubshares[1], sizeof(repaired_pubshare)) != 0) {
        printf("FAILED (the repaired public share differs from the original)\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");

    /* It's best practice to try to clear secrets from memory after using them.
     * This is done because some bugs can allow an attacker to leak memory, for
     * example through "out of bounds" array access (see Heartbleed), or the OS
     * swapping them to disk. Hence, we overwrite secret key material with
     * zeros.
     *
     * The session randomness was already wiped by shares_gen, and the
     * secnonces by frost_sign. */
    secure_erase(threshold_seckey, sizeof(threshold_seckey));
    secure_erase(secshares, sizeof(secshares));
    secure_erase(new_secshare, sizeof(new_secshare));
    secure_erase(repaired_secshare, sizeof(repaired_secshare));
    secure_erase(signer_secshares, sizeof(signer_secshares));
    for (i = 0; i < N_HELPERS; i++) {
        secure_erase(&helpers[i], sizeof(helpers[i]));
        secure_erase(&repair_helpers[i], sizeof(repair_helpers[i]));
    }
    secp256k1_context_destroy(ctx);
    return EXIT_SUCCESS;
}