xref: /dragonfly/sys/net/wg/wg_noise.c (revision 979e91ed)

/*-
 * SPDX-License-Identifier: ISC
 *
 * Copyright (C) 2015-2021 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 * Copyright (C) 2019-2021 Matt Dunwoodie <ncon@noconroy.net>
 * Copyright (c) 2022 The FreeBSD Foundation
 * Copyright (c) 2023-2024 Aaron LI <aly@aaronly.me>
 *
 * 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.
 */
/*
 * This implements Noise_IKpsk2:
 * <- s
 * ******
 * -> e, es, s, ss, {t}
 * <- e, ee, se, psk, {}
 */

#include <sys/param.h>
#include <sys/systm.h>
#include <sys/bitops.h> /* ilog2() */
#include <sys/endian.h>
#include <sys/kernel.h>
#include <sys/lock.h>
#include <sys/malloc.h>
#include <sys/queue.h>
#include <sys/refcount.h>
#include <sys/time.h>

#include <machine/atomic.h>

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

#include "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_ORDER		ilog2(COUNTER_BITS)
#define COUNTER_WINDOW_SIZE	(COUNTER_BITS_TOTAL - COUNTER_BITS)
#define COUNTER_NUM		(COUNTER_BITS_TOTAL / COUNTER_BITS)
#define COUNTER_MASK		(COUNTER_NUM - 1)

/* 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 */
#define REJECT_INTERVAL_MASK	(~((1ULL << ilog2(REJECT_INTERVAL)) - 1))

/* Constants for the hashtable */
#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 {
	LIST_ENTRY(noise_index)	 i_entry;
	uint32_t		 i_local_index;
	uint32_t		 i_remote_index;
	bool			 i_is_keypair;
};

struct noise_keypair {
	struct noise_index	 kp_index;

	u_int			 kp_refcnt;
	bool			 kp_can_send;
	bool			 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];

	struct lock		 kp_counter_lock;
	uint64_t		 kp_counter_send; /* next counter available */
	uint64_t		 kp_counter_recv; /* max counter received */
	unsigned long		 kp_backtrack[COUNTER_NUM];

#ifdef INVARIANTS
	LIST_ENTRY(noise_keypair) _kp_entry;
#endif
};

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

/* Handshake states of the remote/peer side. */
enum noise_handshake_state {
	HANDSHAKE_DEAD,
	HANDSHAKE_INITIATOR,
	HANDSHAKE_RESPONDER,
};

struct noise_remote {
	struct noise_index		 r_index;

	LIST_ENTRY(noise_remote)	 r_entry;
	bool				 r_entry_inserted;
	uint8_t				 r_public[NOISE_PUBLIC_KEY_LEN];

	struct lock			 r_handshake_lock;
	struct noise_handshake		 r_handshake;
	enum noise_handshake_state	 r_handshake_state;
	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];

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

	struct lock			 r_keypair_lock;
	struct noise_keypair		*r_keypair_next;
	struct noise_keypair		*r_keypair_current;
	struct noise_keypair		*r_keypair_previous;

#ifdef INVARIANTS
	LIST_ENTRY(noise_remote)	 _r_entry;
#endif
};

struct noise_local {
	struct lock			 l_identity_lock;
	bool				 l_has_identity;
	uint8_t				 l_public[NOISE_PUBLIC_KEY_LEN];
	uint8_t				 l_private[NOISE_PUBLIC_KEY_LEN];

	u_int				 l_refcnt;
	uint8_t				 l_hash_key[SIPHASH_KEY_LENGTH];

	/* Hash table to lookup the remote from its public key. */
	struct lock			 l_remote_lock;
	size_t				 l_remote_num;
	LIST_HEAD(, noise_remote)	 l_remote_hash[HT_REMOTE_SIZE];

	/* Hash table to lookup the remote/keypair from its index. */
	struct lock			 l_index_lock;
	LIST_HEAD(, noise_index)	 l_index_hash[HT_INDEX_SIZE];

#ifdef INVARIANTS
	LIST_ENTRY(noise_local)		 _l_entry;
#endif
};


static MALLOC_DEFINE(M_NOISE, "NOISE", "wgnoise");

#ifdef INVARIANTS
static struct lock noise_mtx;
static LIST_HEAD(, noise_local) noise_locals;
static LIST_HEAD(, noise_remote) noise_remotes;
static LIST_HEAD(, noise_keypair) noise_keypairs;
#endif


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

static struct noise_local *
		noise_local_ref(struct noise_local *);
static void	noise_local_put(struct noise_local *);

static uint32_t	noise_remote_index_insert(struct noise_local *,
					  struct noise_remote *);
static struct noise_remote *
		noise_remote_index_lookup(struct noise_local *,
					  uint32_t, bool);
static bool	noise_remote_index_remove(struct noise_local *,
					  struct noise_remote *);
static void	noise_remote_expire_current(struct noise_remote *);

static bool	noise_begin_session(struct noise_remote *);
static void	noise_keypair_drop(struct noise_keypair *);

