/*
 *  Copyright (C) 2021 - This file is part of libecc project
 *
 *  Authors:
 *      Ryad BENADJILA <ryadbenadjila@gmail.com>
 *      Arnaud EBALARD <arnaud.ebalard@ssi.gouv.fr>
 *
 *  This software is licensed under a dual BSD and GPL v2 license.
 *  See LICENSE file at the root folder of the project.
 */
#include <libecc/lib_ecc_config.h>
#if defined(WITH_ECCCDH)

#include <libecc/ecdh/ecccdh.h>

/*
 * This module implements the "Elliptic Curve Cryptography Cofactor Diffie-Hellman (ECC CDH)
 * Primitive" as described in section 5.7.1.2 of the NIST SP 800-56A Rev. 3 standard.
 *
 */

/*
 * Get the size of the shared secret associated to the curve parameters.
 */
int ecccdh_shared_secret_size(const ec_params *params, u8 *size)
{
	int ret;

	MUST_HAVE((params != NULL) && (size != NULL), ret, err);
	MUST_HAVE((BYTECEIL(params->ec_fp.p_bitlen) <= 255), ret, err);

	(*size) = (u8)(BYTECEIL(params->ec_fp.p_bitlen));
	ret = 0;

err:
	return ret;
}

/*
 * Get the size of the serialized public key associated to the curve parameters.
 */
int ecccdh_serialized_pub_key_size(const ec_params *params, u8 *size)
{
	int ret;

	MUST_HAVE((params != NULL) && (size != NULL), ret, err);
	MUST_HAVE(((2 * BYTECEIL(params->ec_fp.p_bitlen)) <= 255), ret, err);

	(*size) = (u8)(2 * BYTECEIL(params->ec_fp.p_bitlen));
	ret = 0;

err:
	return ret;
}


/*
 * Initialize ECCCDH public key from an initialized private key.
 */
int ecccdh_init_pub_key(ec_pub_key *out_pub, const ec_priv_key *in_priv)
{
	prj_pt_src_t G;
	int ret, cmp;
	nn_src_t q;

	MUST_HAVE((out_pub != NULL), ret, err);

	/* Zero init public key to be generated */
	ret = local_memset(out_pub, 0, sizeof(ec_pub_key)); EG(ret, err);

	ret = priv_key_check_initialized_and_type(in_priv, ECCCDH); EG(ret, err);
	q = &(in_priv->params->ec_gen_order);

	/* Sanity check on key compliance */
	MUST_HAVE((!nn_cmp(&(in_priv->x), q, &cmp)) && (cmp < 0), ret, err);

	/* Y = xG */
	G = &(in_priv->params->ec_gen);
	/* Use blinding when computing point scalar multiplication */
	ret = prj_pt_mul_blind(&(out_pub->y), &(in_priv->x), G); EG(ret, err);

	out_pub->key_type = ECCCDH;
	out_pub->params = in_priv->params;
	out_pub->magic = PUB_KEY_MAGIC;

err:
        return ret;
}

/*
 * Generate a key pair for ECCCDH given curve parameters as input.
 */
int ecccdh_gen_key_pair(ec_key_pair *kp, const ec_params *params)
{
	int ret;

	MUST_HAVE((kp != NULL) && (params != NULL), ret, err);

	/* Use our generic key pair generation primitive */
	kp->priv_key.magic = PRIV_KEY_MAGIC;
	kp->priv_key.key_type = ECCCDH;
	kp->priv_key.params = params;
	ret = generic_gen_priv_key(&(kp->priv_key)); EG(ret, err);

	/* Then, derive the public key */
	ret = ecccdh_init_pub_key(&(kp->pub_key), &(kp->priv_key));

err:
	/* If we have failed our generation, uninitialize
	 * the key pair.
	 */
	if(ret && (kp != NULL)){
		IGNORE_RET_VAL(local_memset(kp, 0, sizeof(ec_key_pair)));
	}
	return ret;
}

/*
 * Create a key pair from a serialized private key.
 */
int ecccdh_import_key_pair_from_priv_key_buf(ec_key_pair *kp, const ec_params *params, const u8 *priv_key_buf, u8 priv_key_buf_len)
{
	int ret;

	MUST_HAVE((kp != NULL), ret, err);

	/* Use our import primitive */
	ret = ec_priv_key_import_from_buf(&(kp->priv_key), params, priv_key_buf, priv_key_buf_len, ECCCDH); EG(ret, err);

	/* Now derive the public key from the private one */
	ret = ecccdh_init_pub_key(&(kp->pub_key), &(kp->priv_key));

err:
	return ret;
}

