/* Based on the public domain implementation in
 * crypto_hash/keccakc512/simple/ from http://bench.cr.yp.to/supercop.html
 * by Ronny Van Keer
 * and the public domain "TweetFips202" implementation
 * from https://twitter.com/tweetfips202
 * by Gilles Van Assche, Daniel J. Bernstein, and Peter Schwabe
 * SPDX-License-Identifier: Public domain
 */

#include "fips202_kyber_r2.h"
#include "symmetric.h"

#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define PQC_SHAKEINCCTX_BYTES (sizeof(uint64_t)*26)
#define PQC_SHAKECTX_BYTES (sizeof(uint64_t)*25)

#define NROUNDS 24
#define ROL(a, offset) (((a) << (offset)) ^ ((a) >> (64 - (offset))))

/*************************************************
 * Name:        load64
 *
 * Description: Load 8 bytes into uint64_t in little-endian order
 *
 * Arguments:   - const uint8_t *x: pointer to input byte array
 *
 * Returns the loaded 64-bit unsigned integer
 **************************************************/
static uint64_t load64(const uint8_t *x) {
	uint64_t r = 0;
	for (size_t i = 0; i < 8; ++i) {
		r |= (uint64_t)x[i] << 8 * i;
	}

	return r;
}

/*************************************************
 * Name:        store64
 *
 * Description: Store a 64-bit integer to a byte array in little-endian order
 *
 * Arguments:   - uint8_t *x: pointer to the output byte array
 *              - uint64_t u: input 64-bit unsigned integer
 **************************************************/
static void store64(uint8_t *x, uint64_t u) {
	for (size_t i = 0; i < 8; ++i) {
		x[i] = (uint8_t) (u >> 8 * i);
	}
}

/* Keccak round constants */
static const uint64_t KeccakF_RoundConstants[NROUNDS] = {
	0x0000000000000001ULL, 0x0000000000008082ULL,
	0x800000000000808aULL, 0x8000000080008000ULL,
	0x000000000000808bULL, 0x0000000080000001ULL,
	0x8000000080008081ULL, 0x8000000000008009ULL,
	0x000000000000008aULL, 0x0000000000000088ULL,
	0x0000000080008009ULL, 0x000000008000000aULL,
	0x000000008000808bULL, 0x800000000000008bULL,
	0x8000000000008089ULL, 0x8000000000008003ULL,
	0x8000000000008002ULL, 0x8000000000000080ULL,
	0x000000000000800aULL, 0x800000008000000aULL,
	0x8000000080008081ULL, 0x8000000000008080ULL,
	0x0000000080000001ULL, 0x8000000080008008ULL
};

/*************************************************
 * Name:        KeccakF1600_StatePermute
 *
 * Description: The Keccak F1600 Permutation
 *
 * Arguments:   - uint64_t *state: pointer to input/output Keccak state
 **************************************************/