static void	noise_hmac(uint8_t *, const uint8_t *, const uint8_t *,
			   size_t, size_t, size_t);
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 bool	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 bool	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 bool	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 inline uint64_t
siphash24(const uint8_t key[SIPHASH_KEY_LENGTH], const void *src, size_t len)
{
	SIPHASH_CTX ctx;
	return SipHashX(&ctx, 2, 4, key, src, len);
}

static inline bool
timer_expired(const struct timespec *birthdate, time_t sec, long nsec)
{
	struct timespec uptime;
	struct timespec expire = { .tv_sec = sec, .tv_nsec = nsec };

	if (__predict_false(!timespecisset(birthdate)))
		return (true);

	getnanouptime(&uptime);
	timespecadd(birthdate, &expire, &expire);
	return timespeccmp(&uptime, &expire, >);
}

/*----------------------------------------------------------------------------*/

int
noise_init(void)
{
#ifdef INVARIANTS
	lockinit(&noise_mtx, "noise mtx lock", 0, 0);

	LIST_INIT(&noise_locals);
	LIST_INIT(&noise_remotes);
	LIST_INIT(&noise_keypairs);
#endif

	return (0);
}

void
noise_deinit(void)
{
#ifdef INVARIANTS
	lockmgr(&noise_mtx, LK_EXCLUSIVE);

	if (!LIST_EMPTY(&noise_locals))
		panic("%s: noise_local leaked", __func__);
	if (!LIST_EMPTY(&noise_remotes))
		panic("%s: noise_remote leaked", __func__);
	if (!LIST_EMPTY(&noise_keypairs))
		panic("%s: noise_keypair leaked", __func__);

	lockmgr(&noise_mtx, LK_RELEASE);
	lockuninit(&noise_mtx);
#endif
}

/*----------------------------------------------------------------------------*/
/* Local configuration */

struct noise_local *
noise_local_alloc(void)
{
	struct noise_local *l;
	size_t i;

	l = kmalloc(sizeof(*l), M_NOISE, M_WAITOK | M_ZERO);

	lockinit(&l->l_identity_lock, "noise_identity", 0, 0);
	refcount_init(&l->l_refcnt, 1);
	karc4random_buf(l->l_hash_key, sizeof(l->l_hash_key));

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

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

#ifdef INVARIANTS
	lockmgr(&noise_mtx, LK_EXCLUSIVE);
	LIST_INSERT_HEAD(&noise_locals, l, _l_entry);
	lockmgr(&noise_mtx, LK_RELEASE);
#endif

	return (l);
}

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

static void
noise_local_put(struct noise_local *l)
{
	if (refcount_release(&l->l_refcnt)) {
#ifdef INVARIANTS
		lockmgr(&noise_mtx, LK_EXCLUSIVE);
		LIST_REMOVE(l, _l_entry);
		lockmgr(&noise_mtx, LK_RELEASE);
#endif

		lockuninit(&l->l_identity_lock);
		lockuninit(&l->l_remote_lock);
		lockuninit(&l->l_index_lock);
		explicit_bzero(l, sizeof(*l));
		kfree(l, M_NOISE);
	}
}

void
noise_local_free(struct noise_local *l)
{
	noise_local_put(l);
}

bool
noise_local_set_private(struct noise_local *l,
			const uint8_t private[NOISE_PUBLIC_KEY_LEN])
{
	struct noise_remote *r;
	size_t i;
	bool has_identity;

	lockmgr(&l->l_identity_lock, LK_EXCLUSIVE);

	/* Note: we might be removing the private key. */
	memcpy(l->l_private, private, NOISE_PUBLIC_KEY_LEN);
	curve25519_clamp_secret(l->l_private);
	has_identity = l->l_has_identity =
		curve25519_generate_public(l->l_public, l->l_private);

	/* Invalidate all existing handshakes. */
	lockmgr(&l->l_remote_lock, LK_SHARED);
	for (i = 0; i < HT_REMOTE_SIZE; i++) {
		LIST_FOREACH(r, &l->l_remote_hash[i], r_entry) {
			noise_precompute_ss(l, r);
			noise_remote_expire_current(r);
		}
	}
	lockmgr(&l->l_remote_lock, LK_RELEASE);

	lockmgr(&l->l_identity_lock, LK_RELEASE);
	return (has_identity);
}

bool
noise_local_keys(struct noise_local *l, uint8_t public[NOISE_PUBLIC_KEY_LEN],
		 uint8_t private[NOISE_PUBLIC_KEY_LEN])
{
	bool has_identity;

	lockmgr(&l->l_identity_lock, LK_SHARED);
	has_identity = l->l_has_identity;
	if (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);
	}
	lockmgr(&l->l_identity_lock, LK_RELEASE);

	return (has_identity);
}

