/*	$OpenBSD: wg_noise.c,v 1.5 2021/03/21 18:13:59 sthen Exp $ */
/*
 * Copyright (C) 2015-2020 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 * Copyright (C) 2019-2021 Matt Dunwoodie <ncon@noconroy.net>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/systm.h>
#include <sys/param.h>
#include <sys/atomic.h>
#include <sys/rwlock.h>
#include <sys/malloc.h>
#include <sys/smr.h>

#include <crypto/blake2s.h>
#include <crypto/curve25519.h>
#include <crypto/chachapoly.h>
#include <crypto/siphash.h>

#include <net/wg_noise.h>

/* Protocol string constants */
#define NOISE_HANDSHAKE_NAME	"Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s"
#define NOISE_IDENTIFIER_NAME	"WireGuard v1 zx2c4 Jason@zx2c4.com"

/* Constants for the counter */
#define COUNTER_BITS_TOTAL	8192
#define COUNTER_BITS		(sizeof(unsigned long) * 8)
#define COUNTER_NUM		(COUNTER_BITS_TOTAL / COUNTER_BITS)
#define COUNTER_WINDOW_SIZE	(COUNTER_BITS_TOTAL - COUNTER_BITS)

/* Constants for the keypair */
#define REKEY_AFTER_MESSAGES	(1ull << 60)
#define REJECT_AFTER_MESSAGES	(UINT64_MAX - COUNTER_WINDOW_SIZE - 1)
#define REKEY_AFTER_TIME	120
#define REKEY_AFTER_TIME_RECV	165
#define REJECT_INTERVAL		(1000000000 / 50) /* fifty times per sec */
/* 24 = floor(log2(REJECT_INTERVAL)) */
#define REJECT_INTERVAL_MASK	(~((1ull<<24)-1))
#define TIMER_RESET		(struct timespec){ -(REKEY_TIMEOUT+1), 0 }

#define HT_INDEX_SIZE		(1 << 13)
#define HT_INDEX_MASK		(HT_INDEX_SIZE - 1)
#define HT_REMOTE_SIZE		(1 << 11)
#define HT_REMOTE_MASK		(HT_REMOTE_SIZE - 1)
#define MAX_REMOTE_PER_LOCAL	(1 << 20)

struct noise_index {
	SMR_LIST_ENTRY(noise_index)	 i_entry;
	uint32_t			 i_local_index;
	uint32_t			 i_remote_index;
	int				 i_is_keypair;
};

struct noise_keypair {
	struct noise_index		 kp_index;
	struct refcnt			 kp_refcnt;
	int				 kp_can_send;
	int				 kp_is_initiator;
	struct timespec			 kp_birthdate; /* nanouptime */
	struct noise_remote		*kp_remote;

	uint8_t				 kp_send[NOISE_SYMMETRIC_KEY_LEN];
	uint8_t				 kp_recv[NOISE_SYMMETRIC_KEY_LEN];

	/* Counter elements */
	struct rwlock			 kp_nonce_lock;
	uint64_t			 kp_nonce_send;
	uint64_t			 kp_nonce_recv;
	unsigned long			 kp_backtrack[COUNTER_NUM];

	struct smr_entry		 kp_smr;
};

struct noise_handshake {
	uint8_t	 			 hs_e[NOISE_PUBLIC_KEY_LEN];
	uint8_t	 			 hs_hash[NOISE_HASH_LEN];
	uint8_t	 			 hs_ck[NOISE_HASH_LEN];
};

struct noise_remote {
	struct noise_index		 r_index;

	SMR_LIST_ENTRY(noise_remote) 	 r_entry;
	uint8_t				 r_public[NOISE_PUBLIC_KEY_LEN];

	struct rwlock			 r_handshake_lock;
	struct noise_handshake		 r_handshake;
	int				 r_handshake_alive;
	int				 r_handshake_initiator;
	struct timespec			 r_last_sent; /* nanouptime */
	struct timespec			 r_last_init_recv; /* nanouptime */
	uint8_t				 r_timestamp[NOISE_TIMESTAMP_LEN];
	uint8_t				 r_psk[NOISE_SYMMETRIC_KEY_LEN];
	uint8_t		 		 r_ss[NOISE_PUBLIC_KEY_LEN];

	struct refcnt			 r_refcnt;
	struct noise_local		*r_local;
	void				*r_arg;

	struct rwlock			 r_keypair_lock;
	struct noise_keypair		*r_next, *r_current, *r_previous;

	struct smr_entry		 r_smr;
	void				(*r_cleanup)(struct noise_remote *);
};

struct noise_local {
	struct rwlock			 l_identity_lock;
	int				 l_has_identity;
	uint8_t				 l_public[NOISE_PUBLIC_KEY_LEN];
	uint8_t				 l_private[NOISE_PUBLIC_KEY_LEN];

	struct refcnt			 l_refcnt;
	SIPHASH_KEY			 l_hash_key;
	void				*l_arg;
	void				(*l_cleanup)(struct noise_local *);

	struct rwlock			 l_remote_lock;
	size_t				 l_remote_num;
	SMR_LIST_HEAD(,noise_remote)	 l_remote_hash[HT_REMOTE_SIZE];

	struct rwlock			 l_index_lock;
	SMR_LIST_HEAD(,noise_index)	 l_index_hash[HT_INDEX_SIZE];
};

static void	noise_precompute_ss(struct noise_local *, struct noise_remote *);

static void	noise_remote_index_insert(struct noise_local *, struct noise_remote *);
static int	noise_remote_index_remove(struct noise_local *, struct noise_remote *);
static void	noise_remote_expire_current(struct noise_remote *);


static void	noise_add_new_keypair(struct noise_local *, struct noise_remote *, struct noise_keypair *);
static int	noise_received_with(struct noise_keypair *);
static int	noise_begin_session(struct noise_remote *);
static void	noise_keypair_drop(struct noise_keypair *);

static void	noise_kdf(uint8_t *, uint8_t *, uint8_t *, const uint8_t *,
		    size_t, size_t, size_t, size_t,
		    const uint8_t [NOISE_HASH_LEN]);
static int	noise_mix_dh(uint8_t [NOISE_HASH_LEN], uint8_t [NOISE_SYMMETRIC_KEY_LEN],
		    const uint8_t [NOISE_PUBLIC_KEY_LEN],
		    const uint8_t [NOISE_PUBLIC_KEY_LEN]);
static int	noise_mix_ss(uint8_t ck[NOISE_HASH_LEN], uint8_t [NOISE_SYMMETRIC_KEY_LEN],
		    const uint8_t [NOISE_PUBLIC_KEY_LEN]);