static void KeccakF1600_StatePermute(uint64_t *state) {
	int round;

	uint64_t Aba, Abe, Abi, Abo, Abu;
	uint64_t Aga, Age, Agi, Ago, Agu;
	uint64_t Aka, Ake, Aki, Ako, Aku;
	uint64_t Ama, Ame, Ami, Amo, Amu;
	uint64_t Asa, Ase, Asi, Aso, Asu;
	uint64_t BCa, BCe, BCi, BCo, BCu;

	// copyFromState(A, state)
	Aba = state[0];
	Abe = state[1];
	Abi = state[2];
	Abo = state[3];
	Abu = state[4];
	Aga = state[5];
	Age = state[6];
	Agi = state[7];
	Ago = state[8];
	Agu = state[9];
	Aka = state[10];
	Ake = state[11];
	Aki = state[12];
	Ako = state[13];
	Aku = state[14];
	Ama = state[15];
	Ame = state[16];
	Ami = state[17];
	Amo = state[18];
	Amu = state[19];
	Asa = state[20];
	Ase = state[21];
	Asi = state[22];
	Aso = state[23];
	Asu = state[24];

	for (round = 0; round < NROUNDS; round += 2) {
	    uint64_t Da, De, Di, Do, Du;
	    uint64_t Eba, Ebe, Ebi, Ebo, Ebu;
	    uint64_t Ega, Ege, Egi, Ego, Egu;
	    uint64_t Eka, Eke, Eki, Eko, Eku;
	    uint64_t Ema, Eme, Emi, Emo, Emu;
	    uint64_t Esa, Ese, Esi, Eso, Esu;

	    //    prepareTheta
	    BCa = Aba ^ Aga ^ Aka ^ Ama ^ Asa;
	    BCe = Abe ^ Age ^ Ake ^ Ame ^ Ase;
	    BCi = Abi ^ Agi ^ Aki ^ Ami ^ Asi;
	    BCo = Abo ^ Ago ^ Ako ^ Amo ^ Aso;
	    BCu = Abu ^ Agu ^ Aku ^ Amu ^ Asu;

	    // thetaRhoPiChiIotaPrepareTheta(round  , A, E)
	    Da = BCu ^ ROL(BCe, 1);
	    De = BCa ^ ROL(BCi, 1);
	    Di = BCe ^ ROL(BCo, 1);
	    Do = BCi ^ ROL(BCu, 1);
	    Du = BCo ^ ROL(BCa, 1);

	    Aba ^= Da;
	    BCa = Aba;
	    Age ^= De;
	    BCe = ROL(Age, 44);
	    Aki ^= Di;
	    BCi = ROL(Aki, 43);
	    Amo ^= Do;
	    BCo = ROL(Amo, 21);
	    Asu ^= Du;
	    BCu = ROL(Asu, 14);
	    Eba = BCa ^ ((~BCe) & BCi);
	    Eba ^= KeccakF_RoundConstants[round];
	    Ebe = BCe ^ ((~BCi) & BCo);
	    Ebi = BCi ^ ((~BCo) & BCu);
	    Ebo = BCo ^ ((~BCu) & BCa);
	    Ebu = BCu ^ ((~BCa) & BCe);

	    Abo ^= Do;
	    BCa = ROL(Abo, 28);
	    Agu ^= Du;
	    BCe = ROL(Agu, 20);
	    Aka ^= Da;
	    BCi = ROL(Aka, 3);
	    Ame ^= De;
	    BCo = ROL(Ame, 45);
	    Asi ^= Di;
	    BCu = ROL(Asi, 61);
	    Ega = BCa ^ ((~BCe) & BCi);
	    Ege = BCe ^ ((~BCi) & BCo);
	    Egi = BCi ^ ((~BCo) & BCu);
	    Ego = BCo ^ ((~BCu) & BCa);
	    Egu = BCu ^ ((~BCa) & BCe);

	    Abe ^= De;
	    BCa = ROL(Abe, 1);
	    Agi ^= Di;
	    BCe = ROL(Agi, 6);
	    Ako ^= Do;
	    BCi = ROL(Ako, 25);
	    Amu ^= Du;
	    BCo = ROL(Amu, 8);
	    Asa ^= Da;
	    BCu = ROL(Asa, 18);
	    Eka = BCa ^ ((~BCe) & BCi);
	    Eke = BCe ^ ((~BCi) & BCo);
	    Eki = BCi ^ ((~BCo) & BCu);
	    Eko = BCo ^ ((~BCu) & BCa);
	    Eku = BCu ^ ((~BCa) & BCe);

	    Abu ^= Du;
	    BCa = ROL(Abu, 27);
	    Aga ^= Da;
	    BCe = ROL(Aga, 36);
	    Ake ^= De;
	    BCi = ROL(Ake, 10);
	    Ami ^= Di;
	    BCo = ROL(Ami, 15);
	    Aso ^= Do;
	    BCu = ROL(Aso, 56);
	    Ema = BCa ^ ((~BCe) & BCi);
	    Eme = BCe ^ ((~BCi) & BCo);
	    Emi = BCi ^ ((~BCo) & BCu);
	    Emo = BCo ^ ((~BCu) & BCa);
	    Emu = BCu ^ ((~BCa) & BCe);

	    Abi ^= Di;
	    BCa = ROL(Abi, 62);
	    Ago ^= Do;
	    BCe = ROL(Ago, 55);
	    Aku ^= Du;
	    BCi = ROL(Aku, 39);
	    Ama ^= Da;
	    BCo = ROL(Ama, 41);
	    Ase ^= De;
	    BCu = ROL(Ase, 2);
	    Esa = BCa ^ ((~BCe) & BCi);
	    Ese = BCe ^ ((~BCi) & BCo);
	    Esi = BCi ^ ((~BCo) & BCu);
	    Eso = BCo ^ ((~BCu) & BCa);
	    Esu = BCu ^ ((~BCa) & BCe);

	    //    prepareTheta
	    BCa = Eba ^ Ega ^ Eka ^ Ema ^ Esa;
	    BCe = Ebe ^ Ege ^ Eke ^ Eme ^ Ese;
	    BCi = Ebi ^ Egi ^ Eki ^ Emi ^ Esi;
	    BCo = Ebo ^ Ego ^ Eko ^ Emo ^ Eso;
	    BCu = Ebu ^ Egu ^ Eku ^ Emu ^ Esu;

	    // thetaRhoPiChiIotaPrepareTheta(round+1, E, A)
	    Da = BCu ^ ROL(BCe, 1);
	    De = BCa ^ ROL(BCi, 1);
	    Di = BCe ^ ROL(BCo, 1);
	    Do = BCi ^ ROL(BCu, 1);
	    Du = BCo ^ ROL(BCa, 1);

	    Eba ^= Da;
	    BCa = Eba;
	    Ege ^= De;
	    BCe = ROL(Ege, 44);
	    Eki ^= Di;
	    BCi = ROL(Eki, 43);
	    Emo ^= Do;
	    BCo = ROL(Emo, 21);
	    Esu ^= Du;
	    BCu = ROL(Esu, 14);
	    Aba = BCa ^ ((~BCe) & BCi);
	    Aba ^= KeccakF_RoundConstants[round + 1];
	    Abe = BCe ^ ((~BCi) & BCo);
	    Abi = BCi ^ ((~BCo) & BCu);
	    Abo = BCo ^ ((~BCu) & BCa);
	    Abu = BCu ^ ((~BCa) & BCe);

	    Ebo ^= Do;
	    BCa = ROL(Ebo, 28);
	    Egu ^= Du;
	    BCe = ROL(Egu, 20);
	    Eka ^= Da;
	    BCi = ROL(Eka, 3);
	    Eme ^= De;
	    BCo = ROL(Eme, 45);
	    Esi ^= Di;
	    BCu = ROL(Esi, 61);
	    Aga = BCa ^ ((~BCe) & BCi);
	    Age = BCe ^ ((~BCi) & BCo);
	    Agi = BCi ^ ((~BCo) & BCu);
	    Ago = BCo ^ ((~BCu) & BCa);
	    Agu = BCu ^ ((~BCa) & BCe);

	    Ebe ^= De;
	    BCa = ROL(Ebe, 1);
	    Egi ^= Di;
	    BCe = ROL(Egi, 6);
	    Eko ^= Do;
	    BCi = ROL(Eko, 25);
	    Emu ^= Du;
	    BCo = ROL(Emu, 8);
	    Esa ^= Da;
	    BCu = ROL(Esa, 18);
	    Aka = BCa ^ ((~BCe) & BCi);
	    Ake = BCe ^ ((~BCi) & BCo);
	    Aki = BCi ^ ((~BCo) & BCu);
	    Ako = BCo ^ ((~BCu) & BCa);
	    Aku = BCu ^ ((~BCa) & BCe);

	    Ebu ^= Du;
	    BCa = ROL(Ebu, 27);
	    Ega ^= Da;
	    BCe = ROL(Ega, 36);
	    Eke ^= De;
	    BCi = ROL(Eke, 10);
	    Emi ^= Di;
	    BCo = ROL(Emi, 15);
	    Eso ^= Do;
	    BCu = ROL(Eso, 56);
	    Ama = BCa ^ ((~BCe) & BCi);
	    Ame = BCe ^ ((~BCi) & BCo);
	    Ami = BCi ^ ((~BCo) & BCu);
	    Amo = BCo ^ ((~BCu) & BCa);
	    Amu = BCu ^ ((~BCa) & BCe);

	    Ebi ^= Di;
	    BCa = ROL(Ebi, 62);
	    Ego ^= Do;
	    BCe = ROL(Ego, 55);
	    Eku ^= Du;
	    BCi = ROL(Eku, 39);
	    Ema ^= Da;
	    BCo = ROL(Ema, 41);
	    Ese ^= De;
	    BCu = ROL(Ese, 2);
	    Asa = BCa ^ ((~BCe) & BCi);
	    Ase = BCe ^ ((~BCi) & BCo);
	    Asi = BCi ^ ((~BCo) & BCu);
	    Aso = BCo ^ ((~BCu) & BCa);
	    Asu = BCu ^ ((~BCa) & BCe);
	}

	// copyToState(state, A)
	state[0] = Aba;
	state[1] = Abe;
	state[2] = Abi;
	state[3] = Abo;
	state[4] = Abu;
	state[5] = Aga;
	state[6] = Age;
	state[7] = Agi;
	state[8] = Ago;
	state[9] = Agu;
	state[10] = Aka;
	state[11] = Ake;
	state[12] = Aki;
	state[13] = Ako;
	state[14] = Aku;
	state[15] = Ama;
	state[16] = Ame;
	state[17] = Ami;
	state[18] = Amo;
	state[19] = Amu;
	state[20] = Asa;
	state[21] = Ase;
	state[22] = Asi;
	state[23] = Aso;
	state[24] = Asu;
}