static void
noise_precompute_ss(struct noise_local *l, struct noise_remote *r)
{
	KKASSERT(lockstatus(&l->l_identity_lock, curthread) != 0);

	lockmgr(&r->r_handshake_lock, LK_EXCLUSIVE);
	if (!l->l_has_identity ||
	    !curve25519(r->r_ss, l->l_private, r->r_public))
		bzero(r->r_ss, NOISE_PUBLIC_KEY_LEN);
	lockmgr(&r->r_handshake_lock, LK_RELEASE);
}

/*----------------------------------------------------------------------------*/
/* Remote configuration */

struct noise_remote *
noise_remote_alloc(struct noise_local *l,
		   const uint8_t public[NOISE_PUBLIC_KEY_LEN], void *arg)
{
	struct noise_remote *r;

	r = kmalloc(sizeof(*r), M_NOISE, M_WAITOK | M_ZERO);

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

	lockinit(&r->r_handshake_lock, "noise_handshake", 0, 0);
	r->r_handshake_state = HANDSHAKE_DEAD;

	refcount_init(&r->r_refcnt, 1);
	r->r_local = noise_local_ref(l);
	r->r_arg = arg;

	lockinit(&r->r_keypair_lock, "noise_keypair", 0, 0);

	lockmgr(&l->l_identity_lock, LK_SHARED);
	noise_precompute_ss(l, r);
	lockmgr(&l->l_identity_lock, LK_RELEASE);

#ifdef INVARIANTS
	lockmgr(&noise_mtx, LK_EXCLUSIVE);
	LIST_INSERT_HEAD(&noise_remotes, r, _r_entry);
	lockmgr(&noise_mtx, LK_RELEASE);
#endif

	return (r);
}

int
noise_remote_enable(struct noise_remote *r)
{
	struct noise_local *l = r->r_local;
	uint64_t idx;
	int ret = 0;

	idx = siphash24(l->l_hash_key, r->r_public, NOISE_PUBLIC_KEY_LEN);
	idx &= HT_REMOTE_MASK;

	lockmgr(&l->l_remote_lock, LK_EXCLUSIVE);
	if (!r->r_entry_inserted) {
		/* Insert to hashtable */
		if (l->l_remote_num < MAX_REMOTE_PER_LOCAL) {
			r->r_entry_inserted = true;
			l->l_remote_num++;
			LIST_INSERT_HEAD(&l->l_remote_hash[idx], r, r_entry);
		} else {
			ret = ENOSPC;
		}
	}
	lockmgr(&l->l_remote_lock, LK_RELEASE);

	return (ret);
}

void
noise_remote_disable(struct noise_remote *r)
{
	struct noise_local *l = r->r_local;

	/* Remove from hashtable */
	lockmgr(&l->l_remote_lock, LK_EXCLUSIVE);
	if (r->r_entry_inserted) {
		r->r_entry_inserted = false;
		LIST_REMOVE(r, r_entry);
		l->l_remote_num--;
	};
	lockmgr(&l->l_remote_lock, LK_RELEASE);
}

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);
	idx &= HT_REMOTE_MASK;

	lockmgr(&l->l_remote_lock, LK_SHARED);
	LIST_FOREACH(r, &l->l_remote_hash[idx], r_entry) {
		if (timingsafe_bcmp(r->r_public, public, NOISE_PUBLIC_KEY_LEN)
		    == 0) {
			ret = noise_remote_ref(r);
			break;
		}
	}
	lockmgr(&l->l_remote_lock, LK_RELEASE);

	return (ret);
}

static uint32_t
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);

	lockmgr(&l->l_index_lock, LK_EXCLUSIVE);
retry:
	r_i->i_local_index = karc4random();
	idx = r_i->i_local_index & HT_INDEX_MASK;
	LIST_FOREACH(i, &l->l_index_hash[idx], i_entry) {
		if (i->i_local_index == r_i->i_local_index)
			goto retry;
	}
	LIST_INSERT_HEAD(&l->l_index_hash[idx], r_i, i_entry);
	lockmgr(&l->l_index_lock, LK_RELEASE);

	return (r_i->i_local_index);
}

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

	lockmgr(&l->l_index_lock, LK_SHARED);
	LIST_FOREACH(i, &l->l_index_hash[idx], i_entry) {
		if (i->i_local_index == idx0) {
			if (!i->i_is_keypair) {
				r = (struct noise_remote *) i;
			} else if (lookup_keypair) {
				/* Also include keypair entries. */
				kp = (struct noise_keypair *) i;
				r = kp->kp_remote;
			} else {
				break;
			}
			ret = noise_remote_ref(r);
			break;
		}
	}
	lockmgr(&l->l_index_lock, LK_RELEASE);

	return (ret);
}

struct noise_remote *
noise_remote_index(struct noise_local *l, uint32_t idx)
{
	return noise_remote_index_lookup(l, idx, true);
}

static bool
noise_remote_index_remove(struct noise_local *l, struct noise_remote *r)
{
	KKASSERT(lockstatus(&r->r_handshake_lock, curthread) == LK_EXCLUSIVE);

	if (r->r_handshake_state != HANDSHAKE_DEAD) {
		lockmgr(&l->l_index_lock, LK_EXCLUSIVE);
		r->r_handshake_state = HANDSHAKE_DEAD;
		LIST_REMOVE(&r->r_index, i_entry);
		lockmgr(&l->l_index_lock, LK_RELEASE);
		return (true);
	}

	return (false);
}

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

