/*************************************************************************
 * 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 ChillDKG module to run a 2-of-3
 *  distributed key generation (DKG) session for FROST (BIP 445) threshold
 *  signatures. Additionally, see the documentation in
 *  include/secp256k1_chilldkg.h and src/modules/chilldkg/chilldkg.md.
 *
 *  The example runs all roles (participants and coordinator) in a single
 *  process. In a real deployment these roles are performed by different
 *  parties communicating over secure/authenticated channels, and the
 *  coordinator is untrusted.
 *
 *  The resulting key material (secret share, threshold public key and public
 *  shares) can be used with the FROST signing module (see examples/frost.c).
 */

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

#include <secp256k1.h>
#include <secp256k1_chilldkg.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

struct participant {
    /* Long-term secret key; the participant's identity. All session outputs
     * can be recovered from the hostseckey and the recovery data. */
    unsigned char hostseckey[32];
    unsigned char hostpubkey[33];
    secp256k1_chilldkg_participant_state1 state1;
    secp256k1_chilldkg_participant_state2 state2;
    /* Buffers for the messages sent to the coordinator */
    unsigned char pmsg1[33 * THRESHOLD + 32 * N_PARTICIPANTS + 97];
    unsigned char pmsg2[64];
    /* Outputs */
    unsigned char secshare[32];
    unsigned char recovery[4 + 33 * THRESHOLD + 162 * N_PARTICIPANTS];
};