static void	noise_mix_hash(uint8_t [NOISE_HASH_LEN], const uint8_t *, size_t);
static void	noise_mix_psk(uint8_t [NOISE_HASH_LEN], uint8_t [NOISE_HASH_LEN],
		    uint8_t [NOISE_SYMMETRIC_KEY_LEN], const uint8_t [NOISE_SYMMETRIC_KEY_LEN]);
static void	noise_param_init(uint8_t [NOISE_HASH_LEN], uint8_t [NOISE_HASH_LEN],
		    const uint8_t [NOISE_PUBLIC_KEY_LEN]);
static void	noise_msg_encrypt(uint8_t *, const uint8_t *, size_t,
		    uint8_t [NOISE_SYMMETRIC_KEY_LEN], uint8_t [NOISE_HASH_LEN]);
static int	noise_msg_decrypt(uint8_t *, const uint8_t *, size_t,
		    uint8_t [NOISE_SYMMETRIC_KEY_LEN], uint8_t [NOISE_HASH_LEN]);
static void	noise_msg_ephemeral(uint8_t [NOISE_HASH_LEN], uint8_t [NOISE_HASH_LEN],
		    const uint8_t [NOISE_PUBLIC_KEY_LEN]);
static void	noise_tai64n_now(uint8_t [NOISE_TIMESTAMP_LEN]);
static int	noise_timer_expired(struct timespec *, time_t, long);

/* Make rwlock spin locks */
#define rw_enter_write_spin(l) while (rw_enter(l, RW_WRITE | RW_NOSLEEP) != 0)
#define rw_enter_read_spin(l) while (rw_enter(l, RW_READ | RW_NOSLEEP) != 0)

/* Local configuration */
struct noise_local *
noise_local_alloc(void *arg)
{
	struct noise_local *l;
	size_t i;

	if ((l = malloc(sizeof(*l), M_DEVBUF, M_NOWAIT)) == NULL)
		return NULL;

	rw_init(&l->l_identity_lock, "noise_identity");
	l->l_has_identity = 0;
	bzero(l->l_public, NOISE_PUBLIC_KEY_LEN);
	bzero(l->l_private, NOISE_PUBLIC_KEY_LEN);

	refcnt_init(&l->l_refcnt);
	arc4random_buf(&l->l_hash_key, sizeof(l->l_hash_key));
	l->l_arg = arg;
	l->l_cleanup = NULL;

	rw_init(&l->l_remote_lock, "noise_remote");
	l->l_remote_num = 0;
	for (i = 0; i < HT_REMOTE_SIZE; i++)
		SMR_LIST_INIT(&l->l_remote_hash[i]);

	rw_init(&l->l_index_lock, "noise_index");
	for (i = 0; i < HT_INDEX_SIZE; i++)
		SMR_LIST_INIT(&l->l_index_hash[i]);

	return l;
}

struct noise_local *
noise_local_ref(struct noise_local *l)
{
	refcnt_take(&l->l_refcnt);
	return l;
}

void
noise_local_put(struct noise_local *l)
{
	if (refcnt_rele(&l->l_refcnt)) {
		if (l->l_cleanup != NULL)
			l->l_cleanup(l);
		explicit_bzero(l, sizeof(*l));
		free(l, M_DEVBUF, sizeof(*l));
	}
}

void
noise_local_free(struct noise_local *l, void (*cleanup)(struct noise_local *))
{
	l->l_cleanup = cleanup;
	noise_local_put(l);
}

void *
noise_local_arg(struct noise_local *l)
{
	return l->l_arg;
}

void
noise_local_private(struct noise_local *l, const uint8_t private[NOISE_PUBLIC_KEY_LEN])
{
	struct noise_remote *r;
	size_t i;

	rw_enter_write_spin(&l->l_identity_lock);
	memcpy(l->l_private, private, NOISE_PUBLIC_KEY_LEN);
	curve25519_clamp_secret(l->l_private);
	l->l_has_identity = curve25519_generate_public(l->l_public, l->l_private);

	smr_read_enter();
	for (i = 0; i < HT_REMOTE_SIZE; i++) {
		SMR_LIST_FOREACH(r, &l->l_remote_hash[i], r_entry) {
			noise_precompute_ss(l, r);
			noise_remote_expire_current(r);
		}
	}
	smr_read_leave();
	rw_exit_write(&l->l_identity_lock);
}

int
noise_local_keys(struct noise_local *l, uint8_t public[NOISE_PUBLIC_KEY_LEN],
    uint8_t private[NOISE_PUBLIC_KEY_LEN])
{
	int has_identity;
	rw_enter_read_spin(&l->l_identity_lock);
	if ((has_identity = l->l_has_identity)) {
		if (public != NULL)
			memcpy(public, l->l_public, NOISE_PUBLIC_KEY_LEN);
		if (private != NULL)
			memcpy(private, l->l_private, NOISE_PUBLIC_KEY_LEN);
	}
	rw_exit_read(&l->l_identity_lock);
	return has_identity ? 0 : ENXIO;
}

static void
noise_precompute_ss(struct noise_local *l, struct noise_remote *r)
{
	rw_enter_write_spin(&r->r_handshake_lock);
	if (!l->l_has_identity ||
	    !curve25519(r->r_ss, l->l_private, r->r_public))
		bzero(r->r_ss, NOISE_PUBLIC_KEY_LEN);
	rw_exit_write(&r->r_handshake_lock);
}

/* Remote configuration */
struct noise_remote *
noise_remote_alloc(struct noise_local *l, void *arg,
    const uint8_t public[NOISE_PUBLIC_KEY_LEN],
    const uint8_t psk[NOISE_PUBLIC_KEY_LEN])
{
	struct noise_remote *r, *ri;
	uint64_t idx;

	if ((r = malloc(sizeof(*r), M_DEVBUF, M_NOWAIT)) == NULL)
		return NULL;

	r->r_index.i_is_keypair = 0;

	memcpy(r->r_public, public, NOISE_PUBLIC_KEY_LEN);

	rw_init(&r->r_handshake_lock, "noise_handshake");
	bzero(&r->r_handshake, sizeof(r->r_handshake));
	r->r_handshake_alive = 0;
	r->r_handshake_initiator = 0;
	r->r_last_sent = TIMER_RESET;
	r->r_last_init_recv = TIMER_RESET;
	bzero(r->r_timestamp, NOISE_TIMESTAMP_LEN);
	noise_remote_set_psk(r, psk);
	noise_precompute_ss(l, r);

	refcnt_init(&r->r_refcnt);
	r->r_local = noise_local_ref(l);
	r->r_arg = arg;

	rw_init(&r->r_keypair_lock, "noise_keypair");
	r->r_next = r->r_current = r->r_previous = NULL;

	smr_init(&r->r_smr);

	/* Insert to hashtable */
	idx = SipHash24(&l->l_hash_key, public, NOISE_PUBLIC_KEY_LEN) & HT_REMOTE_MASK;

	rw_enter_write_spin(&l->l_remote_lock);
	SMR_LIST_FOREACH_LOCKED(ri, &l->l_remote_hash[idx], r_entry)
		if (timingsafe_bcmp(ri->r_public, public, NOISE_PUBLIC_KEY_LEN) == 0)
			goto free;
	if (l->l_remote_num < MAX_REMOTE_PER_LOCAL) {
		l->l_remote_num++;
		SMR_LIST_INSERT_HEAD_LOCKED(&l->l_remote_hash[idx], r, r_entry);
	} else {
free:
		free(r, M_DEVBUF, sizeof(*r));
		noise_local_put(l);
		r = NULL;
	}
	rw_exit_write(&l->l_remote_lock);

	return r;
}