void
noise_remote_put(struct noise_remote *r)
{
	if (refcount_release(&r->r_refcnt)) {
#ifdef INVARIANTS
		lockmgr(&noise_mtx, LK_EXCLUSIVE);
		LIST_REMOVE(r, _r_entry);
		lockmgr(&noise_mtx, LK_RELEASE);
#endif

		noise_local_put(r->r_local);
		lockuninit(&r->r_handshake_lock);
		lockuninit(&r->r_keypair_lock);
		explicit_bzero(r, sizeof(*r));
		kfree(r, M_NOISE);
	}
}

void
noise_remote_free(struct noise_remote *r)
{
	noise_remote_disable(r);
	noise_remote_handshake_clear(r);
	noise_remote_keypairs_clear(r);
	noise_remote_put(r);
}

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])
{
	lockmgr(&r->r_handshake_lock, LK_EXCLUSIVE);
	if (psk == NULL)
		bzero(r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
	else
		memcpy(r->r_psk, psk, NOISE_SYMMETRIC_KEY_LEN);
	lockmgr(&r->r_handshake_lock, LK_RELEASE);
}

bool
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];
	bool has_psk = false;

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

	lockmgr(&r->r_handshake_lock, LK_SHARED);
	if (timingsafe_bcmp(r->r_psk, null_psk, NOISE_SYMMETRIC_KEY_LEN) != 0) {
		has_psk = true;
		if (psk != NULL)
			memcpy(psk, r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
	}
	lockmgr(&r->r_handshake_lock, LK_RELEASE);

	return (has_psk);
}

bool
noise_remote_initiation_expired(struct noise_remote *r)
{
	bool expired;

	lockmgr(&r->r_handshake_lock, LK_SHARED);
	expired = timer_expired(&r->r_last_sent, REKEY_TIMEOUT, 0);
	lockmgr(&r->r_handshake_lock, LK_RELEASE);

	return (expired);
}

void
noise_remote_handshake_clear(struct noise_remote *r)
{
	lockmgr(&r->r_handshake_lock, LK_EXCLUSIVE);
	if (noise_remote_index_remove(r->r_local, r))
		bzero(&r->r_handshake, sizeof(r->r_handshake));
	bzero(&r->r_last_sent, sizeof(r->r_last_sent));
	lockmgr(&r->r_handshake_lock, LK_RELEASE);
}

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

	lockmgr(&r->r_keypair_lock, LK_EXCLUSIVE);

	/*
	 * We zero the "next" keypair before zeroing the others, so that
	 * noise_keypair_received_with() returns early before subsequent
	 * ones are zeroed.
	 */
	kp = atomic_load_ptr(&r->r_keypair_next);
	atomic_store_ptr(&r->r_keypair_next, NULL);
	noise_keypair_drop(kp);

	kp = atomic_load_ptr(&r->r_keypair_current);
	atomic_store_ptr(&r->r_keypair_current, NULL);
	noise_keypair_drop(kp);

	kp = atomic_load_ptr(&r->r_keypair_previous);
	atomic_store_ptr(&r->r_keypair_previous, NULL);
	noise_keypair_drop(kp);

	lockmgr(&r->r_keypair_lock, LK_RELEASE);
}

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

	noise_remote_handshake_clear(r);

	lockmgr(&r->r_keypair_lock, LK_SHARED);
	kp = atomic_load_ptr(&r->r_keypair_next);
	if (kp != NULL)
		atomic_store_bool(&kp->kp_can_send, false);
	kp = atomic_load_ptr(&r->r_keypair_current);
	if (kp != NULL)
		atomic_store_bool(&kp->kp_can_send, false);
	lockmgr(&r->r_keypair_lock, LK_RELEASE);
}

/*----------------------------------------------------------------------------*/
/* Keypair functions */

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;

	idx = idx0 & HT_INDEX_MASK;

	lockmgr(&l->l_index_lock, LK_SHARED);
	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;
			ret = noise_keypair_ref(kp);
			break;
		}
	}
	lockmgr(&l->l_index_lock, LK_RELEASE);

	return (ret);
}

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

	lockmgr(&r->r_keypair_lock, LK_SHARED);
	kp = atomic_load_ptr(&r->r_keypair_current);
	if (kp != NULL && atomic_load_bool(&kp->kp_can_send)) {
		if (timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME, 0))
			atomic_store_bool(&kp->kp_can_send, false);
		else
			ret = noise_keypair_ref(kp);
	}
	lockmgr(&r->r_keypair_lock, LK_RELEASE);

	return (ret);
}

/*
 * Check whether the keypair <kp> associated with the received packet
 * is the "next" keypair of the remote peer.  If true, it means the
 * keypair has received the confirmation from the initiator and thus
 * becomes the "current".
 */
