/***********************************************************************
 * Copyright (c) 2015 Andrew Poelstra                                  *
 * Distributed under the MIT software license, see the accompanying    *
 * file COPYING or https://www.opensource.org/licenses/mit-license.php.*
 ***********************************************************************/

#ifndef SECP256K1_MODULE_ECDH_TESTS_H
#define SECP256K1_MODULE_ECDH_TESTS_H

#include "../../unit_test.h"
#include "../../testutil.h"

static int ecdh_hash_function_test_xpassthru(unsigned char *output, const unsigned char *x, const unsigned char *y, void *data) {
    (void)y;
    (void)data;
    memcpy(output, x, 32);
    return 1;
}

static int ecdh_hash_function_test_fail(unsigned char *output, const unsigned char *x, const unsigned char *y, void *data) {
    (void)output;
    (void)x;
    (void)y;
    (void)data;
    return 0;
}

static int ecdh_hash_function_custom(unsigned char *output, const unsigned char *x, const unsigned char *y, void *data) {
    (void)data;
    /* Save x and y as uncompressed public key */
    output[0] = 0x04;
    memcpy(output + 1, x, 32);
    memcpy(output + 33, y, 32);
    return 1;
}

static void test_ecdh_api(void) {
    secp256k1_pubkey point;
    unsigned char res[32];
    unsigned char s_one[32] = { 0 };
    s_one[31] = 1;

    CHECK(secp256k1_ec_pubkey_create(CTX, &point, s_one) == 1);

    /* Check all NULLs are detected */
    CHECK(secp256k1_ecdh(CTX, res, &point, s_one, NULL, NULL) == 1);
    CHECK_ILLEGAL(CTX, secp256k1_ecdh(CTX, NULL, &point, s_one, NULL, NULL));
    CHECK_ILLEGAL(CTX, secp256k1_ecdh(CTX, res, NULL, s_one, NULL, NULL));
    CHECK_ILLEGAL(CTX, secp256k1_ecdh(CTX, res, &point, NULL, NULL, NULL));
    CHECK(secp256k1_ecdh(CTX, res, &point, s_one, NULL, NULL) == 1);
}

static void test_ecdh_generator_basepoint(void) {
    unsigned char s_one[32] = { 0 };
    secp256k1_pubkey point[2];
    int i;

    s_one[31] = 1;
    /* Check against pubkey creation when the basepoint is the generator */
    for (i = 0; i < 2 * COUNT; ++i) {
        secp256k1_sha256 sha;
        unsigned char s_b32[32];
        unsigned char output_ecdh[65];
        unsigned char output_ser[32];
        unsigned char point_ser[65];
        size_t point_ser_len = sizeof(point_ser);
        secp256k1_scalar s;

        testutil_random_scalar_order(&s);
        secp256k1_scalar_get_b32(s_b32, &s);

        CHECK(secp256k1_ec_pubkey_create(CTX, &point[0], s_one) == 1);
        CHECK(secp256k1_ec_pubkey_create(CTX, &point[1], s_b32) == 1);

        /* compute using ECDH function with custom hash function */
        CHECK(secp256k1_ecdh(CTX, output_ecdh, &point[0], s_b32, ecdh_hash_function_custom, NULL) == 1);
        /* compute "explicitly" */
        CHECK(secp256k1_ec_pubkey_serialize(CTX, point_ser, &point_ser_len, &point[1], SECP256K1_EC_UNCOMPRESSED) == 1);
        /* compare */
        CHECK(secp256k1_memcmp_var(output_ecdh, point_ser, 65) == 0);

        /* compute using ECDH function with default hash function */
        CHECK(secp256k1_ecdh(CTX, output_ecdh, &point[0], s_b32, NULL, NULL) == 1);
        /* compute "explicitly" */
        CHECK(secp256k1_ec_pubkey_serialize(CTX, point_ser, &point_ser_len, &point[1], SECP256K1_EC_COMPRESSED) == 1);
        secp256k1_sha256_initialize(&sha);
        secp256k1_sha256_write(secp256k1_get_hash_context(CTX), &sha, point_ser, point_ser_len);
        secp256k1_sha256_finalize(secp256k1_get_hash_context(CTX), &sha, output_ser);
        /* compare */
        CHECK(secp256k1_memcmp_var(output_ecdh, output_ser, 32) == 0);
    }
}

DEFINE_SHA256_TRANSFORM_PROBE(sha256_ecdh)
static void test_ecdh_ctx_sha256(void) {
    /* Check ctx-provided SHA256 compression override takes effect */
    secp256k1_context *ctx = secp256k1_context_clone(CTX);
    unsigned char out_default[65], out_custom[65];
    const unsigned char sk[32] = {1};
    secp256k1_pubkey pubkey;
    CHECK(secp256k1_ec_pubkey_create(ctx, &pubkey, sk) == 1);

    /* Default behavior */
    CHECK(secp256k1_ecdh(ctx, out_default, &pubkey, sk, NULL, NULL) == 1);
    CHECK(!sha256_ecdh_called);

    /* Override SHA256 compression directly, bypassing the ctx setter sanity checks */
    ctx->hash_ctx.fn_sha256_compression = sha256_ecdh;
    CHECK(secp256k1_ecdh(ctx, out_custom, &pubkey, sk, NULL, NULL) == 1);

    /* Outputs must differ if custom compression was used */
    CHECK(secp256k1_memcmp_var(out_default, out_custom, 32) != 0);
    CHECK(sha256_ecdh_called);

    secp256k1_context_destroy(ctx);
}

