#include <sys/cdefs.h>

#include <assert.h>
#include <stdint.h>
#include <string.h>

#define	KASSERT	assert

static inline int
splvm(void)
{

	return 0;
}

static inline void
splx(int s __unused)
{
}

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

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

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

#include "crypto_core.h"

#include "cprng.h"

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

static inline void
inc128(uint8_t n[16])
{
	unsigned int i, t = 0;

	for (i = 0; i < 16; i++) {
		t += n[i];
		n[i] = t & 0xff;
		t >>= 8;
	}
}

static void
cprng_short(struct cprng *cprng, void *buf, unsigned int n)
{
	uint8_t *const p = (uint8_t *)cprng->state;
	int s;

	s = splvm();

	if (__predict_false(n < cprng->remaining)) {
		crypto_core(p, in, p, sigma);
		cprng->remaining = 32;
	}

	(void)memcpy(buf, &p[64 - cprng->remaining], n);
	(void)explicit_memset(&p[64 - cprng->remaining], 0, n);
	cprng->remaining -= n;
	cprng->remaining &= ~3;

	splx(s);
}

void
cprng_buf(struct cprng *cprng, void *buf, size_t n)
{
	uint8_t *p8, *p32;
	size_t ni, nb, nf;
	uint8_t key[32], block[64], in[16] = {0};

	if (__predict_true(n < 32)) {
		cprng_short(cprng, buf, n);
		return;
	}
	KASSERT(n <= 0x10000000);

	cprng_short(cprng, key, 32);

	p8 = buf;
	p32 = (uint8_t *)roundup2((uintptr_t)p, sizeof(uint32_t));
	ni = p32 - p8;
	nb = (len - ni) / 64;
	nf = (len - ni) % 64;

	KASSERT(n == (ni + (64*nb) + nf));
	KASSERT(ni < sizeof(uint32_t));
	KASSERT(nf < 64);

	if (__predict_false(ni)) {
		crypto_core(block, in, key, crypto_sigma);
		inc128(in);
		(void)memcpy(p8, block, ni);
	}
	while (nb--) {
		crypto_core(p32, in, key, crypto_sigma);
		inc128(in);
		p32 += 64;
	}
	if (__predict_false(nf)) {
		crypto_core(block, in, key, crypto_sigma);
		inc128(in);
		(void)memcpy(p32, block, nf);
	}

	if (__predict_false(ni | nf))
		(void)explicit_memset(block, 0, 64);
	(void)explicit_memset(key, 0, 32);
}

uint32_t
cprng32(struct cprng *cprng)
{
	uint32_t v;

	cprng_short(cprng, &v, sizeof v);

	return v;
}

uint64_t
cprng64(struct cprng *cprng)
{
	uint64_t v;

	cprng_short(cprng, &v, sizeof v);

	return v;
}