/*************************************************************************
 * 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 module (BIP 445) to create a
 *  2-of-3 threshold signature. Additionally, see the documentation in
 *  include/secp256k1_frost.h and src/modules/frost/frost.md.
 *
 *  The example runs all roles (trusted dealer, signers, coordinator) in a
 *  single process. In a real deployment these roles are performed by
 *  different parties communicating over secure/authenticated channels.
 */

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

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

#include "examples_util.h"

/* Total number of participants n */
#define N_PARTICIPANTS 3
/* Threshold t: the minimum number of signers required to produce a
 * signature */
#define THRESHOLD 2
/* The signers of this signing session (u = 2, participants 0 and 2) */
#define N_SIGNERS 2
static const uint32_t SIGNER_IDS[N_SIGNERS] = { 0, 2 };

struct signer_secrets {
    /* The signer's secret share, received from the dealer over a secure
     * channel */
    unsigned char secshare[32];
    /* Secret nonce. It is wiped by secp256k1_frost_sign; never reuse it. */
    secp256k1_frost_secnonce secnonce;
};

struct signer {
    uint32_t id;
    secp256k1_pubkey pubshare;
    secp256k1_frost_pubnonce pubnonce;
    secp256k1_frost_partial_sig partial_sig;
};

/* Run the trusted dealer key generation: from a random threshold secret key,
 * derive the secret share of every participant, the threshold public key and
 * the public shares.
 *
 * WARNING: The trusted dealer knows the threshold secret key and all secret
 * shares, and must erase them securely after distributing the shares. A
 * distributed key generation (DKG) protocol avoids a trusted dealer entirely
 * but is out of scope for BIP 445 and for this example. */
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;
    }
    /* Everyone can check that the public shares and the threshold public key
     * are consistent. Note that this does NOT validate the security of the
     * key generation that produced them. */
    if (!secp256k1_frost_threshold_info_validate(ctx, thresh_pk, pubshares, N_PARTICIPANTS, THRESHOLD)) {
        return 0;
    }
    return 1;
}

/* Apply an x-only tweak to the threshold public key, e.g. a BIP 341
 * ("Taproot") tweak committing to a script tree. All participants (and the
 * coordinator) must apply exactly the same tweaks in the same order to their
 * own tweak cache. */
static int tweak(const secp256k1_context* ctx, secp256k1_frost_tweak_cache *cache, secp256k1_xonly_pubkey *tweaked_pk) {
    /* For Taproot tweaking the tweak is set to the TapTweak hash as defined
     * in BIP 341. */
    unsigned char xonly_tweak[32] = "this could be a Taproot tweak..";

    if (!secp256k1_frost_pubkey_xonly_tweak_add(ctx, tweaked_pk, cache, xonly_tweak)) {
        return 0;
    }
    /* For BIP 32-style plain tweaking, use
     * secp256k1_frost_pubkey_ec_tweak_add instead. */
    return 1;
}

/* Sign a message with the given signers and store the resulting BIP340
 * signature in sig64. */