/*************************************************
 * Name:        keccak_absorb
 *
 * Description: Absorb step of Keccak;
 *              non-incremental, starts by zeroeing the state.
 *
 * Arguments:   - uint64_t *s: pointer to (uninitialized) output Keccak state
 *              - uint32_t r: rate in bytes (e.g., 168 for SHAKE128)
 *              - const uint8_t *m: pointer to input to be absorbed into s
 *              - size_t mlen: length of input in bytes
 *              - uint8_t p: domain-separation byte for different
 *                                 Keccak-derived functions
 **************************************************/
static void keccak_absorb(uint64_t *s, uint32_t r, const uint8_t *m,
                          size_t mlen, uint8_t p) {
	size_t i;
	uint8_t t[200];

	/* Zero state */
	for (i = 0; i < 25; ++i) {
		s[i] = 0;
	}

	while (mlen >= r) {
		for (i = 0; i < r / 8; ++i) {
			s[i] ^= load64(m + 8 * i);
		}

		KeccakF1600_StatePermute(s);
		mlen -= r;
		m += r;
	}

	for (i = 0; i < r; ++i) {
		t[i] = 0;
	}
	for (i = 0; i < mlen; ++i) {
		t[i] = m[i];
	}
	t[i] = p;
	t[r - 1] |= 128;
	for (i = 0; i < r / 8; ++i) {
		s[i] ^= load64(t + 8 * i);
	}
}

