/*

 * Copyright (C) 2012 Vincent Hanquez <vincent@snarc.org>

 *

 * Redistribution and use in source and binary forms, with or without

 * modification, are permitted provided that the following conditions

 * are met:

 * 1. Redistributions of source code must retain the above copyright

 *    notice, this list of conditions and the following disclaimer.

 * 2. Redistributions in binary form must reproduce the above copyright

 *    notice, this list of conditions and the following disclaimer in the

 *    documentation and/or other materials provided with the distribution.

 *

 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR

 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES

 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.

 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,

 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT

 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,

 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY

 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT

 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF

 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 */

#include <stdint.h>

#include <string.h>

#include "cryptonite_bitfn.h"

#include "cryptonite_align.h"

#include "cryptonite_sha3.h"

#define KECCAK_NB_ROUNDS 24

/* rounds constants */

static const uint64_t keccak_rndc[24] =

{

	0x0000000000000001ULL, 0x0000000000008082ULL, 0x800000000000808aULL,

	0x8000000080008000ULL, 0x000000000000808bULL, 0x0000000080000001ULL,

	0x8000000080008081ULL, 0x8000000000008009ULL, 0x000000000000008aULL,

	0x0000000000000088ULL, 0x0000000080008009ULL, 0x000000008000000aULL,

	0x000000008000808bULL, 0x800000000000008bULL, 0x8000000000008089ULL,

	0x8000000000008003ULL, 0x8000000000008002ULL, 0x8000000000000080ULL,

	0x000000000000800aULL, 0x800000008000000aULL, 0x8000000080008081ULL,

	0x8000000000008080ULL, 0x0000000080000001ULL, 0x8000000080008008ULL,

};

/* triangular numbers constants */

static const int keccak_rotc[24] =

	{ 1,3,6,10,15,21,28,36,45,55,2,14,27,41,56,8,25,43,62,18,39,61,20,44 };

static const int keccak_piln[24] =

	{ 10,7,11,17,18,3,5,16,8,21,24,4,15,23,19,13,12,2,20,14,22,9,6,1 };

static inline void sha3_do_chunk(uint64_t state[25], uint64_t buf[], int bufsz)

{

	int i, j, r;

	uint64_t tmp, bc[5];

	/* merge buf with state */

	for (i = 0; i < bufsz; i++)

		state[i] ^= le64_to_cpu(buf[i]);

	/* run keccak rounds */

	for (r = 0; r < KECCAK_NB_ROUNDS; r++) {

		/* compute the parity of each columns */

		for (i = 0; i < 5; i++)

			bc[i] = state[i] ^ state[i+5] ^ state[i+10] ^ state[i+15] ^ state[i+20];

		for (i = 0; i < 5; i++) {

			tmp = bc[(i + 4) % 5] ^ rol64(bc[(i + 1) % 5], 1);

			for (j = 0; j < 25; j += 5)

				state[j + i] ^= tmp;

		}

		/* rho pi */

		tmp = state[1];

		for (i = 0; i < 24; i++) {

			j = keccak_piln[i];

			bc[0] = state[j];

			state[j] = rol64(tmp, keccak_rotc[i]);

			tmp = bc[0];

		}

		/* bitwise combine along rows using a = a xor (not b and c) */

		for (j = 0; j < 25; j += 5) {

			for (i = 0; i < 5; i++)

				bc[i] = state[j + i];

			#define andn(b,c) (~(b) & (c))

			state[j + 0] ^= andn(bc[1], bc[2]);

			state[j + 1] ^= andn(bc[2], bc[3]);

			state[j + 2] ^= andn(bc[3], bc[4]);

			state[j + 3] ^= andn(bc[4], bc[0]);

			state[j + 4] ^= andn(bc[0], bc[1]);

			#undef andn

		}

		/* xor the round constant */

		state[0] ^= keccak_rndc[r];

	}

}

/*

 * Initialize a SHA-3 / SHAKE context: hashlen is the security level (and

 * half the capacity) in bits

 */

void cryptonite_sha3_init(struct sha3_ctx *ctx, uint32_t hashlen)

{

	/* assert(hashlen >= SHA3_BITSIZE_MIN && hashlen <= SHA3_BITSIZE_MAX) */

	int bufsz = SHA3_BUF_SIZE(hashlen);

	memset(ctx, 0, sizeof(*ctx) + bufsz);

	ctx->bufsz = bufsz;

}

/* Update a SHA-3 / SHAKE context */

void cryptonite_sha3_update(struct sha3_ctx *ctx, const uint8_t *data, uint32_t len)