static int sign(const secp256k1_context* ctx, struct signer_secrets *signer_secrets, struct signer *signers, const secp256k1_frost_tweak_cache *cache, const unsigned char *msg, size_t msglen, unsigned char *sig64) {
    int i;
    const secp256k1_frost_pubnonce *pubnonces[N_SIGNERS];
    const secp256k1_frost_partial_sig *partial_sigs[N_SIGNERS];
    secp256k1_pubkey signer_pubshares[N_SIGNERS];
    uint32_t ids[N_SIGNERS];
    /* The same for all signers and the coordinator */
    secp256k1_frost_aggnonce aggnonce;
    secp256k1_frost_session session;
    /* The x-only encoding of the tweaked threshold public key, bound into the
     * nonce derivation */
    secp256k1_xonly_pubkey tweaked_pk;
    unsigned char tweaked_pk32[32];

    if (!secp256k1_frost_tweaked_pubkey_get(ctx, &tweaked_pk, cache)) {
        return 0;
    }
    if (!secp256k1_xonly_pubkey_serialize(ctx, tweaked_pk32, &tweaked_pk)) {
        return 0;
    }

    /* Every signer creates a nonce pair and sends the pubnonce to the
     * coordinator. */
    for (i = 0; i < N_SIGNERS; i++) {
        unsigned char session_secrand[32];

        ids[i] = signers[i].id;
        signer_pubshares[i] = signers[i].pubshare;
        pubnonces[i] = &signers[i].pubnonce;
        partial_sigs[i] = &signers[i].partial_sig;

        /* Create random session randomness. It is absolutely necessary that
         * this is unique for every call of secp256k1_frost_nonce_gen.
         * Otherwise it's trivial for an attacker to extract the secret share!
         * nonce_gen wipes session_secrand before returning. */
        if (!fill_random(session_secrand, sizeof(session_secrand))) {
            printf("Failed to generate randomness\n");
            return 0;
        }
        /* Bind the secret share (defense-in-depth against bad randomness),
         * the public share, the tweaked threshold public key and the message
         * into the nonce derivation. */
        if (!secp256k1_frost_nonce_gen(ctx, &signer_secrets[i].secnonce, &signers[i].pubnonce, session_secrand, signer_secrets[i].secshare, &signers[i].pubshare, tweaked_pk32, msg, msglen, NULL, 0)) {
            return 0;
        }
        secure_erase(session_secrand, sizeof(session_secrand));
    }

    /* Communication round 1: the coordinator aggregates the pubnonces and
     * sends the aggregate nonce to the signers. */
    if (!secp256k1_frost_nonce_agg(ctx, &aggnonce, NULL, pubnonces, N_SIGNERS)) {
        return 0;
    }

    /* Every signer and the coordinator run session_init locally with
     * identical arguments (aside from the session object), which yields
     * identical sessions. */
    if (!secp256k1_frost_session_init(ctx, &session, &aggnonce, ids, signer_pubshares, N_SIGNERS, N_PARTICIPANTS, THRESHOLD, cache, msg, msglen)) {
        return 0;
    }

    /* Every signer creates a partial signature and sends it to the
     * coordinator. frost_sign wipes the secnonce. That's because you must
     * _never_ reuse the secnonce (or use the same session_secrand to create a
     * secnonce). If you do, you effectively reuse the nonce and leak the
     * secret share. */
    for (i = 0; i < N_SIGNERS; i++) {
        if (!secp256k1_frost_sign(ctx, &signers[i].partial_sig, &signer_secrets[i].secnonce, signer_secrets[i].secshare, &session, ids, signer_pubshares, N_SIGNERS, signers[i].id)) {
            return 0;
        }
    }

    /* Communication round 2: the coordinator verifies the partial signatures
     * and aggregates them. Verifying the individual partial signatures (as
     * opposed to only verifying the final signature) allows the coordinator
     * to identify which signer misbehaved if the protocol run fails. */
    for (i = 0; i < N_SIGNERS; i++) {
        if (!secp256k1_frost_partial_sig_verify(ctx, &signers[i].partial_sig, &signers[i].pubnonce, &signers[i].pubshare, &session, ids, N_SIGNERS, i)) {
            return 0;
        }
    }
    return secp256k1_frost_partial_sig_agg(ctx, sig64, NULL, &session, partial_sigs, N_SIGNERS);
}

int main(void) {
    secp256k1_context* ctx;
    int i;
    /* Secret key material known to the trusted dealer. The dealer must erase
     * it after the participants received their shares (done 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];
    secp256k1_frost_tweak_cache cache;
    secp256k1_xonly_pubkey tweaked_pk;
    struct signer_secrets signer_secrets[N_SIGNERS];
    struct signer signers[N_SIGNERS];
    unsigned char msg[32] = "this_could_be_the_hash_of_a_msg";
    unsigned char sig[64];
    unsigned char buf[33];
    size_t outputlen;

    /* Create a secp256k1 context */
    ctx = secp256k1_context_create(SECP256K1_CONTEXT_NONE);

    printf("Generating threshold key material...");
    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 dealer hands each participant their secret share (over a secure
     * channel). This example signs with the participants 0 and 2. */
    for (i = 0; i < N_SIGNERS; i++) {
        signers[i].id = SIGNER_IDS[i];
        signers[i].pubshare = pubshares[SIGNER_IDS[i]];
        memcpy(signer_secrets[i].secshare, &secshares[32 * SIGNER_IDS[i]], 32);
    }

    printf("Applying x-only tweak (Taproot-style)...");
    fflush(stdout);
    /* Every participant and the coordinator initializes a tweak cache from
     * the threshold public key and applies the same tweaks. */
    if (!secp256k1_frost_tweak_cache_init(ctx, &cache, &thresh_pk)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    if (!tweak(ctx, &cache, &tweaked_pk)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");
    if (!secp256k1_xonly_pubkey_serialize(ctx, buf, &tweaked_pk)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("Tweaked threshold public key: ");
    print_hex(buf, 32);
    fflush(stdout);

    printf("Signing message with %d-of-%d signers...", N_SIGNERS, N_PARTICIPANTS);
    fflush(stdout);
    if (!sign(ctx, signer_secrets, signers, &cache, msg, sizeof(msg), sig)) {
        printf("FAILED\n");
        return EXIT_FAILURE;
    }
    printf("ok\n");

    printf("Verifying signature.....");
    fflush(stdout);
    if (!secp256k1_schnorrsig_verify(ctx, sig, msg, sizeof(msg), &tweaked_pk)) {
        printf("FAILED\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.
     *
     * Here we are preventing these writes from being optimized out, as any good compiler
     * will remove any writes that aren't used.
     *
     * The secnonces have already been wiped by secp256k1_frost_sign. */
    secure_erase(threshold_seckey, sizeof(threshold_seckey));
    secure_erase(secshares, sizeof(secshares));
    for (i = 0; i < N_SIGNERS; i++) {
        secure_erase(&signer_secrets[i], sizeof(signer_secrets[i]));
    }
    secp256k1_context_destroy(ctx);
    return EXIT_SUCCESS;
}