/*************************************************
 * Name:        keccak_squeezeblocks
 *
 * Description: Squeeze step of Keccak. Squeezes full blocks of r bytes each.
 *              Modifies the state. Can be called multiple times to keep
 *              squeezing, i.e., is incremental.
 *
 * Arguments:   - uint8_t *h: pointer to output blocks
 *              - size_t nblocks: number of blocks to be
 *                                                squeezed (written to h)
 *              - uint64_t *s: pointer to input/output Keccak state
 *              - uint32_t r: rate in bytes (e.g., 168 for SHAKE128)
 **************************************************/
static void keccak_squeezeblocks(uint8_t *h, size_t nblocks,
                                 uint64_t *s, uint32_t r) {
	while (nblocks > 0) {
		KeccakF1600_StatePermute(s);
		for (size_t i = 0; i < (r >> 3); i++) {
			store64(h + 8 * i, s[i]);
		}
		h += r;
		nblocks--;
	}
}

/*************************************************
 * Name:        shake128_absorb
 *
 * Description: Absorb step of the SHAKE128 XOF.
 *              non-incremental, starts by zeroeing the state.
 *
 * Arguments:   - uint64_t *s: pointer to (uninitialized) output Keccak state
 *              - const uint8_t *input: pointer to input to be absorbed
 *                                            into s
 *              - size_t inlen: length of input in bytes
 **************************************************/