struct noise_remote *
noise_remote_lookup(struct noise_local *l, const uint8_t public[NOISE_PUBLIC_KEY_LEN])
{
	struct noise_remote *r, *ret = NULL;
	uint64_t idx;

	idx = SipHash24(&l->l_hash_key, public, NOISE_PUBLIC_KEY_LEN) & HT_REMOTE_MASK;

	smr_read_enter();
	SMR_LIST_FOREACH(r, &l->l_remote_hash[idx], r_entry) {
		if (timingsafe_bcmp(r->r_public, public, NOISE_PUBLIC_KEY_LEN) == 0) {
			if (refcnt_take_if_gt(&r->r_refcnt, 0))
				ret = r;
			break;
		}
	}
	smr_read_leave();
	return ret;
}

static void
noise_remote_index_insert(struct noise_local *l, struct noise_remote *r)
{
	struct noise_index *i, *r_i = &r->r_index;
	uint32_t idx;

	noise_remote_index_remove(l, r);

	rw_enter_write_spin(&l->l_index_lock);
assign_id:
	r_i->i_local_index = arc4random();
	idx = r_i->i_local_index & HT_INDEX_MASK;
	SMR_LIST_FOREACH_LOCKED(i, &l->l_index_hash[idx], i_entry)
		if (i->i_local_index == r_i->i_local_index)
			goto assign_id;

	SMR_LIST_INSERT_HEAD_LOCKED(&l->l_index_hash[idx], r_i, i_entry);
	rw_exit_write(&l->l_index_lock);

	r->r_handshake_alive = 1;
}

struct noise_remote *
noise_remote_index_lookup(struct noise_local *l, uint32_t idx0)
{
	struct noise_index *i;
	struct noise_remote *r, *ret = NULL;
	uint32_t idx = idx0 & HT_INDEX_MASK;

	smr_read_enter();
	SMR_LIST_FOREACH(i, &l->l_index_hash[idx], i_entry) {
		if (i->i_local_index == idx0 && !i->i_is_keypair) {
			r = (struct noise_remote *) i;
			if (refcnt_take_if_gt(&r->r_refcnt, 0))
				ret = r;
			break;
		}
	}
	smr_read_leave();
	return ret;
}

static int
noise_remote_index_remove(struct noise_local *l, struct noise_remote *r)
{
	rw_assert_wrlock(&r->r_handshake_lock);
	if (r->r_handshake_alive) {
		rw_enter_write_spin(&l->l_index_lock);
		SMR_LIST_REMOVE_LOCKED(&r->r_index, i_entry);
		rw_exit_write(&l->l_index_lock);
		r->r_handshake_alive = 0;
		return 1;
	}
	return 0;
}

struct noise_remote *
noise_remote_ref(struct noise_remote *r)
{
	refcnt_take(&r->r_refcnt);
	return r;
}

static void
noise_remote_smr_free(void *_r)
{
	struct noise_remote *r = _r;
	if (r->r_cleanup != NULL)
		r->r_cleanup(r);
	noise_local_put(r->r_local);
	explicit_bzero(r, sizeof(*r));
	free(r, M_DEVBUF, sizeof(*r));
}

void
noise_remote_put(struct noise_remote *r)
{
	if (refcnt_rele(&r->r_refcnt))
		smr_call(&r->r_smr, noise_remote_smr_free, r);
}

void
noise_remote_free(struct noise_remote *r, void (*cleanup)(struct noise_remote *))
{
	struct noise_local *l = r->r_local;

	r->r_cleanup = cleanup;

	/* remove from hashtable */
	rw_enter_write_spin(&l->l_remote_lock);
	SMR_LIST_REMOVE_LOCKED(r, r_entry);
	l->l_remote_num--;
	rw_exit_write(&l->l_remote_lock);

	/* now clear all keypairs and handshakes, then put this reference */
	noise_remote_handshake_clear(r);
	noise_remote_keypairs_clear(r);
	noise_remote_put(r);
}

struct noise_local *
noise_remote_local(struct noise_remote *r)
{
	return noise_local_ref(r->r_local);
}

void *
noise_remote_arg(struct noise_remote *r)
{
	return r->r_arg;
}