bool
noise_keypair_received_with(struct noise_keypair *kp)
{
	struct noise_keypair *old;
	struct noise_remote *r = kp->kp_remote;

	if (kp->kp_is_initiator)
		return (false);
	if (kp != atomic_load_ptr(&r->r_keypair_next))
		return (false);

	lockmgr(&r->r_keypair_lock, LK_EXCLUSIVE);

	/* Double check after locking. */
	if (kp != atomic_load_ptr(&r->r_keypair_next)) {
		lockmgr(&r->r_keypair_lock, LK_RELEASE);
		return (false);
	}

	/*
	 * Received the confirming data packet, so move the keypair from
	 * the "next" slot to the "current".
	 */
	old = atomic_load_ptr(&r->r_keypair_previous);
	atomic_store_ptr(&r->r_keypair_previous,
			 atomic_load_ptr(&r->r_keypair_current));
	noise_keypair_drop(old);
	atomic_store_ptr(&r->r_keypair_current, kp);
	atomic_store_ptr(&r->r_keypair_next, NULL);

	lockmgr(&r->r_keypair_lock, LK_RELEASE);

	return (true);
}

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

void
noise_keypair_put(struct noise_keypair *kp)
{
	if (refcount_release(&kp->kp_refcnt)) {
#ifdef INVARIANTS
		lockmgr(&noise_mtx, LK_EXCLUSIVE);
		LIST_REMOVE(kp, _kp_entry);
		lockmgr(&noise_mtx, LK_RELEASE);
#endif

		noise_remote_put(kp->kp_remote);
		lockuninit(&kp->kp_counter_lock);
		explicit_bzero(kp, sizeof(*kp));
		kfree(kp, M_NOISE);
	}
}

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;

	lockmgr(&l->l_index_lock, LK_EXCLUSIVE);
	LIST_REMOVE(&kp->kp_index, i_entry);
	lockmgr(&l->l_index_lock, LK_RELEASE);

	KKASSERT(lockstatus(&r->r_keypair_lock, curthread) == LK_EXCLUSIVE);
	noise_keypair_put(kp);
}

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

bool
noise_keypair_counter_next(struct noise_keypair *kp, uint64_t *send)
{
	if (!atomic_load_bool(&kp->kp_can_send))
		return (false);

#ifdef __LP64__
	*send = atomic_fetchadd_64(&kp->kp_counter_send, 1);
#else
	lockmgr(&kp->kp_counter_lock, LK_EXCLUSIVE);
	*send = kp->kp_counter_send++;
	lockmgr(&kp->kp_counter_lock, LK_RELEASE);
#endif
	if (*send >= REJECT_AFTER_MESSAGES) {
		atomic_store_bool(&kp->kp_can_send, false);
		return (false);
	}

	return (true);
}

/*
 * Validate the received counter to avoid replay attacks.  A sliding window
 * is used to keep track of the received counters, since the UDP messages
 * can arrive out of order.
 *
 * NOTE: Validate the counter only *after* successful decryption, which
 *       ensures that the message and counter is authentic.
 *
 * This implements the algorithm from RFC 6479:
 * "IPsec Anti-Replay Algorithm without Bit Shifting"
 */
int
noise_keypair_counter_check(struct noise_keypair *kp, uint64_t recv)
{
	unsigned long index, index_current, top, i, bit;
	int ret;

	lockmgr(&kp->kp_counter_lock, LK_EXCLUSIVE);

	if (__predict_false(kp->kp_counter_recv >= REJECT_AFTER_MESSAGES ||
			    recv >= REJECT_AFTER_MESSAGES)) {
		ret = EINVAL;
		goto out;
	}

	if (__predict_false(recv + COUNTER_WINDOW_SIZE < kp->kp_counter_recv)) {
		ret = ESTALE;
		goto out;
	}

	index = recv >> COUNTER_ORDER;

	if (__predict_true(recv > kp->kp_counter_recv)) {
		/*
		 * The new counter is ahead of the current counter, so need
		 * to zero out the bitmap that has previously been used.
		 */
		index_current = kp->kp_counter_recv >> COUNTER_ORDER;
		top = MIN(index - index_current, COUNTER_NUM);
		for (i = 1; i <= top; i++)
			kp->kp_backtrack[(i+index_current) & COUNTER_MASK] = 0;
#ifdef __LP64__
		atomic_store_64(&kp->kp_counter_recv, recv);
#else
		kp->kp_counter_recv = recv;
#endif
	}

	index &= COUNTER_MASK;
	bit = 1UL << (recv & (COUNTER_BITS - 1));
	if (kp->kp_backtrack[index] & bit) {
		ret = EEXIST;
		goto out;
	}

	kp->kp_backtrack[index] |= bit;
	ret = 0;

out:
	lockmgr(&kp->kp_counter_lock, LK_RELEASE);
	return (ret);
}

/*
 * Check whether the current keypair of the given remote <r> is expiring soon
 * or already expired, and thus should do a refreshing.
 */