{

	uint32_t to_fill;

	to_fill = ctx->bufsz - ctx->bufindex;

	if (ctx->bufindex == ctx->bufsz) {

		sha3_do_chunk(ctx->state, (uint64_t *) ctx->buf, ctx->bufsz / 8);

		ctx->bufindex = 0;

	}

	/* process partial buffer if there's enough data to make a block */

	if (ctx->bufindex && len >= to_fill) {

		memcpy(ctx->buf + ctx->bufindex, data, to_fill);

		sha3_do_chunk(ctx->state, (uint64_t *) ctx->buf, ctx->bufsz / 8);

		len -= to_fill;

		data += to_fill;

		ctx->bufindex = 0;

	}

	if (need_alignment(data, 8)) {

		uint64_t tramp[SHA3_BUF_SIZE_MAX/8];

		ASSERT_ALIGNMENT(tramp, 8);

		for (; len >= ctx->bufsz; len -= ctx->bufsz, data += ctx->bufsz) {

			memcpy(tramp, data, ctx->bufsz);

			sha3_do_chunk(ctx->state, tramp, ctx->bufsz / 8);

		}

	} else {

		/* process as much ctx->bufsz-block */

		for (; len >= ctx->bufsz; len -= ctx->bufsz, data += ctx->bufsz)

			sha3_do_chunk(ctx->state, (uint64_t *) data, ctx->bufsz / 8);

	}

	/* append data into buf */

	if (len) {

		memcpy(ctx->buf + ctx->bufindex, data, len);

		ctx->bufindex += len;

	}

}

void cryptonite_sha3_finalize_with_pad_byte(struct sha3_ctx *ctx, uint8_t pad_byte)

{

	/* process full buffer if needed */

	if (ctx->bufindex == ctx->bufsz) {

		sha3_do_chunk(ctx->state, (uint64_t *) ctx->buf, ctx->bufsz / 8);

		ctx->bufindex = 0;

	}

	/* add the 10*1 padding */

	ctx->buf[ctx->bufindex++] = pad_byte;

	memset(ctx->buf + ctx->bufindex, 0, ctx->bufsz - ctx->bufindex);

	ctx->buf[ctx->bufsz - 1] |= 0x80;

	/* process */

	sha3_do_chunk(ctx->state, (uint64_t *) ctx->buf, ctx->bufsz / 8);

	ctx->bufindex = 0;

}

/*

 * Extract some bytes from a finalized SHA-3 / SHAKE context.

 * May be called multiple times.

 */

void cryptonite_sha3_output(struct sha3_ctx *ctx, uint8_t *out, uint32_t len)

{

	uint64_t w[25];

	uint8_t *wptr = (uint8_t *) w;

	uint32_t still_avail;

	still_avail = ctx->bufsz - ctx->bufindex;

	if (ctx->bufindex == ctx->bufsz) {

		/* squeeze the sponge again, without any input */

		sha3_do_chunk(ctx->state, NULL, 0);

		ctx->bufindex = 0;

	}

	/* use bytes already available if this block is fully consumed */

	if (ctx->bufindex && len >= still_avail) {

		cpu_to_le64_array(w, ctx->state, 25);

		memcpy(out, wptr + ctx->bufindex, still_avail);

		sha3_do_chunk(ctx->state, NULL, 0);

		len -= still_avail;

		out += still_avail;

		ctx->bufindex = 0;

	}

	/* output as much ctx->bufsz-block */

	for (; len > ctx->bufsz; len -= ctx->bufsz, out += ctx->bufsz) {

		cpu_to_le64_array(w, ctx->state, 25);

		memcpy(out, w, ctx->bufsz);

		sha3_do_chunk(ctx->state, NULL, 0);

	}

	/* output from partial buffer */

	if (len) {

		cpu_to_le64_array(w, ctx->state, 25);

		memcpy(out, wptr + ctx->bufindex, len);

		ctx->bufindex += len;

	}

}

/* Finalize a SHA-3 context and return the digest value */

void cryptonite_sha3_finalize(struct sha3_ctx *ctx, uint32_t hashlen, uint8_t *out)

{

	cryptonite_sha3_finalize_with_pad_byte(ctx, 0x06);

	cryptonite_sha3_output(ctx, out, hashlen / 8);

}

/* Finalize a SHAKE context. Output is read using cryptonite_sha3_output. */

void cryptonite_sha3_finalize_shake(struct sha3_ctx *ctx)

{

	cryptonite_sha3_finalize_with_pad_byte(ctx, 0x1F);

}

void cryptonite_keccak_init(struct sha3_ctx *ctx, uint32_t hashlen)

{

	cryptonite_sha3_init(ctx, hashlen);

}

void cryptonite_keccak_update(struct sha3_ctx *ctx, uint8_t *data, uint32_t len)

{

	cryptonite_sha3_update(ctx, data, len);

}

void cryptonite_keccak_finalize(struct sha3_ctx *ctx, uint32_t hashlen, uint8_t *out)

{

	cryptonite_sha3_finalize_with_pad_byte(ctx, 1);

	cryptonite_sha3_output(ctx, out, hashlen / 8);

}