void
noise_remote_set_psk(struct noise_remote *r,
    const uint8_t psk[NOISE_SYMMETRIC_KEY_LEN])
{
	rw_enter_write_spin(&r->r_handshake_lock);
	if (psk == NULL)
		bzero(r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
	else
		memcpy(r->r_psk, psk, NOISE_SYMMETRIC_KEY_LEN);
	rw_exit_write(&r->r_handshake_lock);
}

int
noise_remote_keys(struct noise_remote *r, uint8_t public[NOISE_PUBLIC_KEY_LEN],
    uint8_t psk[NOISE_SYMMETRIC_KEY_LEN])
{
	static uint8_t null_psk[NOISE_SYMMETRIC_KEY_LEN];
	int ret;

	if (public != NULL)
		memcpy(public, r->r_public, NOISE_PUBLIC_KEY_LEN);

	rw_enter_read_spin(&r->r_handshake_lock);
	if (psk != NULL)
		memcpy(psk, r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
	ret = timingsafe_bcmp(r->r_psk, null_psk, NOISE_SYMMETRIC_KEY_LEN);
	rw_exit_read(&r->r_handshake_lock);

	return ret ? 0 : ENOENT;
}

int
noise_remote_initiation_expired(struct noise_remote *r)
{
	int expired;
	rw_enter_read_spin(&r->r_handshake_lock);
	expired = noise_timer_expired(&r->r_last_sent, REKEY_TIMEOUT, 0);
	rw_exit_read(&r->r_handshake_lock);
	return expired;
}

void
noise_remote_handshake_clear(struct noise_remote *r)
{
	rw_enter_write_spin(&r->r_handshake_lock);
	if (noise_remote_index_remove(r->r_local, r))
		bzero(&r->r_handshake, sizeof(r->r_handshake));
	r->r_last_sent = TIMER_RESET;
	rw_exit_write(&r->r_handshake_lock);
}

void
noise_remote_keypairs_clear(struct noise_remote *r)
{
	struct noise_keypair *kp;

	rw_enter_write_spin(&r->r_keypair_lock);
	kp = SMR_PTR_GET_LOCKED(&r->r_next);
	SMR_PTR_SET_LOCKED(&r->r_next, NULL);
	noise_keypair_drop(kp);

	kp = SMR_PTR_GET_LOCKED(&r->r_current);
	SMR_PTR_SET_LOCKED(&r->r_current, NULL);
	noise_keypair_drop(kp);

	kp = SMR_PTR_GET_LOCKED(&r->r_previous);
	SMR_PTR_SET_LOCKED(&r->r_previous, NULL);
	noise_keypair_drop(kp);
	rw_exit_write(&r->r_keypair_lock);
}

static void
noise_remote_expire_current(struct noise_remote *r)
{
	struct noise_keypair *kp;

	noise_remote_handshake_clear(r);

	smr_read_enter();
	kp = SMR_PTR_GET(&r->r_next);
	if (kp != NULL) WRITE_ONCE(kp->kp_can_send, 0);
	kp = SMR_PTR_GET(&r->r_current);
	if (kp != NULL) WRITE_ONCE(kp->kp_can_send, 0);
	smr_read_leave();
}

/* Keypair functions */
static void
noise_add_new_keypair(struct noise_local *l, struct noise_remote *r,
    struct noise_keypair *kp)
{
	struct noise_keypair *next, *current, *previous;
	struct noise_index *r_i = &r->r_index;

	/* Insert into the keypair table */
	rw_enter_write_spin(&r->r_keypair_lock);
	next = SMR_PTR_GET_LOCKED(&r->r_next);
	current = SMR_PTR_GET_LOCKED(&r->r_current);
	previous = SMR_PTR_GET_LOCKED(&r->r_previous);

	if (kp->kp_is_initiator) {
		if (next != NULL) {
			SMR_PTR_SET_LOCKED(&r->r_next, NULL);
			SMR_PTR_SET_LOCKED(&r->r_previous, next);
			noise_keypair_drop(current);
		} else {
			SMR_PTR_SET_LOCKED(&r->r_previous, current);
		}
		noise_keypair_drop(previous);
		SMR_PTR_SET_LOCKED(&r->r_current, kp);
	} else {
		SMR_PTR_SET_LOCKED(&r->r_next, kp);
		noise_keypair_drop(next);
		SMR_PTR_SET_LOCKED(&r->r_previous, NULL);
		noise_keypair_drop(previous);

	}
	rw_exit_write(&r->r_keypair_lock);

	/* Insert into index table */
	rw_assert_wrlock(&r->r_handshake_lock);

	kp->kp_index.i_is_keypair = 1;
	kp->kp_index.i_local_index = r_i->i_local_index;
	kp->kp_index.i_remote_index = r_i->i_remote_index;

	rw_enter_write_spin(&l->l_index_lock);
	SMR_LIST_INSERT_BEFORE_LOCKED(r_i, &kp->kp_index, i_entry);
	SMR_LIST_REMOVE_LOCKED(r_i, i_entry);
	rw_exit_write(&l->l_index_lock);

	explicit_bzero(&r->r_handshake, sizeof(r->r_handshake));
}

static int
noise_received_with(struct noise_keypair *kp)
{
	struct noise_keypair *old;
	struct noise_remote *r = kp->kp_remote;

	smr_read_enter();
	if (kp != SMR_PTR_GET(&r->r_next)) {
		smr_read_leave();
		return 0;
	}
	smr_read_leave();

	rw_enter_write_spin(&r->r_keypair_lock);
	if (kp != SMR_PTR_GET_LOCKED(&r->r_next)) {
		rw_exit_write(&r->r_keypair_lock);
		return 0;
	}

	old = SMR_PTR_GET_LOCKED(&r->r_previous);
	SMR_PTR_SET_LOCKED(&r->r_previous, SMR_PTR_GET_LOCKED(&r->r_current));
	noise_keypair_drop(old);
	SMR_PTR_SET_LOCKED(&r->r_current, kp);
	SMR_PTR_SET_LOCKED(&r->r_next, NULL);
	rw_exit_write(&r->r_keypair_lock);

	return ECONNRESET;
}

static int
noise_begin_session(struct noise_remote *r)
{
	struct noise_keypair *kp;

	rw_assert_wrlock(&r->r_handshake_lock);

	if ((kp = malloc(sizeof(*kp), M_DEVBUF, M_NOWAIT)) == NULL)
		return ENOSPC;

	refcnt_init(&kp->kp_refcnt);
	kp->kp_can_send = 1;
	kp->kp_is_initiator = r->r_handshake_initiator;
	getnanouptime(&kp->kp_birthdate);
	kp->kp_remote = noise_remote_ref(r);

	if (kp->kp_is_initiator)
		noise_kdf(kp->kp_send, kp->kp_recv, NULL, NULL,
		    NOISE_SYMMETRIC_KEY_LEN, NOISE_SYMMETRIC_KEY_LEN, 0, 0,
		    r->r_handshake.hs_ck);
	else
		noise_kdf(kp->kp_recv, kp->kp_send, NULL, NULL,
		    NOISE_SYMMETRIC_KEY_LEN, NOISE_SYMMETRIC_KEY_LEN, 0, 0,
		    r->r_handshake.hs_ck);

	rw_init(&kp->kp_nonce_lock, "noise_nonce");
	kp->kp_nonce_send = 0;
	kp->kp_nonce_recv = 0;
	bzero(kp->kp_backtrack, sizeof(kp->kp_backtrack));
	smr_init(&kp->kp_smr);

	noise_add_new_keypair(r->r_local, r, kp);
	return 0;
}

struct noise_keypair *
noise_keypair_lookup(struct noise_local *l, uint32_t idx0)
{
	struct noise_index *i;
	struct noise_keypair *kp, *ret = NULL;
	uint32_t idx = idx0 & HT_INDEX_MASK;

	smr_read_enter();
	SMR_LIST_FOREACH(i, &l->l_index_hash[idx], i_entry) {
		if (i->i_local_index == idx0 && i->i_is_keypair) {
			kp = (struct noise_keypair *) i;
			if (refcnt_take_if_gt(&kp->kp_refcnt, 0))
				ret = kp;
			break;
		}
	}
	smr_read_leave();
	return ret;
}

struct noise_keypair *
noise_keypair_current(struct noise_remote *r)
{
	struct noise_keypair *kp, *ret = NULL;

	smr_read_enter();
	kp = SMR_PTR_GET(&r->r_current);
	if (kp != NULL && READ_ONCE(kp->kp_can_send)) {
		if (noise_timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME, 0))
			WRITE_ONCE(kp->kp_can_send, 0);
		else if (refcnt_take_if_gt(&kp->kp_refcnt, 0))
			ret = kp;
	}
	smr_read_leave();
	return ret;
}

struct noise_keypair *
noise_keypair_ref(struct noise_keypair *kp)
{
	refcnt_take(&kp->kp_refcnt);
	return kp;
}

static void
noise_keypair_smr_free(void *_kp)
{
	struct noise_keypair *kp = _kp;
	noise_remote_put(kp->kp_remote);
	explicit_bzero(kp, sizeof(*kp));
	free(kp, M_DEVBUF, sizeof(*kp));
}


void
noise_keypair_put(struct noise_keypair *kp)
{
	if (refcnt_rele(&kp->kp_refcnt))
		smr_call(&kp->kp_smr, noise_keypair_smr_free, kp);
}

static void
noise_keypair_drop(struct noise_keypair *kp)
{
	struct noise_remote *r;
	struct noise_local *l;

	if (kp == NULL)
		return;

	r = kp->kp_remote;
	l = r->r_local;

	rw_enter_write_spin(&l->l_index_lock);
	SMR_LIST_REMOVE_LOCKED(&kp->kp_index, i_entry);
	rw_exit_write(&l->l_index_lock);

	noise_keypair_put(kp);
}

struct noise_remote *
noise_keypair_remote(struct noise_keypair *kp)
{
	return noise_remote_ref(kp->kp_remote);
}

int
noise_keypair_nonce_next(struct noise_keypair *kp, uint64_t *send)
{
#ifdef __LP64__
	*send = atomic_inc_long_nv((u_long *)&kp->kp_nonce_send) - 1;
#else
	rw_enter_write_spin(&kp->kp_nonce_lock);
	*send = ctr->c_send++;
	rw_exit_write(&kp->kp_nonce_lock);
#endif
	if (*send < REJECT_AFTER_MESSAGES)
		return 0;
	WRITE_ONCE(kp->kp_can_send, 0);
	return EINVAL;
}

int
noise_keypair_nonce_check(struct noise_keypair *kp, uint64_t recv)
{
	uint64_t i, top, index_recv, index_ctr;
	unsigned long bit;
	int ret = EEXIST;

	rw_enter_write_spin(&kp->kp_nonce_lock);

	/* Check that the recv counter is valid */
	if (kp->kp_nonce_recv >= REJECT_AFTER_MESSAGES ||
	    recv >= REJECT_AFTER_MESSAGES)
		goto error;

	/* If the packet is out of the window, invalid */
	if (recv + COUNTER_WINDOW_SIZE < kp->kp_nonce_recv)
		goto error;

	/* If the new counter is ahead of the current counter, we'll need to
	 * zero out the bitmap that has previously been used */
	index_recv = recv / COUNTER_BITS;
	index_ctr = kp->kp_nonce_recv / COUNTER_BITS;

	if (recv > kp->kp_nonce_recv) {
		top = MIN(index_recv - index_ctr, COUNTER_NUM);
		for (i = 1; i <= top; i++)
			kp->kp_backtrack[
			    (i + index_ctr) & (COUNTER_NUM - 1)] = 0;
		WRITE_ONCE(kp->kp_nonce_recv, recv);
	}

	index_recv %= COUNTER_NUM;
	bit = 1ul << (recv % COUNTER_BITS);

	if (kp->kp_backtrack[index_recv] & bit)
		goto error;

	kp->kp_backtrack[index_recv] |= bit;

	ret = 0;
error:
	rw_exit_write(&kp->kp_nonce_lock);
	return ret;
}

int
noise_keep_key_fresh_send(struct noise_remote *r)
{
	struct noise_keypair *current;
	int keep_key_fresh;

	smr_read_enter();
	current = SMR_PTR_GET(&r->r_current);
	keep_key_fresh = current != NULL && READ_ONCE(current->kp_can_send) && (
	    READ_ONCE(current->kp_nonce_send) > REKEY_AFTER_MESSAGES ||
	    (current->kp_is_initiator && noise_timer_expired(&current->kp_birthdate, REKEY_AFTER_TIME, 0)));
	smr_read_leave();

	return keep_key_fresh ? ESTALE : 0;
}

int
noise_keep_key_fresh_recv(struct noise_remote *r)
{
	struct noise_keypair *current;
	int keep_key_fresh;

	smr_read_enter();
	current = SMR_PTR_GET(&r->r_current);
	keep_key_fresh = current != NULL && READ_ONCE(current->kp_can_send) &&
	    current->kp_is_initiator && noise_timer_expired(&current->kp_birthdate,
			REJECT_AFTER_TIME - KEEPALIVE_TIMEOUT - REKEY_TIMEOUT, 0);
	smr_read_leave();

	return keep_key_fresh ? ESTALE : 0;
}

void
noise_keypair_encrypt(struct noise_keypair *kp, uint32_t *r_idx, uint64_t nonce,
    uint8_t *buf, size_t buflen)
{
	chacha20poly1305_encrypt(buf, buf, buflen, NULL, 0, nonce, kp->kp_send);
	*r_idx = kp->kp_index.i_remote_index;
}

int
noise_keypair_decrypt(struct noise_keypair *kp, uint64_t nonce, uint8_t *buf,
    size_t buflen)
{
	if (READ_ONCE(kp->kp_nonce_recv) >= REJECT_AFTER_MESSAGES ||
	    noise_timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME, 0))
		return EINVAL;

	if (chacha20poly1305_decrypt(buf, buf, buflen, NULL, 0, nonce, kp->kp_recv) == 0)
		return EINVAL;

	if (noise_received_with(kp) != 0)
		return ECONNRESET;

	return 0;
}


