#include <sys/param.h>

#include <string.h>

#include "cprng.h"

#define	crypto_core_ROUNDS	20

static inline uint32_t
rotate(uint32_t v, unsigned int c)
{

	return ((v << c) | (v >> (32 - c)));
}

#if 1
#include "salsa20.h"
#else
#include "chacha.h"
#endif

void *(*volatile memset_impl)(void *, int, size_t) = &memset;

static void *
explicit_memset(void *buf, int c, size_t n)
{

	return (*memset_impl)(buf, c, n);
}

void
cprng_seed(struct cprng *cprng, uint8_t seed[32])
{

	memcpy(cprng->state, seed, 32);
	cprng->cached = 0;
}

static inline void
inc64(unsigned char n[8])
{
	unsigned int i, u;

	for (i = 0, u = 1; i < 8; i++, u >>= 8) {
		u += n[i];
		n[i] = (u & 0xffU);
	}
}

static const unsigned char sigma[16] = "expand 32-byte k";

void
cprng_buf(struct cprng *cprng, void *buf, size_t n)
{
	uint8_t *const volatile p8 = buf;
	unsigned char in[16] = {0};
	unsigned char block[64];
	uint32_t *volatile p32;
	size_t ni, nchunk, nf;

	/* Service short requests by buffering a single call to the core.  */
	if (__predict_true(n <= 32)) {
		if (__predict_false(cprng->cached < n)) {
			crypto_core(cprng->state, in, cprng->state, sigma);
			cprng->cached = 32;
		}
		memcpy(buf, cprng->state + (64 - cprng->cached), n);
#if 0
		(void)explicit_memset(cprng->state + (64 - cprng->cached), 0,
		    n);
#endif
		cprng->cached = (cprng->cached - n) & ~3;
		return;
	}

	/* Align the input.  */
	p32 = (uint32_t *)roundup2((uintptr_t)p8, (uintptr_t)sizeof(uint32_t));
	ni = (uint8_t *)p32 - p8;

	/* Split the output into 64-bit chunks and a trailer.  */
	nchunk = ((n - ni) / 64);
	nf = ((n - ni) % 64);

	/* Generate the unaligned part of the buffer.  */
	if (__predict_false(ni)) {
		crypto_core(block, in, cprng->state, sigma);
		inc64(&in[8]);
		memcpy(p8, block, ni);
	}

	/* Generate 512-bit aligned chunks.  */
	while (nchunk--) {
		crypto_core((uint8_t *)p32, in, cprng->state, sigma);
		inc64(&in[8]);
		p32 += (64 / sizeof(uint32_t));
	}

	/* Generate any remaining <512-bit part.  */
	if (nf) {
		crypto_core(block, in, cprng->state, sigma);
		inc64(&in[8]);
		memcpy(p32, block, nf);
	}

	/* Clear the temporary output if we used it.  */
	if (__predict_false(ni || nf))
		(void)explicit_memset(block, 0, sizeof(block));

	/* Rekey.  */
	crypto_core(cprng->state, in, cprng->state, sigma);
	cprng->cached = 32;
}

inline uint32_t
cprng32(struct cprng *cprng)
{
	unsigned char in[16] = {0};
	uint32_t v;

	if (__predict_false(cprng->cached < sizeof(uint32_t))) {
		crypto_core(cprng->state, in, cprng->state, sigma);
		cprng->cached = 32;
	}

	v = *(uint32_t *)(cprng->state + (64 - cprng->cached));
	cprng->cached -= sizeof(uint32_t);

	return v;
}

uint64_t
cprng64(struct cprng *cprng)
{

	return ((uint64_t)cprng32(cprng) << 32) | cprng32(cprng);
}