int main(void) {
    secp256k1_context *ctx;
    struct participant participants[N_PARTICIPANTS];
    secp256k1_chilldkg_coordinator_state coord_state;
    unsigned char hostpubkeys[N_PARTICIPANTS * 33];
    unsigned char params_hash[32];
    unsigned char cmsg1[162 * N_PARTICIPANTS + 33 * (THRESHOLD - 1)];
    unsigned char cmsg2[64 * N_PARTICIPANTS];
    const unsigned char *pmsg1_ptrs[N_PARTICIPANTS];
    const unsigned char *pmsg2_ptrs[N_PARTICIPANTS];
    unsigned char thresh_pk[33];
    unsigned char pubshares[N_PARTICIPANTS * 33];
    unsigned char coord_recovery[sizeof(participants[0].recovery)];
    uint32_t fault_index;
    secp256k1_chilldkg_fault fault;
    int i;

    ctx = secp256k1_context_create(SECP256K1_CONTEXT_NONE);

    /* Every participant generates a long-term host key pair. The hostseckey
     * must be generated with a cryptographically secure random number
     * generator and stored securely. */
    printf("Generating host keys...\n");
    for (i = 0; i < N_PARTICIPANTS; i++) {
        if (!fill_random(participants[i].hostseckey, 32)) {
            printf("Failed to generate randomness\n");
            return EXIT_FAILURE;
        }
        if (!secp256k1_chilldkg_hostpubkey_gen(ctx, participants[i].hostpubkey, participants[i].hostseckey)) {
            printf("Failed to generate host public key\n");
            return EXIT_FAILURE;
        }
        memcpy(&hostpubkeys[33 * i], participants[i].hostpubkey, 33);
    }

    /* The participants exchange their host public keys (over authenticated
     * channels) and compare the parameters hash out of band to ensure they
     * all agree on the session parameters. */
    if (!secp256k1_chilldkg_params_hash(ctx, params_hash, hostpubkeys, N_PARTICIPANTS, THRESHOLD)) {
        printf("Invalid session parameters\n");
        return EXIT_FAILURE;
    }
    printf("Session parameters hash: ");
    print_hex(params_hash, 32);

    /* Step 1: every participant creates its first message and sends it to the
     * coordinator. The randomness must be FRESH for every session. */
    printf("Running participant step 1...\n");
    for (i = 0; i < N_PARTICIPANTS; i++) {
        unsigned char random32[32];
        if (!fill_random(random32, 32)) {
            printf("Failed to generate randomness\n");
            return EXIT_FAILURE;
        }
        if (!secp256k1_chilldkg_participant_step1(ctx, &participants[i].state1, participants[i].pmsg1, participants[i].hostseckey, hostpubkeys, N_PARTICIPANTS, THRESHOLD, random32)) {
            printf("participant_step1 failed\n");
            return EXIT_FAILURE;
        }
        secure_erase(random32, sizeof(random32));
        pmsg1_ptrs[i] = participants[i].pmsg1;
    }

    /* Step 1 (coordinator): the coordinator aggregates the first messages and
     * broadcasts cmsg1 to all participants. */
    printf("Running coordinator step 1...\n");
    fault = secp256k1_chilldkg_coordinator_step1(ctx, &coord_state, cmsg1, &fault_index, pmsg1_ptrs, hostpubkeys, N_PARTICIPANTS, THRESHOLD);
    if (fault != SECP256K1_CHILLDKG_OK) {
        printf("coordinator_step1 failed (fault %d, index %u)\n", fault, fault_index);
        return EXIT_FAILURE;
    }

    /* Step 2: every participant verifies cmsg1, computes its DKG output, and
     * sends a CertEq signature over the session transcript to the
     * coordinator. */
    printf("Running participant step 2...\n");
    for (i = 0; i < N_PARTICIPANTS; i++) {
        unsigned char aux_rand32[32];
        if (!fill_random(aux_rand32, 32)) {
            printf("Failed to generate randomness\n");
            return EXIT_FAILURE;
        }
        fault = secp256k1_chilldkg_participant_step2(ctx, &participants[i].state2, participants[i].pmsg2, &fault_index, NULL, &participants[i].state1, participants[i].hostseckey, cmsg1, aux_rand32);
        secure_erase(aux_rand32, sizeof(aux_rand32));
        if (fault != SECP256K1_CHILLDKG_OK) {
            printf("participant_step2 failed for participant %d (fault %d, index %u)\n", i, fault, fault_index);
            return EXIT_FAILURE;
        }
        pmsg2_ptrs[i] = participants[i].pmsg2;
    }

    /* Step 2 (coordinator): the coordinator collects the CertEq signatures
     * into the certificate and broadcasts it to all participants. */
    printf("Running coordinator finalize...\n");
    fault = secp256k1_chilldkg_coordinator_finalize(ctx, cmsg2, thresh_pk, pubshares, coord_recovery, &fault_index, &coord_state, pmsg2_ptrs);
    if (fault != SECP256K1_CHILLDKG_OK) {
        printf("coordinator_finalize failed (fault %d, index %u)\n", fault, fault_index);
        return EXIT_FAILURE;
    }

    /* Finalize: every participant verifies the certificate and outputs the
     * DKG result and the recovery data. */
    printf("Running participant finalize...\n");
    for (i = 0; i < N_PARTICIPANTS; i++) {
        fault = secp256k1_chilldkg_participant_finalize(ctx, participants[i].secshare, thresh_pk, pubshares, participants[i].recovery, &fault_index, &participants[i].state2, cmsg2);
        if (fault != SECP256K1_CHILLDKG_OK) {
            printf("participant_finalize failed for participant %d (fault %d, index %u)\n", i, fault, fault_index);
            return EXIT_FAILURE;
        }
        /* All participants (and the coordinator) hold identical recovery
         * data. Keep it safe: anyone with the recovery data and their
         * hostseckey can recover the DKG output, e.g. after data loss. */
        if (memcmp(participants[i].recovery, coord_recovery, sizeof(coord_recovery)) != 0) {
            printf("recovery data mismatch\n");
            return EXIT_FAILURE;
        }
    }

    printf("Threshold public key:    ");
    print_hex(thresh_pk, 33);
    for (i = 0; i < N_PARTICIPANTS; i++) {
        printf("Public share of participant %d: ", i);
        print_hex(&pubshares[33 * i], 33);
    }

    /* Recovery: a participant can recover its DKG output from a backup of its
     * hostseckey and the recovery data at any time. */
    printf("Recovering participant 0 from recovery data...\n");
    {
        unsigned char rec_secshare[32];
        unsigned char rec_thresh_pk[33];
        unsigned char rec_pubshares[N_PARTICIPANTS * 33];
        unsigned char rec_hostpubkeys[N_PARTICIPANTS * 33];
        size_t n_rec;
        uint32_t t_rec;
        fault = secp256k1_chilldkg_participant_recover(ctx, rec_secshare, rec_thresh_pk, rec_pubshares, rec_hostpubkeys, &n_rec, &t_rec, &fault_index, participants[0].hostseckey, participants[0].recovery, sizeof(participants[0].recovery));
        if (fault != SECP256K1_CHILLDKG_OK) {
            printf("participant_recover failed (fault %d, index %u)\n", fault, fault_index);
            return EXIT_FAILURE;
        }
        if (n_rec != N_PARTICIPANTS || t_rec != THRESHOLD
            || memcmp(rec_secshare, participants[0].secshare, 32) != 0
            || memcmp(rec_thresh_pk, thresh_pk, 33) != 0) {
            printf("recovered output mismatch\n");
            return EXIT_FAILURE;
        }
        secure_erase(rec_secshare, sizeof(rec_secshare));
    }

    printf("DKG session completed successfully.\n");

    /* Clear secrets from memory (see examples/frost.c for why). The
     * participants' state2 objects have been consumed by
     * participant_finalize. */
    for (i = 0; i < N_PARTICIPANTS; i++) {
        secure_erase(participants[i].hostseckey, 32);
        secure_erase(participants[i].secshare, 32);
    }
    secp256k1_context_destroy(ctx);
    return EXIT_SUCCESS;
}