bool
noise_keypair_should_refresh(struct noise_remote *r, bool sending)
{
	struct noise_keypair *kp;
	uint64_t counter;
	bool refresh;

	lockmgr(&r->r_keypair_lock, LK_SHARED);

	kp = atomic_load_ptr(&r->r_keypair_current);
	refresh = (kp != NULL && atomic_load_bool(&kp->kp_can_send));
	if (__predict_false(!refresh))
		goto out;

	if (sending) {
		/* sending path */
#ifdef __LP64__
		counter = atomic_load_64(&kp->kp_counter_send);
#else
		lockmgr(&kp->kp_counter_lock, LK_SHARED);
		counter = kp->kp_counter_send;
		lockmgr(&kp->kp_counter_lock, LK_RELEASE);
#endif
		refresh = (counter > REKEY_AFTER_MESSAGES ||
			   (kp->kp_is_initiator &&
			    timer_expired(&kp->kp_birthdate,
					  REKEY_AFTER_TIME, 0)));
	} else {
		/* receiving path */
		refresh = (kp->kp_is_initiator &&
			   timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME -
					 KEEPALIVE_TIMEOUT - REKEY_TIMEOUT, 0));
	}

out:
	lockmgr(&r->r_keypair_lock, LK_RELEASE);
	return (refresh);
}

int
noise_keypair_encrypt(struct noise_keypair *kp, uint32_t *r_idx,
		      uint64_t counter, struct mbuf *m)
{
	uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE];
	int ret;

	/* 32 bits of zeros + 64-bit little-endian value of the counter */
	*(uint32_t *)nonce = 0;
	*(uint64_t *)(nonce + 4) = htole64(counter);

	ret = chacha20poly1305_encrypt_mbuf(m, NULL, 0, nonce, kp->kp_send);
	if (ret == 0)
		*r_idx = kp->kp_index.i_remote_index;

	return (ret);
}

int
noise_keypair_decrypt(struct noise_keypair *kp, uint64_t counter,
		      struct mbuf *m)
{
	uint64_t cur_counter;
	uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE];

#ifdef __LP64__
	cur_counter = atomic_load_64(&kp->kp_counter_recv);
#else
	lockmgr(&kp->kp_counter_lock, LK_SHARED);
	cur_counter = kp->kp_counter_recv;
	lockmgr(&kp->kp_counter_lock, LK_RELEASE);
#endif
	if (cur_counter >= REJECT_AFTER_MESSAGES ||
	    timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME, 0))
		return (EINVAL);

	*(uint32_t *)nonce = 0;
	*(uint64_t *)(nonce + 4) = htole64(counter);

	return chacha20poly1305_decrypt_mbuf(m, NULL, 0, nonce, kp->kp_recv);
}

/*----------------------------------------------------------------------------*/
/* Handshake functions */

bool
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];
	bool ok = false;

	lockmgr(&l->l_identity_lock, LK_SHARED);
	lockmgr(&r->r_handshake_lock, LK_EXCLUSIVE);

	if (!l->l_has_identity)
		goto error;
	if (!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))
		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))
		goto error;

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

	*s_idx = noise_remote_index_insert(l, r);
	r->r_handshake_state = HANDSHAKE_INITIATOR;
	getnanouptime(&r->r_last_sent);
	ok = true;

error:
	lockmgr(&r->r_handshake_lock, LK_RELEASE);
	lockmgr(&l->l_identity_lock, LK_RELEASE);
	explicit_bzero(key, sizeof(key));
	return (ok);
}

struct noise_remote *
noise_consume_initiation(struct noise_local *l, 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, *ret = NULL;
	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];

	lockmgr(&l->l_identity_lock, LK_SHARED);

	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))
		goto error;

	/* s */
	if (!noise_msg_decrypt(r_public, es,
			       NOISE_PUBLIC_KEY_LEN + NOISE_AUTHTAG_LEN,
			       key, hs.hs_hash))
		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))
		goto error_put;

	/* {t} */
	if (!noise_msg_decrypt(timestamp, ets,
			       NOISE_TIMESTAMP_LEN + NOISE_AUTHTAG_LEN,
			       key, hs.hs_hash))
		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.
	 */
	lockmgr(&r->r_handshake_lock, LK_EXCLUSIVE);

	/* 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 (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_state = HANDSHAKE_RESPONDER;
	r->r_handshake = hs;
	ret = noise_remote_ref(r);

error_set:
	lockmgr(&r->r_handshake_lock, LK_RELEASE);
error_put:
	noise_remote_put(r);
error:
	lockmgr(&l->l_identity_lock, LK_RELEASE);
	explicit_bzero(key, sizeof(key));
	explicit_bzero(&hs, sizeof(hs));
	return (ret);
}

bool
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];
	bool ok = false;

	lockmgr(&l->l_identity_lock, LK_SHARED);
	lockmgr(&r->r_handshake_lock, LK_EXCLUSIVE);

	if (r->r_handshake_state != HANDSHAKE_RESPONDER)
		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))
		goto error;

	/* se */
	if (!noise_mix_dh(hs->hs_ck, NULL, e, r->r_public))
		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 (noise_begin_session(r)) {
		getnanouptime(&r->r_last_sent);
		*s_idx = r->r_index.i_local_index;
		*r_idx = r->r_index.i_remote_index;
		ok = true;
	}