void shake128_absorb(shake128ctx *state, const uint8_t *input, size_t inlen) {
	state->ctx = malloc(PQC_SHAKECTX_BYTES);
	if (state->ctx == NULL) {
		exit(111);
	}
	keccak_absorb(state->ctx, SHAKE128_RATE, input, inlen, 0x1F);
}

/*************************************************
 * Name:        shake128_squeezeblocks
 *
 * Description: Squeeze step of SHAKE128 XOF. Squeezes full blocks of
 *              SHAKE128_RATE bytes each. Modifies the state. Can be called
 *              multiple times to keep squeezing, i.e., is incremental.
 *
 * Arguments:   - uint8_t *output: pointer to output blocks
 *              - size_t nblocks: number of blocks to be squeezed
 *                                            (written to output)
 *              - shake128ctx *state: pointer to input/output Keccak state
 **************************************************/
void shake128_squeezeblocks(uint8_t *output, size_t nblocks, shake128ctx *state) {
	keccak_squeezeblocks(output, nblocks, state->ctx, SHAKE128_RATE);
}

void shake128_ctx_clone(shake128ctx *dest, const shake128ctx *src) {
	dest->ctx = malloc(PQC_SHAKECTX_BYTES);
	if (dest->ctx == NULL) {
		exit(111);
	}
	memcpy(dest->ctx, src->ctx, PQC_SHAKECTX_BYTES);
}

/** Release the allocated state. Call only once. */
void shake128_ctx_release(shake128ctx *state) {
	free(state->ctx); // IGNORE free-check
}

/*************************************************
 * Name:        shake256_absorb
 *
 * Description: Absorb step of the SHAKE256 XOF.
 *              non-incremental, starts by zeroeing the state.
 *
 * Arguments:   - shake256ctx *state: pointer to (uninitialized) output Keccak state
 *              - const uint8_t *input: pointer to input to be absorbed
 *                                            into s
 *              - size_t inlen: length of input in bytes
 **************************************************/
void shake256_absorb(shake256ctx *state, const uint8_t *input, size_t inlen) {
	state->ctx = malloc(PQC_SHAKECTX_BYTES);
	if (state->ctx == NULL) {
		exit(111);
	}
	keccak_absorb(state->ctx, SHAKE256_RATE, input, inlen, 0x1F);
}

/*************************************************
 * Name:        shake256_squeezeblocks
 *
 * Description: Squeeze step of SHAKE256 XOF. Squeezes full blocks of
 *              SHAKE256_RATE bytes each. Modifies the state. Can be called
 *              multiple times to keep squeezing, i.e., is incremental.
 *
 * Arguments:   - uint8_t *output: pointer to output blocks
 *              - size_t nblocks: number of blocks to be squeezed
 *                                (written to output)
 *              - shake256ctx *state: pointer to input/output Keccak state
 **************************************************/
void shake256_squeezeblocks(uint8_t *output, size_t nblocks, shake256ctx *state) {
	keccak_squeezeblocks(output, nblocks, state->ctx, SHAKE256_RATE);
}

