#include <sys/cdefs.h>

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

#define	KASSERT	assert

__noinline
static int
splvm(void)
{

	return 0;
}

__noinline
static void
splx(int s)
{
	volatile int v __unused = s;

}

static struct cprng *volatile cprng_percpu;

__noinline
static struct cprng *
cprng_percpu_get(struct cprng *cprng)
{

	cprng_percpu = cprng;
	return cprng_percpu;
}

__noinline
static void
cprng_percpu_put(struct cprng *cprng)
{

	cprng_percpu = cprng;
}

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);
}

static void *cprng_softint = NULL;

static void __noinline
softint_schedule(void *cookie)
{
	void *volatile v __unused = cookie;

}

#include "crypto_core.h"

#include "cprng.h"

#define	CPRNG_KEYSTREAM_MAX	UINT_MAX

/* `expand 32-byte k' */
static const uint32_t sigma[4] = {
	0x61707865U, 0x3320646eU, 0x79622d32U, 0x6b206574U,
};

static void
inc128(uint32_t n[4])
{
	uint64_t t = 1;
	unsigned int i;

	for (i = 0; i < 4; i++) {
		t += n[i];
		n[i] = t & 0xffffffffU;
		t >>= 32;
	}
}

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

	(void)memcpy(cprng->key, seed, sizeof(cprng->key));
	(void)memset(cprng->nonce, 0, sizeof(cprng->nonce));
	cprng->buffered = 0;
}

#define	roundup2(x, n)	((((x) - 1) | ((n) - 1)) + 1)
#define	howmany(x, n)	(((x) + ((n) - 1)) / (n))

static inline void
cprng_short(struct cprng *cprng, void *buf, unsigned int n)
{
	int s;

	cprng = cprng_percpu_get(cprng);
	s = splvm();

	if (__predict_false(cprng->buffered < n)) {
		crypto_core(cprng->buffer, cprng->nonce, cprng->key, sigma);
		inc128(cprng->nonce);
		if (__predict_false(cprng->nonce[1] != 0))
			softint_schedule(cprng_softint);
		cprng->buffered = 16;
	}

	(void)memcpy(buf, &cprng->buffer[16 - cprng->buffered], n);
	cprng->buffered -= howmany(n, 4);

	splx(s);
	cprng_percpu_put(cprng);
}

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

	if (__predict_true(n <= 64)) {
		cprng_short(cprng, buf, n);
		return;
	}
	KASSERT(n <= CPRNG_MAX);

	cprng_short(cprng, key, sizeof key);

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

	KASSERT(((uintptr_t)p32 & 3) == 0);
	KASSERT(n == (ni + (64*nb) + nf));
	KASSERT(ni < sizeof(uint32_t));
	KASSERT(nf < 64);

	if (__predict_false(ni)) {
		crypto_core(block, nonce, key, sigma);
		inc128(nonce);
		(void)memcpy(p8, block, ni);
	}
	while (nb--) {
		crypto_core(p32, nonce, key, sigma);
		inc128(nonce);
		p32 += 16;
	}
	if (__predict_false(nf)) {
		crypto_core(block, nonce, key, sigma);
		inc128(nonce);
		(void)memcpy(p32, block, nf);
	}

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

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;
}