error:
	lockmgr(&r->r_handshake_lock, LK_RELEASE);
	lockmgr(&l->l_identity_lock, LK_RELEASE);
	explicit_bzero(key, sizeof(key));
	explicit_bzero(e, sizeof(e));
	return (ok);
}

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

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

	lockmgr(&l->l_identity_lock, LK_SHARED);
	if (!l->l_has_identity)
		goto error;

	lockmgr(&r->r_handshake_lock, LK_SHARED);
	if (r->r_handshake_state != HANDSHAKE_INITIATOR) {
		lockmgr(&r->r_handshake_lock, LK_RELEASE);
		goto error;
	}
	memcpy(preshared_key, r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
	hs = r->r_handshake;
	lockmgr(&r->r_handshake_lock, LK_RELEASE);

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

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

	/* se */
	if (!noise_mix_dh(hs.hs_ck, NULL, l->l_private, ue))
		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))
		goto error_zero;

	lockmgr(&r->r_handshake_lock, LK_EXCLUSIVE);
	if (r->r_handshake_state == HANDSHAKE_INITIATOR &&
	    r->r_index.i_local_index == r_idx) {
		r->r_handshake = hs;
		r->r_index.i_remote_index = s_idx;
		if (noise_begin_session(r))
			ret = noise_remote_ref(r);
	}
	lockmgr(&r->r_handshake_lock, LK_RELEASE);

error_zero:
	explicit_bzero(preshared_key, sizeof(preshared_key));
	explicit_bzero(key, sizeof(key));
	explicit_bzero(&hs, sizeof(hs));
error:
	lockmgr(&l->l_identity_lock, LK_RELEASE);
	noise_remote_put(r);
	return (ret);
}

/*----------------------------------------------------------------------------*/
/* Handshake helper functions */

