/*
 * Copyright (C) 2019 Red Hat, Inc.
 *
 * Author: Daiki Ueno
 *
 * This file is part of GNUTLS.
 *
 * The GNUTLS library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * as published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>
 *
 */

#if HAVE_CONFIG_H
# include "config.h"
#endif

#include "dsa-compute-k.h"

#include "gnutls_int.h"
#include "mem.h"
#include "mpn-base256.h"
#include <string.h>

#define BITS_TO_LIMBS(bits) (((bits) + GMP_NUMB_BITS - 1) / GMP_NUMB_BITS)

/* The maximum size of q, chosen from the fact that we support
 * 521-bit elliptic curve generator and 512-bit DSA subgroup at
 * maximum. */
#define MAX_Q_BITS 521
#define MAX_Q_SIZE ((MAX_Q_BITS + 7) / 8)
#define MAX_Q_LIMBS BITS_TO_LIMBS(MAX_Q_BITS)

#define MAX_HASH_BITS (MAX_HASH_SIZE * 8)
#define MAX_HASH_LIMBS BITS_TO_LIMBS(MAX_HASH_BITS)

int
_gnutls_dsa_compute_k(mpz_t k,
		      const mpz_t q,
		      const mpz_t x,
		      gnutls_mac_algorithm_t mac,
		      const uint8_t *digest,
		      size_t length)
{
	uint8_t V[MAX_HASH_SIZE];
	uint8_t K[MAX_HASH_SIZE];
	uint8_t xp[MAX_Q_SIZE];
	uint8_t tp[MAX_Q_SIZE];
	mp_limb_t h[MAX(MAX_Q_LIMBS, MAX_HASH_LIMBS)];
	mp_bitcnt_t q_bits = mpz_sizeinbase (q, 2);
	mp_size_t qn = mpz_size(q);
	mp_bitcnt_t h_bits = length * 8;
	mp_size_t hn = BITS_TO_LIMBS(h_bits);
	size_t nbytes = (q_bits + 7) / 8;
	const uint8_t c0 = 0x00;
	const uint8_t c1 = 0x01;
	mp_limb_t cy;
	gnutls_hmac_hd_t hd;
	int ret = 0;

	if (unlikely(q_bits > MAX_Q_BITS))
		return gnutls_assert_val(GNUTLS_E_INVALID_REQUEST);
	if (unlikely(length > MAX_HASH_SIZE))
		return gnutls_assert_val(GNUTLS_E_INVALID_REQUEST);

	/* int2octets(x) */
	mpn_get_base256(xp, nbytes, mpz_limbs_read(x), qn);

	/* bits2octets(h) */
	mpn_set_base256(h, hn, digest, length);

	if (hn < qn)
		/* qlen > blen: add zero bits to the left */
		mpn_zero(&h[hn], qn - hn);
	else if (h_bits > q_bits) {
		/* qlen < blen: keep the leftmost qlen bits.  We do this in 2
		 * steps because mpn_rshift only accepts shift count in the
		 * range 1 to mp_bits_per_limb-1.
		 */
		mp_bitcnt_t shift = h_bits - q_bits;

		if (shift / GMP_NUMB_BITS > 0) {
			mpn_copyi(h, &h[shift / GMP_NUMB_BITS], qn);
			hn -= shift / GMP_NUMB_BITS;
		}

		if (shift % GMP_NUMB_BITS > 0)
			mpn_rshift(h, h, hn, shift % GMP_NUMB_BITS);
	}

	cy = mpn_sub_n(h, h, mpz_limbs_read(q), qn);
	/* Fall back to addmul_1, if nettle is linked with mini-gmp. */
#ifdef mpn_cnd_add_n
	mpn_cnd_add_n(cy, h, h, mpz_limbs_read(q), qn);
#else
	mpn_addmul_1(h, mpz_limbs_read(q), qn, cy != 0);
#endif
	mpn_get_base256(tp, nbytes, h, qn);

	/* Step b */
	memset(V, c1, length);

	/* Step c */
	memset(K, c0, length);

	/* Step d */
	ret = gnutls_hmac_init(&hd, mac, K, length);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, V, length);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, &c0, 1);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, xp, nbytes);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, tp, nbytes);
	if (ret < 0)
		goto out;
	gnutls_hmac_deinit(hd, K);

	/* Step e */
	ret = gnutls_hmac_fast(mac, K, length, V, length, V);
	if (ret < 0)
		goto out;

	/* Step f */
	ret = gnutls_hmac_init(&hd, mac, K, length);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, V, length);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, &c1, 1);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, xp, nbytes);
	if (ret < 0)
		goto out;
	ret = gnutls_hmac(hd, tp, nbytes);
	if (ret < 0)
		goto out;
	gnutls_hmac_deinit(hd, K);

	/* Step g */
	ret = gnutls_hmac_fast(mac, K, length, V, length, V);
	if (ret < 0)
		goto out;

	/* Step h */
	for (;;) {
		/* Step 1 */
		size_t tlen = 0;

		/* Step 2 */
		while (tlen < nbytes) {
			size_t remaining = MIN(nbytes - tlen, length);
			ret = gnutls_hmac_fast(mac, K, length, V, length, V);
			if (ret < 0)
				goto out;
			memcpy (&tp[tlen], V, remaining);
			tlen += remaining;
		}

		/* Step 3 */
		mpn_set_base256 (h, qn, tp, tlen);
		if (tlen * 8 > q_bits)
			mpn_rshift (h, h, qn, tlen * 8 - q_bits);
		/* Check if k is in [1,q-1] */
		if (!mpn_zero_p (h, qn) &&
		    mpn_cmp (h, mpz_limbs_read(q), qn) < 0) {
			mpn_copyi(mpz_limbs_write(k, qn), h, qn);
			mpz_limbs_finish(k, qn);
			break;
		}

		ret = gnutls_hmac_init(&hd, mac, K, length);
		if (ret < 0)
			goto out;
		ret = gnutls_hmac(hd, V, length);
		if (ret < 0)
			goto out;
		ret = gnutls_hmac(hd, &c0, 1);
		if (ret < 0)
			goto out;
		gnutls_hmac_deinit(hd, K);

		ret = gnutls_hmac_fast(mac, K, length, V, length, V);
		if (ret < 0)
			goto out;
	}

 out:
	zeroize_key(xp, sizeof(xp));
	zeroize_key(tp, sizeof(tp));

	return ret;
}