/* Handshake functions */
int
noise_create_initiation(struct noise_remote *r,
    uint32_t *s_idx,
    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
    uint8_t es[NOISE_PUBLIC_KEY_LEN + NOISE_AUTHTAG_LEN],
    uint8_t ets[NOISE_TIMESTAMP_LEN + NOISE_AUTHTAG_LEN])
{
	struct noise_handshake *hs = &r->r_handshake;
	struct noise_local *l = r->r_local;
	uint8_t key[NOISE_SYMMETRIC_KEY_LEN];
	int ret = EINVAL;

	rw_enter_read_spin(&l->l_identity_lock);
	rw_enter_write_spin(&r->r_handshake_lock);
	if (!l->l_has_identity)
		goto error;
	if (!noise_timer_expired(&r->r_last_sent, REKEY_TIMEOUT, 0))
		goto error;
	noise_param_init(hs->hs_ck, hs->hs_hash, r->r_public);

	/* e */
	curve25519_generate_secret(hs->hs_e);
	if (curve25519_generate_public(ue, hs->hs_e) == 0)
		goto error;
	noise_msg_ephemeral(hs->hs_ck, hs->hs_hash, ue);

	/* es */
	if (noise_mix_dh(hs->hs_ck, key, hs->hs_e, r->r_public) != 0)
		goto error;

	/* s */
	noise_msg_encrypt(es, l->l_public,
	    NOISE_PUBLIC_KEY_LEN, key, hs->hs_hash);

	/* ss */
	if (noise_mix_ss(hs->hs_ck, key, r->r_ss) != 0)
		goto error;

	/* {t} */
	noise_tai64n_now(ets);
	noise_msg_encrypt(ets, ets,
	    NOISE_TIMESTAMP_LEN, key, hs->hs_hash);

	noise_remote_index_insert(l, r);
	getnanouptime(&r->r_last_sent);
	*s_idx = r->r_index.i_local_index;
	r->r_handshake_initiator = 1;
	ret = 0;
error:
	rw_exit_write(&r->r_handshake_lock);
	rw_exit_read(&l->l_identity_lock);
	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
	return ret;
}