/*
 * Serialize our public key in a buffer.
 */
int ecccdh_serialize_pub_key(const ec_pub_key *our_pub_key, u8 *buf, u8 buf_len)
{
	int ret, iszero;

	/* Sanity check */
	ret = pub_key_check_initialized_and_type(our_pub_key, ECCCDH); EG(ret, err);

	/* Reject the point at infinity */
	ret = prj_pt_iszero(&(our_pub_key->y), &iszero); EG(ret, err);
	MUST_HAVE((!iszero), ret, err);

	/* Export our public key as an affine point
	 * NOTE: sanity checks on buf_len are performed in the lower layers.
	 */
	ret = ec_pub_key_export_to_aff_buf(our_pub_key, buf, buf_len);

err:
	return ret;
}

/*
 * Derive the ECCCDH shared secret and store it in a buffer given the peer
 * public key and our private key.
 *
 * The shared_secret_len length MUST be exactly equal to the expected shared secret size:
 * the function fails otherwise.
 */
int ecccdh_derive_secret(const ec_priv_key *our_priv_key, const u8 *peer_pub_key_buf, u8 peer_pub_key_buf_len, u8 *shared_secret, u8 shared_secret_len)
{
	int ret, iszero, isone;
	ec_pub_key peer_pub_key;
	prj_pt_t Q;
	nn_src_t cofactor;
	u8 expected_shared_secret_len;
	peer_pub_key.magic = WORD(0);

	/* Sanity checks */
	MUST_HAVE((shared_secret != NULL), ret, err);
	ret = priv_key_check_initialized_and_type(our_priv_key, ECCCDH); EG(ret, err);

	/* Try to import the peer public key.
	 * NOTE: the check that this public key is indeed on the curve is performed in the lower layer
	 * functions.
	 */
	ret = ec_pub_key_import_from_aff_buf(&peer_pub_key, our_priv_key->params, peer_pub_key_buf, peer_pub_key_buf_len, ECCCDH); EG(ret, err);
	Q = &(peer_pub_key.y);

	cofactor = &(our_priv_key->params->ec_gen_cofactor);
	ret = nn_isone(cofactor, &isone); EG(ret, err);
	if(!isone){
		/* Perform a cofactor multiplication if necessary.
		 * NOTE: since the cofactor and the base point are public, we perform an unprotected
		 * scalar multiplication here.
		 */
		ret = _prj_pt_unprotected_mult(Q, cofactor, Q); EG(ret, err);
	}

	/*
	 * Reject the point at infinity or low order point as input as a trivial wrong public key.
	 * This would be rejected in any case by the check post scalar multiplication below, but we
	 * do not want to use and possibly leak the secret scalar if not necessary!
	 */
	ret = prj_pt_iszero(Q, &iszero); EG(ret, err);
	MUST_HAVE((!iszero), ret, err);

	/* Compute the shared secret using scalar multiplication */
#ifdef USE_SIG_BLINDING
	ret = prj_pt_mul_blind(Q, &(our_priv_key->x), Q); EG(ret, err);
#else
	ret = prj_pt_mul(Q, &(our_priv_key->x), Q); EG(ret, err);
#endif

	/* NOTE: scalar multiplication primitive checks that the resulting point is on
	 * the curve.
	 */
	/* Reject the point at infinity */
	ret = prj_pt_iszero(Q, &iszero); EG(ret, err);
	MUST_HAVE((!iszero), ret, err);

	/* Get the unique affine representation of the resulting point */
	ret = prj_pt_unique(Q, Q); EG(ret, err);
	/* Now export the X coordinate as the shared secret in the output buffer */
	ret = ecccdh_shared_secret_size(our_priv_key->params, &expected_shared_secret_len); EG(ret, err);
	MUST_HAVE((shared_secret_len == expected_shared_secret_len), ret, err);
	ret = fp_export_to_buf(shared_secret, shared_secret_len, &(Q->X));

err:
	PTR_NULLIFY(Q);
	PTR_NULLIFY(cofactor);
	/* Uninit local peer pub key and zeroize intermediate computations */
	IGNORE_RET_VAL(local_memset(&peer_pub_key, 0, sizeof(ec_pub_key)));

	return ret;
}

#else /* !defined(WITH_ECCCDH) */

/*
 * Dummy definition to avoid the empty translation unit ISO C warning
 */
typedef int dummy;

#endif /* WITH_ECCCDH */