static void test_bad_scalar(void) {
    unsigned char s_zero[32] = { 0 };
    unsigned char s_overflow[32] = { 0 };
    unsigned char s_rand[32] = { 0 };
    unsigned char output[32];
    secp256k1_scalar rand;
    secp256k1_pubkey point;

    /* Create random point */
    testutil_random_scalar_order(&rand);
    secp256k1_scalar_get_b32(s_rand, &rand);
    CHECK(secp256k1_ec_pubkey_create(CTX, &point, s_rand) == 1);

    /* Try to multiply it by bad values */
    memcpy(s_overflow, secp256k1_group_order_bytes, 32);
    CHECK(secp256k1_ecdh(CTX, output, &point, s_zero, NULL, NULL) == 0);
    CHECK(secp256k1_ecdh(CTX, output, &point, s_overflow, NULL, NULL) == 0);
    /* ...and a good one */
    s_overflow[31] -= 1;
    CHECK(secp256k1_ecdh(CTX, output, &point, s_overflow, NULL, NULL) == 1);

    /* Hash function failure results in ecdh failure */
    CHECK(secp256k1_ecdh(CTX, output, &point, s_overflow, ecdh_hash_function_test_fail, NULL) == 0);
}

/** Test that ECDH(sG, 1/s) == ECDH((1/s)G, s) == ECDH(G, 1) for a few random s. */
static void test_result_basepoint(void) {
    secp256k1_pubkey point;
    secp256k1_scalar rand;
    unsigned char s[32];
    unsigned char s_inv[32];
    unsigned char out[32];
    unsigned char out_inv[32];
    unsigned char out_base[32];
    int i;

    unsigned char s_one[32] = { 0 };
    s_one[31] = 1;
    CHECK(secp256k1_ec_pubkey_create(CTX, &point, s_one) == 1);
    CHECK(secp256k1_ecdh(CTX, out_base, &point, s_one, NULL, NULL) == 1);

    for (i = 0; i < 2 * COUNT; i++) {
        testutil_random_scalar_order(&rand);
        secp256k1_scalar_get_b32(s, &rand);
        secp256k1_scalar_inverse(&rand, &rand);
        secp256k1_scalar_get_b32(s_inv, &rand);

        CHECK(secp256k1_ec_pubkey_create(CTX, &point, s) == 1);
        CHECK(secp256k1_ecdh(CTX, out, &point, s_inv, NULL, NULL) == 1);
        CHECK(secp256k1_memcmp_var(out, out_base, 32) == 0);

        CHECK(secp256k1_ec_pubkey_create(CTX, &point, s_inv) == 1);
        CHECK(secp256k1_ecdh(CTX, out_inv, &point, s, NULL, NULL) == 1);
        CHECK(secp256k1_memcmp_var(out_inv, out_base, 32) == 0);
    }
}

static void test_ecdh_wycheproof(void) {
#include "../../wycheproof/ecdh_secp256k1_test.h"
    int t;
    for (t = 0; t < SECP256K1_ECDH_WYCHEPROOF_NUMBER_TESTVECTORS; t++) {
        int parsed_ok;
        secp256k1_pubkey point;
        const unsigned char *pk;
        const unsigned char *sk;
        const unsigned char *expected_shared_secret;
        unsigned char output_ecdh[65] = { 0 };

        int expected_result;

        memset(&point, 0, sizeof(point));
        pk = &wycheproof_ecdh_public_keys[testvectors[t].pk_offset];
        parsed_ok = secp256k1_ec_pubkey_parse(CTX, &point, pk, testvectors[t].pk_len);

        expected_result = testvectors[t].expected_result;
        CHECK(parsed_ok == expected_result);
        if (!parsed_ok) {
            continue;
        }

        sk = &wycheproof_ecdh_private_keys[testvectors[t].sk_offset];
        CHECK(testvectors[t].sk_len == 32);

        CHECK(secp256k1_ecdh(CTX, output_ecdh, &point, sk, ecdh_hash_function_test_xpassthru, NULL) == 1);
        expected_shared_secret = &wycheproof_ecdh_shared_secrets[testvectors[t].shared_offset];

        CHECK(secp256k1_memcmp_var(output_ecdh, expected_shared_secret, testvectors[t].shared_len) == 0);
    }
}

/* --- Test registry --- */
static const struct tf_test_entry tests_ecdh[] = {
    CASE1(test_ecdh_api),
    CASE1(test_ecdh_generator_basepoint),
    CASE1(test_bad_scalar),
    CASE1(test_result_basepoint),
    CASE1(test_ecdh_wycheproof),
    CASE1(test_ecdh_ctx_sha256),
};

#endif /* SECP256K1_MODULE_ECDH_TESTS_H */