static bool
noise_begin_session(struct noise_remote *r)
{
	struct noise_local *l = r->r_local;
	struct noise_keypair *kp, *next, *current, *previous;
	struct noise_index *r_i;

	KKASSERT(lockstatus(&r->r_handshake_lock, curthread) == LK_EXCLUSIVE);

	kp = kmalloc(sizeof(*kp), M_NOISE, M_NOWAIT | M_ZERO);
	if (kp == NULL)
		return (false);

	/*
	 * Initialize the new keypair.
	 */
	refcount_init(&kp->kp_refcnt, 1);
	kp->kp_can_send = true;
	kp->kp_is_initiator = (r->r_handshake_state == HANDSHAKE_INITIATOR);
	kp->kp_remote = noise_remote_ref(r);
	getnanouptime(&kp->kp_birthdate);
	lockinit(&kp->kp_counter_lock, "noise_counter", 0, 0);

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

	/*
	 * Rotate existing keypairs and load the new one.
	 */
	lockmgr(&r->r_keypair_lock, LK_EXCLUSIVE);
	next = atomic_load_ptr(&r->r_keypair_next);
	current = atomic_load_ptr(&r->r_keypair_current);
	previous = atomic_load_ptr(&r->r_keypair_previous);
	if (kp->kp_is_initiator) {
		/*
		 * Received a confirmation response, which means this new
		 * keypair can now be used.
		 *
		 * Rotate the existing keypair ("current" or "next" slot)
		 * to the "previous" slot, and load the new keypair to the
		 * "current" slot.
		 */
		if (next != NULL) {
			atomic_store_ptr(&r->r_keypair_next, NULL);
			atomic_store_ptr(&r->r_keypair_previous, next);
			noise_keypair_drop(current);
		} else {
			atomic_store_ptr(&r->r_keypair_previous, current);
		}
		noise_keypair_drop(previous);
		atomic_store_ptr(&r->r_keypair_current, kp);
	} else {
		/*
		 * This new keypair cannot be used until we receive a
		 * confirmation via the first data packet.
		 *
		 * So drop the "previous" keypair, the possibly existing
		 * "next" one, and load the new keypair to the "next" slot.
		 */
		atomic_store_ptr(&r->r_keypair_next, kp);
		noise_keypair_drop(next);
		atomic_store_ptr(&r->r_keypair_previous, NULL);
		noise_keypair_drop(previous);
	}
	lockmgr(&r->r_keypair_lock, LK_RELEASE);

	/*
	 * Insert into index hashtable, replacing the existing remote index
	 * (added with handshake initiation creation/consumption).
	 */
	r_i = &r->r_index;
	kp->kp_index.i_is_keypair = true;
	kp->kp_index.i_local_index = r_i->i_local_index;
	kp->kp_index.i_remote_index = r_i->i_remote_index;

	KKASSERT(lockstatus(&r->r_handshake_lock, curthread) == LK_EXCLUSIVE);
	lockmgr(&l->l_index_lock, LK_EXCLUSIVE);
	LIST_INSERT_BEFORE(r_i, &kp->kp_index, i_entry);
	r->r_handshake_state = HANDSHAKE_DEAD;
	LIST_REMOVE(r_i, i_entry);
	lockmgr(&l->l_index_lock, LK_RELEASE);

#ifdef INVARIANTS
	lockmgr(&noise_mtx, LK_EXCLUSIVE);
	LIST_INSERT_HEAD(&noise_keypairs, kp, _kp_entry);
	lockmgr(&noise_mtx, LK_RELEASE);
#endif

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

static void
noise_hmac(uint8_t *out, const uint8_t *in, const uint8_t *key,
	   size_t outlen, size_t inlen, size_t keylen)
{
	struct blake2s_state blake;
	uint8_t x_key[BLAKE2S_BLOCK_SIZE] __aligned(__alignof__(uint32_t));
	uint8_t i_hash[BLAKE2S_HASH_SIZE] __aligned(__alignof__(uint32_t));
	int i;

	KKASSERT(out != NULL && outlen <= BLAKE2S_HASH_SIZE &&
		 key != NULL && keylen > 0);

	memset(x_key, 0, BLAKE2S_BLOCK_SIZE);
	if (keylen > BLAKE2S_BLOCK_SIZE) {
		blake2s_init(&blake, BLAKE2S_HASH_SIZE);
		blake2s_update(&blake, key, keylen);
		blake2s_final(&blake, x_key);
	} else {
		memcpy(x_key, key, keylen);
	}

	for (i = 0; i < BLAKE2S_BLOCK_SIZE; ++i)
		x_key[i] ^= 0x36;

	blake2s_init(&blake, BLAKE2S_HASH_SIZE);
	blake2s_update(&blake, x_key, BLAKE2S_BLOCK_SIZE);
	blake2s_update(&blake, in, inlen);
	blake2s_final(&blake, i_hash);

	for (i = 0; i < BLAKE2S_BLOCK_SIZE; ++i)
		x_key[i] ^= 0x5c ^ 0x36;

	blake2s_init(&blake, BLAKE2S_HASH_SIZE);
	blake2s_update(&blake, x_key, BLAKE2S_BLOCK_SIZE);
	blake2s_update(&blake, i_hash, BLAKE2S_HASH_SIZE);
	blake2s_final(&blake, i_hash);

	memcpy(out, i_hash, outlen);
	explicit_bzero(x_key, BLAKE2S_BLOCK_SIZE);
	explicit_bzero(i_hash, BLAKE2S_HASH_SIZE);
}

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

	KKASSERT(a != NULL && a_len > 0);
	KKASSERT(a_len <= BLAKE2S_HASH_SIZE &&
		 b_len <= BLAKE2S_HASH_SIZE &&
		 c_len <= BLAKE2S_HASH_SIZE);

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

	/* Expand first key: key = sec, data = 0x1 */
	out[0] = 1;
	noise_hmac(out, out, sec, BLAKE2S_HASH_SIZE /* outlen */,
		   1 /* inlen */, 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;
	noise_hmac(out, out, sec, BLAKE2S_HASH_SIZE /* outlen */,
		   BLAKE2S_HASH_SIZE + 1 /* inlen */, 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;
	noise_hmac(out, out, sec, BLAKE2S_HASH_SIZE /* outlen */,
		   BLAKE2S_HASH_SIZE + 1 /* inlen */, BLAKE2S_HASH_SIZE);
	memcpy(c, out, c_len);

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

static bool
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 (false);

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

static bool
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 (false);

	noise_kdf(ck, key, NULL, ss,
		  NOISE_HASH_LEN, NOISE_SYMMETRIC_KEY_LEN,
		  0, NOISE_PUBLIC_KEY_LEN, ck);
	return (true);
}

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, NOISE_HANDSHAKE_NAME, NULL,
		NOISE_HASH_LEN, sizeof(NOISE_HANDSHAKE_NAME) - 1, 0);
	blake2s_init(&blake, NOISE_HASH_LEN);
	blake2s_update(&blake, ck, NOISE_HASH_LEN);
	blake2s_update(&blake, NOISE_IDENTIFIER_NAME,
		       sizeof(NOISE_IDENTIFIER_NAME) - 1);
	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 */
	static const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE] = { 0 };

	chacha20poly1305_encrypt(dst, src, src_len, hash, NOISE_HASH_LEN,
				 nonce, key);
	noise_mix_hash(hash, dst, src_len + NOISE_AUTHTAG_LEN);
}

static bool
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 */
	static const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE] = { 0 };

	if (!chacha20poly1305_decrypt(dst, src, src_len,
				      hash, NOISE_HASH_LEN, nonce, key))
		return (false);

	noise_mix_hash(hash, src, src_len);
	return (true);
}

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

#ifdef WG_SELFTESTS
#include "selftest/counter.c"
#endif /* WG_SELFTESTS */