int
noise_consume_initiation(struct noise_local *l, struct noise_remote **rp,
    uint32_t s_idx,
    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
    uint8_t es[NOISE_PUBLIC_KEY_LEN + NOISE_AUTHTAG_LEN],
    uint8_t ets[NOISE_TIMESTAMP_LEN + NOISE_AUTHTAG_LEN])
{
	struct noise_remote *r;
	struct noise_handshake hs;
	uint8_t key[NOISE_SYMMETRIC_KEY_LEN];
	uint8_t r_public[NOISE_PUBLIC_KEY_LEN];
	uint8_t	timestamp[NOISE_TIMESTAMP_LEN];
	int ret = EINVAL;

	rw_enter_read_spin(&l->l_identity_lock);
	if (!l->l_has_identity)
		goto error;
	noise_param_init(hs.hs_ck, hs.hs_hash, l->l_public);

	/* e */
	noise_msg_ephemeral(hs.hs_ck, hs.hs_hash, ue);

	/* es */
	if (noise_mix_dh(hs.hs_ck, key, l->l_private, ue) != 0)
		goto error;

	/* s */
	if (noise_msg_decrypt(r_public, es,
	    NOISE_PUBLIC_KEY_LEN + NOISE_AUTHTAG_LEN, key, hs.hs_hash) != 0)
		goto error;

	/* Lookup the remote we received from */
	if ((r = noise_remote_lookup(l, r_public)) == NULL)
		goto error;

	/* ss */
	if (noise_mix_ss(hs.hs_ck, key, r->r_ss) != 0)
		goto error_put;

	/* {t} */
	if (noise_msg_decrypt(timestamp, ets,
	    NOISE_TIMESTAMP_LEN + NOISE_AUTHTAG_LEN, key, hs.hs_hash) != 0)
		goto error_put;

	memcpy(hs.hs_e, ue, NOISE_PUBLIC_KEY_LEN);

	/* We have successfully computed the same results, now we ensure that
	 * this is not an initiation replay, or a flood attack */
	rw_enter_write_spin(&r->r_handshake_lock);

	/* Replay */
	if (memcmp(timestamp, r->r_timestamp, NOISE_TIMESTAMP_LEN) > 0)
		memcpy(r->r_timestamp, timestamp, NOISE_TIMESTAMP_LEN);
	else
		goto error_set;
	/* Flood attack */
	if (noise_timer_expired(&r->r_last_init_recv, 0, REJECT_INTERVAL))
		getnanouptime(&r->r_last_init_recv);
	else
		goto error_set;

	/* Ok, we're happy to accept this initiation now */
	noise_remote_index_insert(l, r);
	r->r_index.i_remote_index = s_idx;
	r->r_handshake_initiator = 0;
	r->r_handshake = hs;
	*rp = noise_remote_ref(r);
	ret = 0;
error_set:
	rw_exit_write(&r->r_handshake_lock);
error_put:
	noise_remote_put(r);
error:
	rw_exit_read(&l->l_identity_lock);
	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
	explicit_bzero(&hs, sizeof(hs));
	return ret;
}

int
noise_create_response(struct noise_remote *r,
    uint32_t *s_idx, uint32_t *r_idx,
    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
    uint8_t en[0 + NOISE_AUTHTAG_LEN])
{
	struct noise_handshake *hs = &r->r_handshake;
	struct noise_local *l = r->r_local;
	uint8_t key[NOISE_SYMMETRIC_KEY_LEN];
	uint8_t e[NOISE_PUBLIC_KEY_LEN];
	int ret = EINVAL;

	rw_enter_read_spin(&l->l_identity_lock);
	rw_enter_write_spin(&r->r_handshake_lock);

	if (!r->r_handshake_alive || r->r_handshake_initiator)
		goto error;

	/* e */
	curve25519_generate_secret(e);
	if (curve25519_generate_public(ue, e) == 0)
		goto error;
	noise_msg_ephemeral(hs->hs_ck, hs->hs_hash, ue);

	/* ee */
	if (noise_mix_dh(hs->hs_ck, NULL, e, hs->hs_e) != 0)
		goto error;

	/* se */
	if (noise_mix_dh(hs->hs_ck, NULL, e, r->r_public) != 0)
		goto error;

	/* psk */
	noise_mix_psk(hs->hs_ck, hs->hs_hash, key, r->r_psk);

	/* {} */
	noise_msg_encrypt(en, NULL, 0, key, hs->hs_hash);

	if ((ret = noise_begin_session(r)) == 0) {
		getnanouptime(&r->r_last_sent);
		*s_idx = r->r_index.i_local_index;
		*r_idx = r->r_index.i_remote_index;
	}
error:
	rw_exit_write(&r->r_handshake_lock);
	rw_exit_read(&l->l_identity_lock);
	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
	explicit_bzero(e, NOISE_PUBLIC_KEY_LEN);
	return ret;
}