void shake256_ctx_clone(shake256ctx *dest, const shake256ctx *src) {
	dest->ctx = malloc(PQC_SHAKECTX_BYTES);
	if (dest->ctx == NULL) {
		exit(111);
	}
	memcpy(dest->ctx, src->ctx, PQC_SHAKECTX_BYTES);
}

/** Release the allocated state. Call only once. */
void shake256_ctx_release(shake256ctx *state) {
	free(state->ctx); // IGNORE free-check
}

/*************************************************
 * Name:        shake128
 *
 * Description: SHAKE128 XOF with non-incremental API
 *
 * Arguments:   - uint8_t *output: pointer to output
 *              - size_t outlen: requested output length in bytes
 *              - const uint8_t *input: pointer to input
 *              - size_t inlen: length of input in bytes
 **************************************************/
void shake128(uint8_t *output, size_t outlen,
              const uint8_t *input, size_t inlen) {
	size_t nblocks = outlen / SHAKE128_RATE;
	uint8_t t[SHAKE128_RATE];
	shake128ctx s;

	shake128_absorb(&s, input, inlen);
	shake128_squeezeblocks(output, nblocks, &s);

	output += nblocks * SHAKE128_RATE;
	outlen -= nblocks * SHAKE128_RATE;

	if (outlen) {
		shake128_squeezeblocks(t, 1, &s);
		for (size_t i = 0; i < outlen; ++i) {
			output[i] = t[i];
		}
	}
	shake128_ctx_release(&s);
}

/*************************************************
 * Name:        shake256
 *
 * Description: SHAKE256 XOF with non-incremental API
 *
 * Arguments:   - uint8_t *output: pointer to output
 *              - size_t outlen: requested output length in bytes
 *              - const uint8_t *input: pointer to input
 *              - size_t inlen: length of input in bytes
 **************************************************/
void shake256_kyber(uint8_t *output, size_t outlen,
              const uint8_t *input, size_t inlen) {
	size_t nblocks = outlen / SHAKE256_RATE;
	uint8_t t[SHAKE256_RATE];
	shake256ctx s;

	shake256_absorb(&s, input, inlen);
	shake256_squeezeblocks(output, nblocks, &s);

	output += nblocks * SHAKE256_RATE;
	outlen -= nblocks * SHAKE256_RATE;

	if (outlen) {
		shake256_squeezeblocks(t, 1, &s);
		for (size_t i = 0; i < outlen; ++i) {
			output[i] = t[i];
		}
	}
	shake256_ctx_release(&s);
}

/*************************************************
 * Name:        sha3_256
 *
 * Description: SHA3-256 with non-incremental API
 *
 * Arguments:   - uint8_t *output:      pointer to output
 *              - const uint8_t *input: pointer to input
 *              - size_t inlen:   length of input in bytes
 **************************************************/
void sha3_256(uint8_t *output, const uint8_t *input, size_t inlen) {
	uint64_t s[25];
	uint8_t t[SHA3_256_RATE];

	/* Absorb input */
	keccak_absorb(s, SHA3_256_RATE, input, inlen, 0x06);

	/* Squeeze output */
	keccak_squeezeblocks(t, 1, s, SHA3_256_RATE);

	for (size_t i = 0; i < 32; i++) {
		output[i] = t[i];
	}
}


/*************************************************
 * Name:        sha3_512
 *
 * Description: SHA3-512 with non-incremental API
 *
 * Arguments:   - uint8_t *output:      pointer to output
 *              - const uint8_t *input: pointer to input
 *              - size_t inlen:   length of input in bytes
 **************************************************/
void sha3_512(uint8_t *output, const uint8_t *input, size_t inlen) {
	uint64_t s[25];
	uint8_t t[SHA3_512_RATE];

	/* Absorb input */
	keccak_absorb(s, SHA3_512_RATE, input, inlen, 0x06);

	/* Squeeze output */
	keccak_squeezeblocks(t, 1, s, SHA3_512_RATE);

	for (size_t i = 0; i < 64; i++) {
		output[i] = t[i];
	}
}