int
noise_consume_response(struct noise_local *l, struct noise_remote **rp,
    uint32_t s_idx, uint32_t r_idx,
    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
    uint8_t en[0 + NOISE_AUTHTAG_LEN])
{
	uint8_t preshared_key[NOISE_SYMMETRIC_KEY_LEN];
	uint8_t key[NOISE_SYMMETRIC_KEY_LEN];
	struct noise_handshake hs;
	struct noise_remote *r = NULL;
	int ret = EINVAL;

	if ((r = noise_remote_index_lookup(l, r_idx)) == NULL)
		return ret;

	rw_enter_read_spin(&l->l_identity_lock);
	if (!l->l_has_identity)
		goto error;

	rw_enter_read_spin(&r->r_handshake_lock);
	if (!r->r_handshake_alive || !r->r_handshake_initiator) {
		rw_exit_read(&r->r_handshake_lock);
		goto error;
	}
	memcpy(preshared_key, r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
	hs = r->r_handshake;
	rw_exit_read(&r->r_handshake_lock);

	/* e */
	noise_msg_ephemeral(hs.hs_ck, hs.hs_hash, ue);

	/* ee */
	if (noise_mix_dh(hs.hs_ck, NULL, hs.hs_e, ue) != 0)
		goto error_zero;

	/* se */
	if (noise_mix_dh(hs.hs_ck, NULL, l->l_private, ue) != 0)
		goto error_zero;

	/* psk */
	noise_mix_psk(hs.hs_ck, hs.hs_hash, key, preshared_key);

	/* {} */
	if (noise_msg_decrypt(NULL, en,
	    0 + NOISE_AUTHTAG_LEN, key, hs.hs_hash) != 0)
		goto error_zero;

	rw_enter_write_spin(&r->r_handshake_lock);
	if (r->r_handshake_alive && r->r_handshake_initiator &&
	    r->r_index.i_local_index == r_idx) {
		r->r_handshake = hs;
		r->r_index.i_remote_index = s_idx;
		ret = noise_begin_session(r);
		*rp = noise_remote_ref(r);
	}
	rw_exit_write(&r->r_handshake_lock);
error_zero:
	explicit_bzero(preshared_key, NOISE_SYMMETRIC_KEY_LEN);
	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
	explicit_bzero(&hs, sizeof(hs));
error:
	rw_exit_read(&l->l_identity_lock);
	noise_remote_put(r);
	return ret;
}

/* Handshake helper functions */
static void
noise_kdf(uint8_t *a, uint8_t *b, uint8_t *c, const uint8_t *x,
    size_t a_len, size_t b_len, size_t c_len, size_t x_len,
    const uint8_t ck[NOISE_HASH_LEN])
{
	uint8_t out[BLAKE2S_HASH_SIZE + 1];
	uint8_t sec[BLAKE2S_HASH_SIZE];

	/* Extract entropy from "x" into sec */
	blake2s_hmac(sec, x, ck, BLAKE2S_HASH_SIZE, x_len, NOISE_HASH_LEN);

	if (a == NULL || a_len == 0)
		goto out;

	/* Expand first key: key = sec, data = 0x1 */
	out[0] = 1;
	blake2s_hmac(out, out, sec, BLAKE2S_HASH_SIZE, 1, BLAKE2S_HASH_SIZE);
	memcpy(a, out, a_len);

	if (b == NULL || b_len == 0)
		goto out;

	/* Expand second key: key = sec, data = "a" || 0x2 */
	out[BLAKE2S_HASH_SIZE] = 2;
	blake2s_hmac(out, out, sec, BLAKE2S_HASH_SIZE, BLAKE2S_HASH_SIZE + 1,
			BLAKE2S_HASH_SIZE);
	memcpy(b, out, b_len);

	if (c == NULL || c_len == 0)
		goto out;

	/* Expand third key: key = sec, data = "b" || 0x3 */
	out[BLAKE2S_HASH_SIZE] = 3;
	blake2s_hmac(out, out, sec, BLAKE2S_HASH_SIZE, BLAKE2S_HASH_SIZE + 1,
			BLAKE2S_HASH_SIZE);
	memcpy(c, out, c_len);

out:
	/* Clear sensitive data from stack */
	explicit_bzero(sec, BLAKE2S_HASH_SIZE);
	explicit_bzero(out, BLAKE2S_HASH_SIZE + 1);
}

static int
noise_mix_dh(uint8_t ck[NOISE_HASH_LEN], uint8_t key[NOISE_SYMMETRIC_KEY_LEN],
    const uint8_t private[NOISE_PUBLIC_KEY_LEN],
    const uint8_t public[NOISE_PUBLIC_KEY_LEN])
{
	uint8_t dh[NOISE_PUBLIC_KEY_LEN];

	if (!curve25519(dh, private, public))
		return EINVAL;
	noise_kdf(ck, key, NULL, dh,
	    NOISE_HASH_LEN, NOISE_SYMMETRIC_KEY_LEN, 0, NOISE_PUBLIC_KEY_LEN, ck);
	explicit_bzero(dh, NOISE_PUBLIC_KEY_LEN);
	return 0;
}

static int
noise_mix_ss(uint8_t ck[NOISE_HASH_LEN], uint8_t key[NOISE_SYMMETRIC_KEY_LEN],
    const uint8_t ss[NOISE_PUBLIC_KEY_LEN])
{
	static uint8_t null_point[NOISE_PUBLIC_KEY_LEN];
	if (timingsafe_bcmp(ss, null_point, NOISE_PUBLIC_KEY_LEN) == 0)
		return ENOENT;
	noise_kdf(ck, key, NULL, ss,
	    NOISE_HASH_LEN, NOISE_SYMMETRIC_KEY_LEN, 0, NOISE_PUBLIC_KEY_LEN, ck);
	return 0;
}

static void
noise_mix_hash(uint8_t hash[NOISE_HASH_LEN], const uint8_t *src,
    size_t src_len)
{
	struct blake2s_state blake;

	blake2s_init(&blake, NOISE_HASH_LEN);
	blake2s_update(&blake, hash, NOISE_HASH_LEN);
	blake2s_update(&blake, src, src_len);
	blake2s_final(&blake, hash);
}

static void
noise_mix_psk(uint8_t ck[NOISE_HASH_LEN], uint8_t hash[NOISE_HASH_LEN],
    uint8_t key[NOISE_SYMMETRIC_KEY_LEN],
    const uint8_t psk[NOISE_SYMMETRIC_KEY_LEN])
{
	uint8_t tmp[NOISE_HASH_LEN];

	noise_kdf(ck, tmp, key, psk,
	    NOISE_HASH_LEN, NOISE_HASH_LEN, NOISE_SYMMETRIC_KEY_LEN,
	    NOISE_SYMMETRIC_KEY_LEN, ck);
	noise_mix_hash(hash, tmp, NOISE_HASH_LEN);
	explicit_bzero(tmp, NOISE_HASH_LEN);
}

static void
noise_param_init(uint8_t ck[NOISE_HASH_LEN], uint8_t hash[NOISE_HASH_LEN],
    const uint8_t s[NOISE_PUBLIC_KEY_LEN])
{
	struct blake2s_state blake;

	blake2s(ck, (uint8_t *)NOISE_HANDSHAKE_NAME, NULL,
	    NOISE_HASH_LEN, strlen(NOISE_HANDSHAKE_NAME), 0);
	blake2s_init(&blake, NOISE_HASH_LEN);
	blake2s_update(&blake, ck, NOISE_HASH_LEN);
	blake2s_update(&blake, (uint8_t *)NOISE_IDENTIFIER_NAME,
	    strlen(NOISE_IDENTIFIER_NAME));
	blake2s_final(&blake, hash);

	noise_mix_hash(hash, s, NOISE_PUBLIC_KEY_LEN);
}

static void
noise_msg_encrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
    uint8_t key[NOISE_SYMMETRIC_KEY_LEN], uint8_t hash[NOISE_HASH_LEN])
{
	/* Nonce always zero for Noise_IK */
	chacha20poly1305_encrypt(dst, src, src_len,
	    hash, NOISE_HASH_LEN, 0, key);
	noise_mix_hash(hash, dst, src_len + NOISE_AUTHTAG_LEN);
}

static int
noise_msg_decrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
    uint8_t key[NOISE_SYMMETRIC_KEY_LEN], uint8_t hash[NOISE_HASH_LEN])
{
	/* Nonce always zero for Noise_IK */
	if (!chacha20poly1305_decrypt(dst, src, src_len,
	    hash, NOISE_HASH_LEN, 0, key))
		return EINVAL;
	noise_mix_hash(hash, src, src_len);
	return 0;
}

static void
noise_msg_ephemeral(uint8_t ck[NOISE_HASH_LEN], uint8_t hash[NOISE_HASH_LEN],
    const uint8_t src[NOISE_PUBLIC_KEY_LEN])
{
	noise_mix_hash(hash, src, NOISE_PUBLIC_KEY_LEN);
	noise_kdf(ck, NULL, NULL, src, NOISE_HASH_LEN, 0, 0,
		  NOISE_PUBLIC_KEY_LEN, ck);
}

static void
noise_tai64n_now(uint8_t output[NOISE_TIMESTAMP_LEN])
{
	struct timespec time;
	uint64_t sec;
	uint32_t nsec;

	getnanotime(&time);

	/* Round down the nsec counter to limit precise timing leak. */
	time.tv_nsec &= REJECT_INTERVAL_MASK;

	/* https://cr.yp.to/libtai/tai64.html */
	sec = htobe64(0x400000000000000aULL + time.tv_sec);
	nsec = htobe32(time.tv_nsec);

	/* memcpy to output buffer, assuming output could be unaligned. */
	memcpy(output, &sec, sizeof(sec));
	memcpy(output + sizeof(sec), &nsec, sizeof(nsec));
}

static int
noise_timer_expired(struct timespec *birthdate, time_t sec, long nsec)
{
	struct timespec uptime;
	struct timespec expire = { .tv_sec = sec, .tv_nsec = nsec };

	getnanouptime(&uptime);
	timespecadd(birthdate, &expire, &expire);
	return timespeccmp(&uptime, &expire, >) ? ETIMEDOUT : 0;
}

#ifdef WGTEST

#define T_LIM (COUNTER_WINDOW_SIZE + 1)
#define T_INIT do {				\
	bzero(&kp, sizeof(kp));			\
	rw_init(&kp.kp_nonce_lock, "counter");	\
} while (0)
#define T(num, v, e) do {						\
	if (noise_keypair_nonce_check(&kp, v) != e) {			\
		printf("%s, test %d: failed.\n", __func__, num);	\
		return;							\
	}								\
} while (0)
#define T_PASSED printf("%s: passed.\n", __func__)

void
noise_test()
{
	struct noise_keypair kp;
	int i;

	T_INIT;
	/* T(test number, nonce, expected_response) */
	T( 1, 0, 0);
	T( 2, 1, 0);
	T( 3, 1, EEXIST);
	T( 4, 9, 0);
	T( 5, 8, 0);
	T( 6, 7, 0);
	T( 7, 7, EEXIST);
	T( 8, T_LIM, 0);
	T( 9, T_LIM - 1, 0);
	T(10, T_LIM - 1, EEXIST);
	T(11, T_LIM - 2, 0);
	T(12, 2, 0);
	T(13, 2, EEXIST);
	T(14, T_LIM + 16, 0);
	T(15, 3, EEXIST);
	T(16, T_LIM + 16, EEXIST);
	T(17, T_LIM * 4, 0);
	T(18, T_LIM * 4 - (T_LIM - 1), 0);
	T(19, 10, EEXIST);
	T(20, T_LIM * 4 - T_LIM, EEXIST);
	T(21, T_LIM * 4 - (T_LIM + 1), EEXIST);
	T(22, T_LIM * 4 - (T_LIM - 2), 0);
	T(23, T_LIM * 4 + 1 - T_LIM, EEXIST);
	T(24, 0, EEXIST);
	T(25, REJECT_AFTER_MESSAGES, EEXIST);
	T(26, REJECT_AFTER_MESSAGES - 1, 0);
	T(27, REJECT_AFTER_MESSAGES, EEXIST);
	T(28, REJECT_AFTER_MESSAGES - 1, EEXIST);
	T(29, REJECT_AFTER_MESSAGES - 2, 0);
	T(30, REJECT_AFTER_MESSAGES + 1, EEXIST);
	T(31, REJECT_AFTER_MESSAGES + 2, EEXIST);
	T(32, REJECT_AFTER_MESSAGES - 2, EEXIST);
	T(33, REJECT_AFTER_MESSAGES - 3, 0);
	T(34, 0, EEXIST);

	T_INIT;
	for (i = 1; i <= COUNTER_WINDOW_SIZE; ++i)
		T(35, i, 0);
	T(36, 0, 0);
	T(37, 0, EEXIST);

	T_INIT;
	for (i = 2; i <= COUNTER_WINDOW_SIZE + 1; ++i)
		T(38, i, 0);
	T(39, 1, 0);
	T(40, 0, EEXIST);

	T_INIT;
	for (i = COUNTER_WINDOW_SIZE + 1; i-- > 0;)
		T(41, i, 0);

	T_INIT;
	for (i = COUNTER_WINDOW_SIZE + 2; i-- > 1;)
		T(42, i, 0);
	T(43, 0, EEXIST);

	T_INIT;
	for (i = COUNTER_WINDOW_SIZE + 1; i-- > 1;)
		T(44, i, 0);
	T(45, COUNTER_WINDOW_SIZE + 1, 0);
	T(46, 0, EEXIST);

	T_INIT;
	for (i = COUNTER_WINDOW_SIZE + 1; i-- > 1;)
		T(47, i, 0);
	T(48, 0, 0);
	T(49, COUNTER_WINDOW_SIZE + 1, 0);

	T_PASSED;
}
#endif /* WGTEST */