if_wg: import latest wg_noise.{c,h}

Note: this is a partial diff, introducing temporary bugs that will be
resolved in following commits, detailed below.

This commit brings wg_noise.{c,h} up to date with wireguard-openbsd. The
primary motivator for this large patchset is to allow checking nonces
serial, requiring a reference to the receiving keypair across noise_*
calls. Due to requiring reference counting on the keypairs, we also take
this opportunity to throw away the old locking and bring in EPOCH
(roughly equivalent to SMR on OpenBSD and RCU on Linux).

The changes to if_wg.c are purely to allow it to compile, there are most
certainly refcount leaks present (to be addressed in the following
commits). Readers should review wg_noise.{c,h} in their entirety rather
than the diffs, as there are significant changes. if_wg.c can be
reviewed, but must be contextualised with the following commits
(repace wg_tag with wg_packet, encrypt mbuf in place).

Signed-off-by: Matt Dunwoodie <ncon@noconroy.net>

5 files changed, 1137 insertions, 1068 deletions

diff --git a/TODO.md b/TODO.md
index b525083..0f4e634 100644
--- a/TODO.md
+++ b/TODO.md
@@ -12,8 +12,6 @@
 - Resize mbufs once at the beginning, and then encrypt/decrypt in place, rather
   than making a new mbuf and copying. (Remember to clear the tags and other
   pieces of metadata before passing it off to udp sending or netisr receiving.)
-- Check nonces in serial, rather than in parallel. (This requires taking a
-  keypair reference; ncon is working on it.)
 - Audit allowedips / radix tree checks, and make sure it's actually behaving as
   expected. (It might be useful to port [this selftest](https://git.zx2c4.com/wireguard-linux/tree/drivers/net/wireguard/selftest/allowedips.c).)
 - Make code style consistent with one FreeBSD way, rather than a mix of styles.
diff --git a/src/if_wg.c b/src/if_wg.c
index 6f4d225..70f6b4b 100644
--- a/src/if_wg.c
+++ b/src/if_wg.c
@@ -89,13 +89,7 @@ __FBSDID("$FreeBSD$");
 
 #define MAX_QUEUED_HANDSHAKES	4096
 
-#define HASHTABLE_PEER_SIZE	(1 << 11)
-#define HASHTABLE_INDEX_SIZE	(1 << 13)
-#define MAX_PEERS_PER_IFACE	(1 << 20)
-
-#define REKEY_TIMEOUT		5
 #define REKEY_TIMEOUT_JITTER	334 /* 1/3 sec, round for arc4random_uniform */
-#define KEEPALIVE_TIMEOUT	10
 #define MAX_TIMER_HANDSHAKES	(90 / REKEY_TIMEOUT)
 #define NEW_HANDSHAKE_TIMEOUT	(REKEY_TIMEOUT + KEEPALIVE_TIMEOUT)
 #define UNDERLOAD_TIMEOUT	1
@@ -163,19 +157,13 @@ struct wg_endpoint {
 struct wg_tag {
 	struct m_tag		 t_tag;
 	struct wg_endpoint	 t_endpoint;
-	struct wg_peer		*t_peer;
+	struct noise_keypair	*t_keypair;
+	uint64_t		 t_nonce;
 	struct mbuf		*t_mbuf;
 	int			 t_done;
 	int			 t_mtu;
 };
 
-struct wg_index {
-	LIST_ENTRY(wg_index)	 i_entry;
-	SLIST_ENTRY(wg_index)	 i_unused_entry;
-	uint32_t		 i_key;
-	struct noise_remote	*i_value;
-};
-
 struct wg_timers {
 	/* t_lock is for blocking wg_timers_event_* when setting t_disabled. */
 	struct rwlock		 t_lock;
@@ -190,7 +178,6 @@ struct wg_timers {
 	struct callout		 t_persistent_keepalive;
 
 	struct mtx		 t_handshake_mtx;
-	struct timespec		 t_handshake_last_sent;
 	struct timespec		 t_handshake_complete;
 	volatile int		 t_handshake_retries;
 };
@@ -209,26 +196,21 @@ struct wg_queue {
 };
 
 struct wg_peer {
-	CK_LIST_ENTRY(wg_peer)		 p_hash_entry;
-	CK_LIST_ENTRY(wg_peer)		 p_entry;
+	TAILQ_ENTRY(wg_peer)		 p_entry;
 	uint64_t			 p_id;
 	struct wg_softc			*p_sc;
 
-	struct noise_remote		 p_remote;
+	struct noise_remote		*p_remote;
 	struct cookie_maker		 p_cookie;
 	struct wg_timers		 p_timers;
 
 	struct rwlock			 p_endpoint_lock;
 	struct wg_endpoint		 p_endpoint;
 
-	SLIST_HEAD(,wg_index)		 p_unused_index;
-	struct wg_index			 p_index[3];
-
 	struct wg_queue	 		 p_stage_queue;
 	struct wg_queue	 		 p_encap_queue;
 	struct wg_queue	 		 p_decap_queue;
 
-	struct grouptask		 p_clear_secrets;
 	struct grouptask		 p_send_initiation;
 	struct grouptask		 p_send_keepalive;
 	struct grouptask		 p_send;
@@ -238,8 +220,6 @@ struct wg_peer {
 	counter_u64_t			 p_rx_bytes;
 
 	CK_LIST_HEAD(, wg_aip)		 p_aips;
-	struct mtx			 p_lock;
-	struct epoch_context		 p_ctx;
 };
 
 enum route_direction {
@@ -264,15 +244,6 @@ struct wg_allowedip {
 	uint8_t cidr;
 };
 
-struct wg_hashtable {
-	struct mtx			 h_mtx;
-	SIPHASH_KEY			 h_secret;
-	CK_LIST_HEAD(, wg_peer)		 h_peers_list;
-	CK_LIST_HEAD(, wg_peer)		*h_peers;
-	u_long				 h_peers_mask;
-	size_t				 h_num_peers;
-};
-
 struct wg_socket {
 	struct mtx	 so_mtx;
 	struct socket	*so_so4;
@@ -289,13 +260,15 @@ struct wg_softc {
 
 	struct ucred		*sc_ucred;
 	struct wg_socket	 sc_socket;
-	struct wg_hashtable	 sc_hashtable;
 	struct wg_aip_table	 sc_aips;
 
+	TAILQ_HEAD(,wg_peer)	 sc_peers;
+	size_t			 sc_peers_num;
+
 	struct mbufq		 sc_handshake_queue;
 	struct grouptask	 sc_handshake;
 
-	struct noise_local	 sc_local;
+	struct noise_local	*sc_local;
 	struct cookie_checker	 sc_cookie;
 
 	struct buf_ring		*sc_encap_ring;
@@ -304,12 +277,7 @@ struct wg_softc {
 	struct grouptask	*sc_encrypt;
 	struct grouptask	*sc_decrypt;
 
-	struct rwlock		 sc_index_lock;
-	LIST_HEAD(,wg_index)	*sc_index;
-	u_long			 sc_index_mask;
-
 	struct sx		 sc_lock;
-	volatile u_int		 sc_peer_count;
 };
 
 #define	WGF_DYING	0x0001
@@ -369,11 +337,9 @@ static void wg_timers_event_any_authenticated_packet_sent(struct wg_timers *);
 static void wg_timers_event_any_authenticated_packet_received(struct wg_timers *);
 static void wg_timers_event_any_authenticated_packet_traversal(struct wg_timers *);
 static void wg_timers_event_handshake_initiated(struct wg_timers *);
-static void wg_timers_event_handshake_responded(struct wg_timers *);
 static void wg_timers_event_handshake_complete(struct wg_timers *);
 static void wg_timers_event_session_derived(struct wg_timers *);
 static void wg_timers_event_want_initiation(struct wg_timers *);
-static void wg_timers_event_reset_handshake_last_sent(struct wg_timers *);
 static void wg_timers_run_send_initiation(struct wg_timers *, int);
 static void wg_timers_run_retry_handshake(struct wg_timers *);
 static void wg_timers_run_send_keepalive(struct wg_timers *);
@@ -385,8 +351,6 @@ static void wg_timers_enable(struct wg_timers *);
 static void wg_timers_disable(struct wg_timers *);
 static void wg_timers_set_persistent_keepalive(struct wg_timers *, uint16_t);
 static void wg_timers_get_last_handshake(struct wg_timers *, struct wg_timespec64 *);
-static int wg_timers_expired_handshake_last_sent(struct wg_timers *);
-static int wg_timers_check_handshake_last_sent(struct wg_timers *);
 static void wg_queue_init(struct wg_queue *, const char *);
 static void wg_queue_deinit(struct wg_queue *);
 static void wg_queue_purge(struct wg_queue *);
@@ -404,14 +368,9 @@ static int wg_peer_remove(struct radix_node *, void *);
 static void wg_peer_remove_all(struct wg_softc *);
 static int wg_aip_delete(struct wg_aip_table *, struct wg_peer *);
 static struct wg_peer *wg_aip_lookup(struct wg_aip_table *, struct mbuf *, enum route_direction);
-static void wg_hashtable_init(struct wg_hashtable *);
-static void wg_hashtable_destroy(struct wg_hashtable *);
-static void wg_hashtable_peer_insert(struct wg_hashtable *, struct wg_peer *);
-static struct wg_peer *wg_peer_lookup(struct wg_softc *, const uint8_t [32]);
-static void wg_hashtable_peer_remove(struct wg_hashtable *, struct wg_peer *);
 static int wg_cookie_validate_packet(struct cookie_checker *, struct mbuf *, int);
 static struct wg_peer *wg_peer_alloc(struct wg_softc *);
-static void wg_peer_free_deferred(epoch_context_t);
+static void wg_peer_free_deferred(struct noise_remote *);
 static void wg_peer_destroy(struct wg_peer *);
 static void wg_peer_send_buf(struct wg_peer *, uint8_t *, size_t);
 static void wg_send_initiation(struct wg_peer *);
@@ -430,10 +389,6 @@ static void wg_decap(struct wg_softc *, struct mbuf *);
 static void wg_softc_handshake_receive(struct wg_softc *);
 static void wg_softc_decrypt(struct wg_softc *);
 static void wg_softc_encrypt(struct wg_softc *);
-static struct noise_remote *wg_remote_get(struct wg_softc *, uint8_t [NOISE_PUBLIC_KEY_LEN]);
-static uint32_t wg_index_set(struct wg_softc *, struct noise_remote *);
-static struct noise_remote *wg_index_get(struct wg_softc *, uint32_t);
-static void wg_index_drop(struct wg_softc *, uint32_t);
 static int wg_update_endpoint_addrs(struct wg_endpoint *, const struct sockaddr *, struct ifnet *);
 static void wg_input(struct mbuf *, int, struct inpcb *, const struct sockaddr *, void *);
 static void wg_encrypt_dispatch(struct wg_softc *);
@@ -480,9 +435,6 @@ wg_peer_alloc(struct wg_softc *sc)
 	taskqgroup_attach(qgroup_wg_tqg, &peer->p_send_initiation, peer, NULL, NULL, "wg initiation");
 	GROUPTASK_INIT(&peer->p_send_keepalive, 0, (gtask_fn_t *)wg_send_keepalive, peer);
 	taskqgroup_attach(qgroup_wg_tqg, &peer->p_send_keepalive, peer, NULL, NULL, "wg keepalive");
-	GROUPTASK_INIT(&peer->p_clear_secrets, 0, (gtask_fn_t *)noise_remote_clear, &peer->p_remote);
-	taskqgroup_attach(qgroup_wg_tqg, &peer->p_clear_secrets,
-	    &peer->p_remote, NULL, NULL, "wg clear secrets");
 
 	GROUPTASK_INIT(&peer->p_send, 0, (gtask_fn_t *)wg_deliver_out, peer);
 	taskqgroup_attach(qgroup_wg_tqg, &peer->p_send, peer, NULL, NULL, "wg send");
@@ -494,91 +446,13 @@ wg_peer_alloc(struct wg_softc *sc)
 	peer->p_tx_bytes = counter_u64_alloc(M_WAITOK);
 	peer->p_rx_bytes = counter_u64_alloc(M_WAITOK);
 
-	SLIST_INIT(&peer->p_unused_index);
-	SLIST_INSERT_HEAD(&peer->p_unused_index, &peer->p_index[0],
-	    i_unused_entry);
-	SLIST_INSERT_HEAD(&peer->p_unused_index, &peer->p_index[1],
-	    i_unused_entry);
-	SLIST_INSERT_HEAD(&peer->p_unused_index, &peer->p_index[2],
-	    i_unused_entry);
-
 	return (peer);
 }
 
-#define WG_HASHTABLE_PEER_FOREACH(peer, i, ht) \
-	for (i = 0; i < HASHTABLE_PEER_SIZE; i++) \
-		LIST_FOREACH(peer, &(ht)->h_peers[i], p_hash_entry)
-#define WG_HASHTABLE_PEER_FOREACH_SAFE(peer, i, ht, tpeer) \
-	for (i = 0; i < HASHTABLE_PEER_SIZE; i++) \
-		CK_LIST_FOREACH_SAFE(peer, &(ht)->h_peers[i], p_hash_entry, tpeer)
-static void
-wg_hashtable_init(struct wg_hashtable *ht)
-{
-	mtx_init(&ht->h_mtx, "hash lock", NULL, MTX_DEF);
-	arc4random_buf(&ht->h_secret, sizeof(ht->h_secret));
-	ht->h_num_peers = 0;
-	ht->h_peers = hashinit(HASHTABLE_PEER_SIZE, M_DEVBUF,
-			&ht->h_peers_mask);
-}
-
-static void
-wg_hashtable_destroy(struct wg_hashtable *ht)
-{
-	MPASS(ht->h_num_peers == 0);
-	mtx_destroy(&ht->h_mtx);
-	hashdestroy(ht->h_peers, M_DEVBUF, ht->h_peers_mask);
-}
-
 static void
-wg_hashtable_peer_insert(struct wg_hashtable *ht, struct wg_peer *peer)
+wg_peer_free_deferred(struct noise_remote *r)
 {
-	uint64_t key;
-
-	key = siphash24(&ht->h_secret, peer->p_remote.r_public,
-			sizeof(peer->p_remote.r_public));
-
-	mtx_lock(&ht->h_mtx);
-	ht->h_num_peers++;
-	CK_LIST_INSERT_HEAD(&ht->h_peers[key & ht->h_peers_mask], peer, p_hash_entry);
-	CK_LIST_INSERT_HEAD(&ht->h_peers_list, peer, p_entry);
-	mtx_unlock(&ht->h_mtx);
-}
-
-static struct wg_peer *
-wg_peer_lookup(struct wg_softc *sc,
-    const uint8_t pubkey[WG_KEY_SIZE])
-{
-	struct wg_hashtable *ht = &sc->sc_hashtable;
-	uint64_t key;
-	struct wg_peer *i = NULL;
-
-	key = siphash24(&ht->h_secret, pubkey, WG_KEY_SIZE);
-
-	mtx_lock(&ht->h_mtx);
-	CK_LIST_FOREACH(i, &ht->h_peers[key & ht->h_peers_mask], p_hash_entry) {
-		if (timingsafe_bcmp(i->p_remote.r_public, pubkey,
-					WG_KEY_SIZE) == 0)
-			break;
-	}
-	mtx_unlock(&ht->h_mtx);
-
-	return i;
-}
-
-static void
-wg_hashtable_peer_remove(struct wg_hashtable *ht, struct wg_peer *peer)
-{
-	mtx_lock(&ht->h_mtx);
-	ht->h_num_peers--;
-	CK_LIST_REMOVE(peer, p_hash_entry);
-	CK_LIST_REMOVE(peer, p_entry);
-	mtx_unlock(&ht->h_mtx);
-}
-
-static void
-wg_peer_free_deferred(epoch_context_t ctx)
-{
-	struct wg_peer *peer = __containerof(ctx, struct wg_peer, p_ctx);
+	struct wg_peer *peer = noise_remote_arg(r);
 	counter_u64_free(peer->p_tx_bytes);
 	counter_u64_free(peer->p_rx_bytes);
 	rw_destroy(&peer->p_timers.t_lock);
@@ -599,13 +473,11 @@ wg_peer_destroy(struct wg_peer *peer)
 	wg_timers_disable(&peer->p_timers);
 
 	/* Ensure the tasks have finished running */
-	GROUPTASK_DRAIN(&peer->p_clear_secrets);
 	GROUPTASK_DRAIN(&peer->p_send_initiation);
 	GROUPTASK_DRAIN(&peer->p_send_keepalive);
 	GROUPTASK_DRAIN(&peer->p_recv);
 	GROUPTASK_DRAIN(&peer->p_send);
 
-	taskqgroup_detach(qgroup_wg_tqg, &peer->p_clear_secrets);
 	taskqgroup_detach(qgroup_wg_tqg, &peer->p_send_initiation);
 	taskqgroup_detach(qgroup_wg_tqg, &peer->p_send_keepalive);
 	taskqgroup_detach(qgroup_wg_tqg, &peer->p_recv);
@@ -616,10 +488,10 @@ wg_peer_destroy(struct wg_peer *peer)
 	wg_queue_deinit(&peer->p_stage_queue);
 
 	/* Final cleanup */
-	--peer->p_sc->sc_peer_count;
-	noise_remote_clear(&peer->p_remote);
+	peer->p_sc->sc_peers_num--;
+	TAILQ_REMOVE(&peer->p_sc->sc_peers, peer, p_entry);
 	DPRINTF(peer->p_sc, "Peer %llu destroyed\n", (unsigned long long)peer->p_id);
-	NET_EPOCH_CALL(wg_peer_free_deferred, &peer->p_ctx);
+	noise_remote_free(peer->p_remote, wg_peer_free_deferred);
 }
 
 static void
@@ -888,11 +760,8 @@ wg_peer_remove_all(struct wg_softc *sc)
 
 	sx_assert(&sc->sc_lock, SX_XLOCKED);
 
-	CK_LIST_FOREACH_SAFE(peer, &sc->sc_hashtable.h_peers_list,
-			     p_entry, tpeer) {
-		wg_hashtable_peer_remove(&sc->sc_hashtable, peer);
+	TAILQ_FOREACH_SAFE(peer, &sc->sc_peers, p_entry, tpeer)
 		wg_peer_destroy(peer);
-	}
 }
 
 static int
@@ -1253,29 +1122,6 @@ wg_timers_get_last_handshake(struct wg_timers *t, struct wg_timespec64 *time)
 	rw_runlock(&t->t_lock);
 }
 
-static int
-wg_timers_expired_handshake_last_sent(struct wg_timers *t)
-{
-	struct timespec uptime;
-	struct timespec expire = { .tv_sec = REKEY_TIMEOUT, .tv_nsec = 0 };
-
-	getnanouptime(&uptime);
-	timespecadd(&t->t_handshake_last_sent, &expire, &expire);
-	return timespeccmp(&uptime, &expire, >) ? ETIMEDOUT : 0;
-}
-
-static int
-wg_timers_check_handshake_last_sent(struct wg_timers *t)
-{
-	int ret;
-
-	rw_wlock(&t->t_lock);
-	if ((ret = wg_timers_expired_handshake_last_sent(t)) == ETIMEDOUT)
-		getnanouptime(&t->t_handshake_last_sent);
-	rw_wunlock(&t->t_lock);
-	return (ret);
-}
-
 /* Should be called after an authenticated data packet is sent. */
 static void
 wg_timers_event_data_sent(struct wg_timers *t)
@@ -1354,14 +1200,6 @@ wg_timers_event_handshake_initiated(struct wg_timers *t)
 	rw_runlock(&t->t_lock);
 }
 
-static void
-wg_timers_event_handshake_responded(struct wg_timers *t)
-{
-	rw_wlock(&t->t_lock);
-	getnanouptime(&t->t_handshake_last_sent);
-	rw_wunlock(&t->t_lock);
-}
-
 /*
  * Should be called after a handshake response message is received and processed
  * or when getting key confirmation via the first data message.
@@ -1405,20 +1243,12 @@ wg_timers_event_want_initiation(struct wg_timers *t)
 }
 
 static void
-wg_timers_event_reset_handshake_last_sent(struct wg_timers *t)
-{
-	rw_wlock(&t->t_lock);
-	t->t_handshake_last_sent.tv_sec -= (REKEY_TIMEOUT + 1);
-	rw_wunlock(&t->t_lock);
-}
-
-static void
 wg_timers_run_send_initiation(struct wg_timers *t, int is_retry)
 {
 	struct wg_peer	 *peer = __containerof(t, struct wg_peer, p_timers);
 	if (!is_retry)
 		t->t_handshake_retries = 0;
-	if (wg_timers_expired_handshake_last_sent(t) == ETIMEDOUT)
+	if (noise_remote_initiation_expired(peer->p_remote) == ETIMEDOUT)
 		GROUPTASK_ENQUEUE(&peer->p_send_initiation);
 }
 
@@ -1489,7 +1319,7 @@ wg_timers_run_zero_key_material(struct wg_timers *t)
 	DPRINTF(peer->p_sc, "Zeroing out all keys for peer %llu, since we "
 	    "haven't received a new one in %d seconds\n",
 		(unsigned long long)peer->p_id, REJECT_AFTER_TIME * 3);
-	GROUPTASK_ENQUEUE(&peer->p_clear_secrets);
+	noise_remote_keypairs_clear(peer->p_remote);
 }
 
 static void
@@ -1520,15 +1350,14 @@ wg_send_initiation(struct wg_peer *peer)
 	struct wg_pkt_initiation pkt;
 	struct epoch_tracker et;
 
-	if (wg_timers_check_handshake_last_sent(&peer->p_timers) != ETIMEDOUT)
-		return;
-	DPRINTF(peer->p_sc, "Sending handshake initiation to peer %llu\n",
-		(unsigned long long)peer->p_id);
-
 	NET_EPOCH_ENTER(et);
-	if (noise_create_initiation(&peer->p_remote, &pkt.s_idx, pkt.ue,
+	if (noise_create_initiation(peer->p_remote, &pkt.s_idx, pkt.ue,
 	    pkt.es, pkt.ets) != 0)
 		goto out;
+
+	DPRINTF(peer->p_sc, "Sending handshake initiation to peer %llu\n",
+		(unsigned long long)peer->p_id);
+
 	pkt.t = WG_PKT_INITIATION;
 	cookie_maker_mac(&peer->p_cookie, &pkt.m, &pkt,
 	    sizeof(pkt)-sizeof(pkt.m));
@@ -1545,21 +1374,17 @@ wg_send_response(struct wg_peer *peer)
 	struct epoch_tracker et;
 
 	NET_EPOCH_ENTER(et);
+	if (noise_create_response(peer->p_remote, &pkt.s_idx, &pkt.r_idx,
+	    pkt.ue, pkt.en) != 0)
+		goto out;
 
 	DPRINTF(peer->p_sc, "Sending handshake response to peer %llu\n",
 	    (unsigned long long)peer->p_id);
 
-	if (noise_create_response(&peer->p_remote, &pkt.s_idx, &pkt.r_idx,
-	    pkt.ue, pkt.en) != 0)
-		goto out;
-	if (noise_remote_begin_session(&peer->p_remote) != 0)
-		goto out;
-
 	wg_timers_event_session_derived(&peer->p_timers);
 	pkt.t = WG_PKT_RESPONSE;
 	cookie_maker_mac(&peer->p_cookie, &pkt.m, &pkt,
 	     sizeof(pkt)-sizeof(pkt.m));
-	wg_timers_event_handshake_responded(&peer->p_timers);
 	wg_peer_send_buf(peer, (uint8_t*)&pkt, sizeof(pkt));
 out:
 	NET_EPOCH_EXIT(et);
@@ -1600,7 +1425,6 @@ wg_send_keepalive(struct wg_peer *peer)
 		m_freem(m);
 		return;
 	}
-	t->t_peer = peer;
 	t->t_mbuf = NULL;
 	t->t_done = 0;
 	t->t_mtu = 0; /* MTU == 0 OK for keepalive */
@@ -1673,10 +1497,6 @@ wg_handshake(struct wg_softc *sc, struct mbuf *m)
 
 	res = wg_cookie_validate_packet(&sc->sc_cookie, m, underload);
 
-	if (res && res != EAGAIN) {
-		printf("validate_packet got %d\n", res);
-		goto free;
-	}
 	if (res == EINVAL) {
 		DPRINTF(sc, "Invalid initiation MAC\n");
 		goto free;
@@ -1699,13 +1519,13 @@ wg_handshake(struct wg_softc *sc, struct mbuf *m)
 			wg_send_cookie(sc, &init->m, init->s_idx, m);
 			goto free;
 		}
-		if (noise_consume_initiation(&sc->sc_local, &remote,
+		if (noise_consume_initiation(sc->sc_local, &remote,
 		    init->s_idx, init->ue, init->es, init->ets) != 0) {
 			DPRINTF(sc, "Invalid handshake initiation");
 			goto free;
 		}
 
-		peer = __containerof(remote, struct wg_peer, p_remote);
+		peer = noise_remote_arg(remote);
 		DPRINTF(sc, "Receiving handshake initiation from peer %llu\n",
 		    (unsigned long long)peer->p_id);
 		counter_u64_add(peer->p_rx_bytes, sizeof(*init));
@@ -1721,46 +1541,43 @@ wg_handshake(struct wg_softc *sc, struct mbuf *m)
 			wg_send_cookie(sc, &resp->m, resp->s_idx, m);
 			goto free;
 		}
-
-		if ((remote = wg_index_get(sc, resp->r_idx)) == NULL) {
-			DPRINTF(sc, "Unknown handshake response\n");
-			goto free;
-		}
-		peer = __containerof(remote, struct wg_peer, p_remote);
-		if (noise_consume_response(remote, resp->s_idx, resp->r_idx,
-		    resp->ue, resp->en) != 0) {
+		if (noise_consume_response(sc->sc_local, &remote, resp->s_idx,
+		    resp->r_idx, resp->ue, resp->en) != 0) {
 			DPRINTF(sc, "Invalid handshake response\n");
 			goto free;
 		}
 
+		peer = noise_remote_arg(remote);
+
 		DPRINTF(sc, "Receiving handshake response from peer %llu\n",
 				(unsigned long long)peer->p_id);
+
+		wg_timers_event_session_derived(&peer->p_timers);
+		wg_timers_event_handshake_complete(&peer->p_timers);
+
 		counter_u64_add(peer->p_rx_bytes, sizeof(*resp));
 		if_inc_counter(sc->sc_ifp, IFCOUNTER_IPACKETS, 1);
 		if_inc_counter(sc->sc_ifp, IFCOUNTER_IBYTES, sizeof(*resp));
 		wg_peer_set_endpoint_from_tag(peer, t);
-		if (noise_remote_begin_session(&peer->p_remote) == 0) {
-			wg_timers_event_session_derived(&peer->p_timers);
-			wg_timers_event_handshake_complete(&peer->p_timers);
-		}
 		break;
 	case WG_PKT_COOKIE:
 		cook = mtod(m, struct wg_pkt_cookie *);
 
-		if ((remote = wg_index_get(sc, cook->r_idx)) == NULL) {
+		if ((remote = noise_remote_index_lookup(sc->sc_local,
+		    cook->r_idx)) == NULL) {
 			DPRINTF(sc, "Unknown cookie index\n");
 			goto free;
 		}
 
-		peer = __containerof(remote, struct wg_peer, p_remote);
+		peer = noise_remote_arg(remote);
 
 		if (cookie_maker_consume_payload(&peer->p_cookie,
-		    cook->nonce, cook->ec) != 0) {
+		    cook->nonce, cook->ec) == 0) {
+			DPRINTF(sc, "Receiving cookie response\n");
+		} else {
 			DPRINTF(sc, "Could not decrypt cookie response\n");
-			goto free;
 		}
-
-		DPRINTF(sc, "Receiving cookie response\n");
+		noise_remote_put(remote);
 		goto free;
 	default:
 		goto free;
@@ -1788,15 +1605,19 @@ wg_encap(struct wg_softc *sc, struct mbuf *m)
 {
 	struct wg_pkt_data *data;
 	size_t padding_len, plaintext_len, out_len;
+	struct noise_remote *remote;
 	struct mbuf *mc;
 	struct wg_peer *peer;
 	struct wg_tag *t;
 	uint64_t nonce;
-	int res, allocation_order;
+	int allocation_order;
 
 	NET_EPOCH_ASSERT();
 	t = wg_tag_get(m);
-	peer = t->t_peer;
+	remote = noise_keypair_remote(t->t_keypair);
+	peer = noise_remote_arg(remote);
+	/* We can put the remote as we still hold a keypair ref */
+	noise_remote_put(remote);
 
 	plaintext_len = MIN(WG_PKT_WITH_PADDING(m->m_pkthdr.len), t->t_mtu);
 	padding_len = plaintext_len - m->m_pkthdr.len;
@@ -1822,23 +1643,11 @@ wg_encap(struct wg_softc *sc, struct mbuf *m)
 
 	data->t = WG_PKT_DATA;
 
-	res = noise_remote_encrypt(&peer->p_remote, &data->r_idx, &nonce,
+	noise_keypair_encrypt(t->t_keypair, &data->r_idx, t->t_nonce,
 	    data->buf, plaintext_len);
-	nonce = htole64(nonce); /* Wire format is little endian. */
-	memcpy(data->nonce, &nonce, sizeof(data->nonce));
 
-	if (__predict_false(res)) {
-		if (res == EINVAL) {
-			wg_timers_event_want_initiation(&peer->p_timers);
-			m_freem(mc);
-			goto error;
-		} else if (res == ESTALE) {
-			wg_timers_event_want_initiation(&peer->p_timers);
-		} else {
-			m_freem(mc);
-			goto error;
-		}
-	}
+	nonce = htole64(t->t_nonce); /* Wire format is little endian. */
+	memcpy(data->nonce, &nonce, sizeof(data->nonce));
 
 	/* A packet with length 0 is a keepalive packet */
 	if (m->m_pkthdr.len == 0)
@@ -1863,6 +1672,7 @@ wg_decap(struct wg_softc *sc, struct mbuf *m)
 {
 	struct wg_pkt_data *data;
 	struct wg_peer *peer, *routed_peer;
+	struct noise_remote *remote;
 	struct wg_tag *t;
 	size_t plaintext_len;
 	uint8_t version;
@@ -1874,21 +1684,22 @@ wg_decap(struct wg_softc *sc, struct mbuf *m)
 	plaintext_len = m->m_pkthdr.len - sizeof(struct wg_pkt_data);
 
 	t = wg_tag_get(m);
-	peer = t->t_peer;
+	remote = noise_keypair_remote(t->t_keypair);
+	peer = noise_remote_arg(remote);
+	/* We can put the remote as we still hold a keypair ref */
+	noise_remote_put(remote);
 
 	memcpy(&nonce, data->nonce, sizeof(nonce));
-	nonce = le64toh(nonce); /* Wire format is little endian. */
+	t->t_nonce = le64toh(nonce); /* Wire format is little endian. */
 
-	res = noise_remote_decrypt(&peer->p_remote, data->r_idx, nonce,
-	    data->buf, plaintext_len);
+	res = noise_keypair_decrypt(t->t_keypair, t->t_nonce, data->buf,
+	    plaintext_len);
 
 	if (__predict_false(res)) {
 		if (res == EINVAL) {
 			goto error;
 		} else if (res == ECONNRESET) {
 			wg_timers_event_handshake_complete(&peer->p_timers);
-		} else if (res == ESTALE) {
-			wg_timers_event_want_initiation(&peer->p_timers);
 		} else  {
 			panic("unexpected response: %d\n", res);
 		}
@@ -1994,6 +1805,7 @@ wg_deliver_out(struct wg_peer *peer)
 			continue;
 		}
 		len = t->t_mbuf->m_pkthdr.len;
+		noise_keypair_put(t->t_keypair);
 		ret = wg_send(peer->p_sc, &endpoint, t->t_mbuf);
 
 		if (ret == 0) {
@@ -2010,6 +1822,9 @@ wg_deliver_out(struct wg_peer *peer)
 			wg_peer_get_endpoint(peer, &endpoint);
 		}
 		m_freem(m);
+
+		if (noise_keep_key_fresh_send(peer->p_remote) == 0)
+			wg_timers_event_want_initiation(&peer->p_timers);
 	}
 
 	NET_EPOCH_EXIT(et);
@@ -2039,6 +1854,14 @@ wg_deliver_in(struct wg_peer *peer)
 		}
 		MPASS(m == t->t_mbuf);
 
+		if (noise_keypair_nonce_check(t->t_keypair, t->t_nonce) != 0) {
+			if_inc_counter(ifp, IFCOUNTER_IERRORS, 1);
+			noise_keypair_put(t->t_keypair);
+			m_freem(m);
+			continue;
+		}
+		noise_keypair_put(t->t_keypair);
+
 		wg_timers_event_any_authenticated_packet_received(
 		    &peer->p_timers);
 		wg_timers_event_any_authenticated_packet_traversal(
@@ -2048,6 +1871,9 @@ wg_deliver_in(struct wg_peer *peer)
 		if_inc_counter(sc->sc_ifp, IFCOUNTER_IPACKETS, 1);
 		if_inc_counter(sc->sc_ifp, IFCOUNTER_IBYTES, m->m_pkthdr.len + sizeof(struct wg_pkt_data) + NOISE_AUTHTAG_LEN);
 
+		if (noise_keep_key_fresh_recv(peer->p_remote) == 0)
+			wg_timers_event_want_initiation(&peer->p_timers);
+
 		if (m->m_pkthdr.len == 0) {
 			m_freem(m);
 			continue;
@@ -2129,10 +1955,11 @@ wg_queue_out(struct wg_peer *peer)
 	struct buf_ring *parallel = peer->p_sc->sc_encap_ring;
 	struct wg_queue		*serial = &peer->p_encap_queue;
 	struct wg_tag		*t;
+	struct noise_keypair	*keypair;
 	struct mbufq		 staged;
 	struct mbuf		*m;
 
-	if (noise_remote_ready(&peer->p_remote) != 0) {
+	if ((keypair = noise_keypair_current(peer->p_remote)) == NULL) {
 		if (wg_queue_len(&peer->p_stage_queue))
 			wg_timers_event_want_initiation(&peer->p_timers);
 		return;
@@ -2153,10 +1980,18 @@ wg_queue_out(struct wg_peer *peer)
 			m_freem(m);
 			continue;
 		}
-		t->t_peer = peer;
+		if (noise_keypair_nonce_next(keypair, &t->t_nonce) != 0) {
+			/* TODO if we get here, it means we are about to
+			 * overflow this keypair sending nonce. We should place
+			 * this back on the staged queue. */
+			m_freem(m);
+			continue;
+		}
+		t->t_keypair = noise_keypair_ref(keypair);
 		mtx_lock(&serial->q_mtx);
 		if (mbufq_enqueue(&serial->q, m) != 0) {
 			m_freem(m);
+			noise_keypair_put(keypair);
 			if_inc_counter(peer->p_sc->sc_ifp, IFCOUNTER_OQDROPS, 1);
 		} else {
 			m->m_flags |= M_ENQUEUED;
@@ -2215,90 +2050,6 @@ wg_queue_purge(struct wg_queue *q)
 	mtx_unlock(&q->q_mtx);
 }
 
-/* TODO Indexes */
-static struct noise_remote *
-wg_remote_get(struct wg_softc *sc, uint8_t public[NOISE_PUBLIC_KEY_LEN])
-{
-	struct wg_peer *peer;
-
-	if ((peer = wg_peer_lookup(sc, public)) == NULL)
-		return (NULL);
-	return (&peer->p_remote);
-}
-
-static uint32_t
-wg_index_set(struct wg_softc *sc, struct noise_remote *remote)
-{
-	struct wg_index *index, *iter;
-	struct wg_peer	*peer;
-	uint32_t	 key;
-
-	/* We can modify this without a lock as wg_index_set, wg_index_drop are
-	 * guaranteed to be serialised (per remote). */
-	peer = __containerof(remote, struct wg_peer, p_remote);
-	index = SLIST_FIRST(&peer->p_unused_index);
-	MPASS(index != NULL);
-	SLIST_REMOVE_HEAD(&peer->p_unused_index, i_unused_entry);
-
-	index->i_value = remote;
-
-	rw_wlock(&sc->sc_index_lock);
-assign_id:
-	key = index->i_key = arc4random();
-	key &= sc->sc_index_mask;
-	LIST_FOREACH(iter, &sc->sc_index[key], i_entry)
-		if (iter->i_key == index->i_key)
-			goto assign_id;
-
-	LIST_INSERT_HEAD(&sc->sc_index[key], index, i_entry);
-
-	rw_wunlock(&sc->sc_index_lock);
-
-	/* Likewise, no need to lock for index here. */
-	return index->i_key;
-}
-
-static struct noise_remote *
-wg_index_get(struct wg_softc *sc, uint32_t key0)
-{
-	struct wg_index		*iter;
-	struct noise_remote	*remote = NULL;
-	uint32_t		 key = key0 & sc->sc_index_mask;
-
-	rw_enter_read(&sc->sc_index_lock);
-	LIST_FOREACH(iter, &sc->sc_index[key], i_entry)
-		if (iter->i_key == key0) {
-			remote = iter->i_value;
-			break;
-		}
-	rw_exit_read(&sc->sc_index_lock);
-	return remote;
-}
-
-static void
-wg_index_drop(struct wg_softc *sc, uint32_t key0)
-{
-	struct wg_index	*iter;
-	struct wg_peer	*peer = NULL;
-	uint32_t	 key = key0 & sc->sc_index_mask;
-
-	rw_enter_write(&sc->sc_index_lock);
-	LIST_FOREACH(iter, &sc->sc_index[key], i_entry)
-		if (iter->i_key == key0) {
-			LIST_REMOVE(iter, i_entry);
-			break;
-		}
-	rw_exit_write(&sc->sc_index_lock);
-
-	if (iter == NULL)
-		return;
-
-	/* We expect a peer */
-	peer = __containerof(iter->i_value, struct wg_peer, p_remote);
-	MPASS(peer != NULL);
-	SLIST_INSERT_HEAD(&peer->p_unused_index, iter, i_unused_entry);
-}
-
 static int
 wg_update_endpoint_addrs(struct wg_endpoint *e, const struct sockaddr *srcsa,
     struct ifnet *rcvif)
@@ -2339,6 +2090,7 @@ wg_input(struct mbuf *m0, int offset, struct inpcb *inpcb,
 	struct wg_softc *sc = _sc;
 	struct mbuf *m;
 	int pktlen, pkttype;
+	struct noise_keypair *keypair;
 	struct noise_remote *remote;
 	struct wg_tag *t;
 	void *data;
@@ -2406,21 +2158,23 @@ wg_input(struct mbuf *m0, int offset, struct inpcb *inpcb,
 	} else if (pktlen >= sizeof(struct wg_pkt_data) + NOISE_AUTHTAG_LEN
 	    && pkttype == WG_PKT_DATA) {
 		pkt_data = data;
-		remote = wg_index_get(sc, pkt_data->r_idx);
-		if (remote == NULL) {
+		keypair = noise_keypair_lookup(sc->sc_local, pkt_data->r_idx);
+		if (keypair == NULL) {
 			if_inc_counter(sc->sc_ifp, IFCOUNTER_IERRORS, 1);
 			m_freem(m);
 		} else if (buf_ring_count(sc->sc_decap_ring) > MAX_QUEUED_PKT) {
 			if_inc_counter(sc->sc_ifp, IFCOUNTER_IQDROPS, 1);
+			noise_keypair_put(keypair);
 			m_freem(m);
 		} else {
-			t->t_peer = __containerof(remote, struct wg_peer,
-			    p_remote);
+			t->t_keypair = keypair;
 			t->t_mbuf = NULL;
 			t->t_done = 0;
 
-			wg_queue_in(t->t_peer, m);
+			remote = noise_keypair_remote(keypair);
+			wg_queue_in(noise_remote_arg(remote), m);
 			wg_decrypt_dispatch(sc);
+			noise_remote_put(remote);
 		}
 	} else {
 free:
@@ -2465,7 +2219,6 @@ wg_transmit(struct ifnet *ifp, struct mbuf *m)
 		rc = EHOSTUNREACH;
 		goto err;
 	}
-	t->t_peer = peer;
 	t->t_mbuf = NULL;
 	t->t_done = 0;
 	t->t_mtu = ifp->if_mtu;
@@ -2494,10 +2247,11 @@ static int
 wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 {
 	uint8_t			 public[WG_KEY_SIZE];
-	const void *pub_key;
+	const void *pub_key, *preshared_key = NULL;
 	const struct sockaddr *endpoint;
 	int err;
 	size_t size;
+	struct noise_remote *remote;
 	struct wg_peer *peer = NULL;
 	bool need_insert = false;
 
@@ -2510,17 +2264,16 @@ wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 	if (size != WG_KEY_SIZE) {
 		return (EINVAL);
 	}
-	if (noise_local_keys(&sc->sc_local, public, NULL) == 0 &&
+	if (noise_local_keys(sc->sc_local, public, NULL) == 0 &&
 	    bcmp(public, pub_key, WG_KEY_SIZE) == 0) {
 		return (0); // Silently ignored; not actually a failure.
 	}
-	peer = wg_peer_lookup(sc, pub_key);
+	if ((remote = noise_remote_lookup(sc->sc_local, pub_key)) != NULL)
+		peer = noise_remote_arg(remote);
 	if (nvlist_exists_bool(nvl, "remove") &&
 		nvlist_get_bool(nvl, "remove")) {
-		if (peer != NULL) {
-			wg_hashtable_peer_remove(&sc->sc_hashtable, peer);
+		if (peer != NULL)
 			wg_peer_destroy(peer);
-		}
 		return (0);
 	}
 	if (nvlist_exists_bool(nvl, "replace-allowedips") &&
@@ -2530,14 +2283,9 @@ wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 		wg_aip_delete(&peer->p_sc->sc_aips, peer);
 	}
 	if (peer == NULL) {
-		if (sc->sc_peer_count >= MAX_PEERS_PER_IFACE)
-			return (E2BIG);
-		sc->sc_peer_count++;
-
 		need_insert = true;
 		peer = wg_peer_alloc(sc);
 		MPASS(peer != NULL);
-		noise_remote_init(&peer->p_remote, pub_key, &sc->sc_local);
 		cookie_maker_init(&peer->p_cookie, pub_key);
 	}
 	if (nvlist_exists_binary(nvl, "endpoint")) {
@@ -2549,14 +2297,13 @@ wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 		memcpy(&peer->p_endpoint.e_remote, endpoint, size);
 	}
 	if (nvlist_exists_binary(nvl, "preshared-key")) {
-		const void *key;
-
-		key = nvlist_get_binary(nvl, "preshared-key", &size);
+		preshared_key = nvlist_get_binary(nvl, "preshared-key", &size);
 		if (size != WG_KEY_SIZE) {
 			err = EINVAL;
 			goto out;
 		}
-		noise_remote_set_psk(&peer->p_remote, key);
+		if (!need_insert)
+			noise_remote_set_psk(peer->p_remote, preshared_key);
 	}
 	if (nvlist_exists_number(nvl, "persistent-keepalive-interval")) {
 		uint64_t pki = nvlist_get_number(nvl, "persistent-keepalive-interval");
@@ -2606,7 +2353,11 @@ wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 		}
 	}
 	if (need_insert) {
-		wg_hashtable_peer_insert(&sc->sc_hashtable, peer);
+		if ((peer->p_remote = noise_remote_alloc(sc->sc_local, peer,
+		    pub_key, preshared_key)) == NULL)
+			goto out;
+		TAILQ_INSERT_TAIL(&sc->sc_peers, peer, p_entry);
+		sc->sc_peers_num++;
 		if (sc->sc_ifp->if_link_state == LINK_STATE_UP)
 			wg_timers_enable(&peer->p_timers);
 	}
@@ -2672,19 +2423,18 @@ wgc_set(struct wg_softc *sc, struct wg_data_io *wgd)
 			goto out;
 		}
 
-		if (noise_local_keys(&sc->sc_local, NULL, private) != 0 ||
+		if (noise_local_keys(sc->sc_local, NULL, private) != 0 ||
 		    timingsafe_bcmp(private, key, WG_KEY_SIZE) != 0) {
-			struct noise_local *local;
 			struct wg_peer *peer;
-			struct wg_hashtable *ht = &sc->sc_hashtable;
-			bool has_identity;
 
 			if (curve25519_generate_public(public, key)) {
 				/* Peer conflict: remove conflicting peer. */
-				if ((peer = wg_peer_lookup(sc, public)) !=
-				    NULL) {
-					wg_hashtable_peer_remove(ht, peer);
+				struct noise_remote *remote;
+				if ((remote = noise_remote_lookup(sc->sc_local,
+				    public)) != NULL) {
+					peer = noise_remote_arg(remote);
 					wg_peer_destroy(peer);
+					noise_remote_put(remote);
 				}
 			}
 
@@ -2692,21 +2442,12 @@ wgc_set(struct wg_softc *sc, struct wg_data_io *wgd)
 			 * Set the private key and invalidate all existing
 			 * handshakes.
 			 */
-			local = &sc->sc_local;
-			noise_local_lock_identity(local);
 			/* Note: we might be removing the private key. */
-			has_identity = noise_local_set_private(local, key) == 0;
-			mtx_lock(&ht->h_mtx);
-			CK_LIST_FOREACH(peer, &ht->h_peers_list, p_entry) {
-				noise_remote_precompute(&peer->p_remote);
-				wg_timers_event_reset_handshake_last_sent(
-				    &peer->p_timers);
-				noise_remote_expire_current(&peer->p_remote);
-			}
-			mtx_unlock(&ht->h_mtx);
-			cookie_checker_update(&sc->sc_cookie,
-			    has_identity ? public : NULL);
-			noise_local_unlock_identity(local);
+			noise_local_private(sc->sc_local, key);
+			if (noise_local_keys(sc->sc_local, NULL, NULL) == 0)
+				cookie_checker_update(&sc->sc_cookie, public);
+			else
+				cookie_checker_update(&sc->sc_cookie, NULL);
 		}
 	}
 	if (nvlist_exists_number(nvl, "user-cookie")) {
@@ -2742,7 +2483,9 @@ out:
 static int
 wgc_get(struct wg_softc *sc, struct wg_data_io *wgd)
 {
-	uint8_t public_key[WG_KEY_SIZE] = { 0 }, preshared_key[NOISE_SYMMETRIC_KEY_LEN] = { 0 };
+	uint8_t public_key[WG_KEY_SIZE] = { 0 };
+	uint8_t private_key[WG_KEY_SIZE] = { 0 };
+	uint8_t preshared_key[NOISE_SYMMETRIC_KEY_LEN] = { 0 };
 	nvlist_t *nvl, *nvl_peer, *nvl_aip, **nvl_peers, **nvl_aips;
 	size_t size, peer_count, aip_count, i, j;
 	struct wg_timespec64 ts64;
@@ -2761,16 +2504,17 @@ wgc_get(struct wg_softc *sc, struct wg_data_io *wgd)
 		nvlist_add_number(nvl, "listen-port", sc->sc_socket.so_port);
 	if (sc->sc_socket.so_user_cookie != 0)
 		nvlist_add_number(nvl, "user-cookie", sc->sc_socket.so_user_cookie);
-	if (sc->sc_local.l_has_identity) {
-		nvlist_add_binary(nvl, "public-key", sc->sc_local.l_public, WG_KEY_SIZE);
+	if (noise_local_keys(sc->sc_local, public_key, private_key) == 0) {
+		nvlist_add_binary(nvl, "public-key", public_key, WG_KEY_SIZE);
 		if (wgc_privileged(sc))
-			nvlist_add_binary(nvl, "private-key", sc->sc_local.l_private, WG_KEY_SIZE);
+			nvlist_add_binary(nvl, "private-key", private_key, WG_KEY_SIZE);
+		explicit_bzero(private_key, sizeof(private_key));
 	}
-	peer_count = sc->sc_hashtable.h_num_peers;
+	peer_count = sc->sc_peers_num;
 	if (peer_count) {
 		nvl_peers = mallocarray(peer_count, sizeof(void *), M_NVLIST, M_WAITOK | M_ZERO);
 		i = 0;
-		CK_LIST_FOREACH(peer, &sc->sc_hashtable.h_peers_list, p_entry) {
+		TAILQ_FOREACH(peer, &sc->sc_peers, p_entry) {
 			if (i >= peer_count)
 				panic("peers changed from under us");
 
@@ -2780,7 +2524,7 @@ wgc_get(struct wg_softc *sc, struct wg_data_io *wgd)
 				goto err_peer;
 			}
 
-			(void)noise_remote_keys(&peer->p_remote, public_key, preshared_key);
+			(void)noise_remote_keys(peer->p_remote, public_key, preshared_key);
 			nvlist_add_binary(nvl_peer, "public-key", public_key, sizeof(public_key));
 			if (wgc_privileged(sc))
 				nvlist_add_binary(nvl_peer, "preshared-key", preshared_key, sizeof(preshared_key));
@@ -2941,7 +2685,6 @@ out:
 static int
 wg_up(struct wg_softc *sc)
 {
-	struct wg_hashtable *ht = &sc->sc_hashtable;
 	struct ifnet *ifp = sc->sc_ifp;
 	struct wg_peer *peer;
 	int rc = EBUSY;
@@ -2959,13 +2702,10 @@ wg_up(struct wg_softc *sc)
 
 	rc = wg_socket_init(sc, sc->sc_socket.so_port);
 	if (rc == 0) {
-		mtx_lock(&ht->h_mtx);
-		CK_LIST_FOREACH(peer, &ht->h_peers_list, p_entry) {
+		TAILQ_FOREACH(peer, &sc->sc_peers, p_entry) {
 			wg_timers_enable(&peer->p_timers);
 			wg_queue_out(peer);
 		}
-		mtx_unlock(&ht->h_mtx);
-
 		if_link_state_change(sc->sc_ifp, LINK_STATE_UP);
 	} else {
 		ifp->if_drv_flags &= ~IFF_DRV_RUNNING;
@@ -2979,7 +2719,6 @@ out:
 static void
 wg_down(struct wg_softc *sc)
 {
-	struct wg_hashtable *ht = &sc->sc_hashtable;
 	struct ifnet *ifp = sc->sc_ifp;
 	struct wg_peer *peer;
 
@@ -2990,21 +2729,17 @@ wg_down(struct wg_softc *sc)
 	}
 	ifp->if_drv_flags &= ~IFF_DRV_RUNNING;
 
-	mtx_lock(&ht->h_mtx);
-	CK_LIST_FOREACH(peer, &ht->h_peers_list, p_entry) {
+	TAILQ_FOREACH(peer, &sc->sc_peers, p_entry) {
 		wg_queue_purge(&peer->p_stage_queue);
 		wg_timers_disable(&peer->p_timers);
 	}
-	mtx_unlock(&ht->h_mtx);
 
 	mbufq_drain(&sc->sc_handshake_queue);
 
-	mtx_lock(&ht->h_mtx);
-	CK_LIST_FOREACH(peer, &ht->h_peers_list, p_entry) {
-		noise_remote_clear(&peer->p_remote);
-		wg_timers_event_reset_handshake_last_sent(&peer->p_timers);
+	TAILQ_FOREACH(peer, &sc->sc_peers, p_entry) {
+		noise_remote_handshake_clear(peer->p_remote);
+		noise_remote_keypairs_clear(peer->p_remote);
 	}
-	mtx_unlock(&ht->h_mtx);
 
 	if_link_state_change(sc->sc_ifp, LINK_STATE_DOWN);
 	wg_socket_uninit(sc);
@@ -3045,7 +2780,6 @@ wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 {
 	struct wg_softc *sc;
 	struct ifnet *ifp;
-	struct noise_upcall noise_upcall;
 
 	sc = malloc(sizeof(*sc), M_WG, M_WAITOK | M_ZERO);
 	sc->sc_ucred = crhold(curthread->td_ucred);
@@ -3054,14 +2788,15 @@ wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 	ifp->if_softc = sc;
 	if_initname(ifp, wgname, unit);
 
-	noise_upcall.u_arg = sc;
-	noise_upcall.u_remote_get =
-		(struct noise_remote *(*)(void *, uint8_t *))wg_remote_get;
-	noise_upcall.u_index_set =
-		(uint32_t (*)(void *, struct noise_remote *))wg_index_set;
-	noise_upcall.u_index_drop =
-		(void (*)(void *, uint32_t))wg_index_drop;
-	noise_local_init(&sc->sc_local, &noise_upcall);
+	TAILQ_INIT(&sc->sc_peers);
+	sc->sc_peers_num = 0;
+
+	if ((sc->sc_local = noise_local_alloc(sc)) == NULL) {
+		free(sc, M_WG);
+		return ENOMEM;
+	}
+
+	/* TODO check checker_init return value */
 	cookie_checker_init(&sc->sc_cookie, ratelimit_zone);
 
 	sc->sc_socket.so_port = 0;
@@ -3071,8 +2806,6 @@ wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 
 	mbufq_init(&sc->sc_handshake_queue, MAX_QUEUED_HANDSHAKES);
 	sx_init(&sc->sc_lock, "wg softc lock");
-	rw_init(&sc->sc_index_lock, "wg index lock");
-	sc->sc_peer_count = 0;
 	sc->sc_encap_ring = buf_ring_alloc(MAX_QUEUED_PKT, M_WG, M_WAITOK, NULL);
 	sc->sc_decap_ring = buf_ring_alloc(MAX_QUEUED_PKT, M_WG, M_WAITOK, NULL);
 	GROUPTASK_INIT(&sc->sc_handshake, 0,
@@ -3080,8 +2813,6 @@ wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 	taskqgroup_attach(qgroup_wg_tqg, &sc->sc_handshake, sc, NULL, NULL, "wg tx initiation");
 	crypto_taskq_setup(sc);
 
-	wg_hashtable_init(&sc->sc_hashtable);
-	sc->sc_index = hashinit(HASHTABLE_INDEX_SIZE, M_DEVBUF, &sc->sc_index_mask);
 	wg_aip_init(&sc->sc_aips);
 
 	if_setmtu(ifp, ETHERMTU - 80);
@@ -3108,6 +2839,15 @@ wg_clone_create(struct if_clone *ifc, int unit, caddr_t params)
 }
 
 static void
+wg_clone_deferred_free(struct noise_local *l)
+{
+	struct wg_softc *sc = noise_local_arg(l);
+
+	free(sc, M_WG);
+	atomic_add_int(&clone_count, -1);
+}
+
+static void
 wg_clone_destroy(struct ifnet *ifp)
 {
 	struct wg_softc *sc = ifp->if_softc;
@@ -3144,24 +2884,19 @@ wg_clone_destroy(struct ifnet *ifp)
 	epoch_drain_callbacks(net_epoch_preempt);
 	sx_xunlock(&sc->sc_lock);
 	sx_destroy(&sc->sc_lock);
-	rw_destroy(&sc->sc_index_lock);
 	taskqgroup_detach(qgroup_wg_tqg, &sc->sc_handshake);
 	crypto_taskq_destroy(sc);
 	buf_ring_free(sc->sc_encap_ring, M_WG);
 	buf_ring_free(sc->sc_decap_ring, M_WG);
 
 	wg_aip_destroy(&sc->sc_aips);
-	wg_hashtable_destroy(&sc->sc_hashtable);
 
 	if (cred != NULL)
 		crfree(cred);
 	if_detach(sc->sc_ifp);
 	if_free(sc->sc_ifp);
-	/* Ensure any local/private keys are cleaned up */
-	explicit_bzero(sc, sizeof(*sc));
-	free(sc, M_WG);
 
-	atomic_add_int(&clone_count, -1);
+	noise_local_free(sc->sc_local, wg_clone_deferred_free);
 }
 
 static void
diff --git a/src/support.h b/src/support.h
index 5256e62..99cd518 100644
--- a/src/support.h
+++ b/src/support.h
@@ -82,4 +82,25 @@ sogetsockaddr(struct socket *so, struct sockaddr **nam)
 	return (error);
 }
 
+/* These are defined in sys/compat/linuxkpi/common/include/linux/compiler.h,
+ * however I don't really want to include all that. */
+#define barrier()                       __asm__ __volatile__("": : :"memory")
+
+#define ACCESS_ONCE(x)                  (*(volatile __typeof(x) *)&(x))
+
+#define WRITE_ONCE(x,v) do {            \
+        barrier();                      \
+        ACCESS_ONCE(x) = (v);           \
+        barrier();                      \
+} while (0)
+
+#define READ_ONCE(x) ({                 \
+        __typeof(x) __var = ({          \
+                barrier();              \
+                ACCESS_ONCE(x);         \
+        });                             \
+        barrier();                      \
+        __var;                          \
+})
+
 #endif
diff --git a/src/wg_noise.c b/src/wg_noise.c
index b5bd5c5..5d92750 100644
--- a/src/wg_noise.c
+++ b/src/wg_noise.c
@@ -9,147 +9,474 @@
 #include <sys/lock.h>
 #include <sys/rwlock.h>
 #include <sys/systm.h>
+#include <sys/malloc.h>
+#include <sys/refcount.h>
+#include <sys/epoch.h>
+#include <sys/ck.h>
 
-#include "support.h"
+#include "crypto.h"
 #include "wg_noise.h"
+#include "support.h"
 
-/* Private functions */
-static struct noise_keypair *
-		noise_remote_keypair_allocate(struct noise_remote *);
-static void
-		noise_remote_keypair_free(struct noise_remote *,
-			struct noise_keypair *);
-static uint32_t	noise_remote_handshake_index_get(struct noise_remote *);
-static void	noise_remote_handshake_index_drop(struct noise_remote *);
-
-static uint64_t	noise_counter_send(struct noise_counter *);
-static int	noise_counter_recv(struct noise_counter *, uint64_t);
+/* 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 {
+	CK_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;
+	u_int				 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 epoch_context		 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;
+
+	CK_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];
+
+	u_int				 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 epoch_context		 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];
+
+	u_int				 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;
+	CK_LIST_HEAD(,noise_remote)	 l_remote_hash[HT_REMOTE_SIZE];
+
+	struct rwlock			 l_index_lock;
+	CK_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 key[NOISE_SYMMETRIC_KEY_LEN],
-			const uint8_t ss[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]);
-
+		    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]);
+		    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 src[NOISE_PUBLIC_KEY_LEN]);
-
+		    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);
 
-/* Set/Get noise parameters */
-void
-noise_local_init(struct noise_local *l, struct noise_upcall *upcall)
+/* I can't find where FreeBSD defines such behaviours, so that is temporarily here. */
+#define epoch_ptr_read(p) ck_pr_load_ptr(p)
+#define epoch_ptr_write(p, v) ck_pr_store_ptr(p, v)
+/* Back to regular programming... */
+
+MALLOC_DEFINE(M_NOISE, "NOISE", "wgnoise");
+
+/* Local configuration */
+struct noise_local *
+noise_local_alloc(void *arg)
 {
-	bzero(l, sizeof(*l));
-	rw_init(&l->l_identity_lock, "noise_local_identity");
-	l->l_upcall = *upcall;
+	struct noise_local *l;
+	size_t i;
+
+	if ((l = malloc(sizeof(*l), M_NOISE, 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);
+
+	refcount_init(&l->l_refcnt, 1);
+	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++)
+		CK_LIST_INIT(&l->l_remote_hash[i]);
+
+	rw_init(&l->l_index_lock, "noise_index");
+	for (i = 0; i < HT_INDEX_SIZE; i++)
+		CK_LIST_INIT(&l->l_index_hash[i]);
+
+	return (l);
+}
+
+struct noise_local *
+noise_local_ref(struct noise_local *l)
+{
+	refcount_acquire(&l->l_refcnt);
+	return (l);
 }
 
 void
-noise_local_lock_identity(struct noise_local *l)
+noise_local_put(struct noise_local *l)
 {
-	rw_enter_write(&l->l_identity_lock);
+	if (refcount_release(&l->l_refcnt)) {
+		if (l->l_cleanup != NULL)
+			l->l_cleanup(l);
+		explicit_bzero(l, sizeof(*l));
+		free(l, M_NOISE);
+	}
 }
 
 void
-noise_local_unlock_identity(struct noise_local *l)
+noise_local_free(struct noise_local *l, void (*cleanup)(struct noise_local *))
 {
-	rw_exit_write(&l->l_identity_lock);
+	l->l_cleanup = cleanup;
+	noise_local_put(l);
 }
 
-int
-noise_local_set_private(struct noise_local *l,
-    const uint8_t private[NOISE_PUBLIC_KEY_LEN])
+void *
+noise_local_arg(struct noise_local *l)
 {
-	rw_assert_wrlock(&l->l_identity_lock);
+	return (l->l_arg);
+}
 
+void
+noise_local_private(struct noise_local *l, const uint8_t private[NOISE_PUBLIC_KEY_LEN])
+{
+	struct epoch_tracker et;
+	struct noise_remote *r;
+	size_t i;
+
+	rw_wlock(&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, private);
+	l->l_has_identity = curve25519_generate_public(l->l_public, l->l_private);
 
-	return l->l_has_identity ? 0 : ENXIO;
+	NET_EPOCH_ENTER(et);
+	for (i = 0; i < HT_REMOTE_SIZE; i++) {
+		CK_LIST_FOREACH(r, &l->l_remote_hash[i], r_entry) {
+			noise_precompute_ss(l, r);
+			noise_remote_expire_current(r);
+		}
+	}
+	NET_EPOCH_EXIT(et);
+	rw_wunlock(&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 ret = 0;
-	rw_enter_read(&l->l_identity_lock);
-	if (l->l_has_identity) {
+	int has_identity;
+	rw_rlock(&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);
-	} else {
-		ret = ENXIO;
 	}
-	rw_exit_read(&l->l_identity_lock);
-	return ret;
+	rw_runlock(&l->l_identity_lock);
+	return (has_identity ? 0 : ENXIO);
 }
 
-void
-noise_remote_init(struct noise_remote *r,
-    const uint8_t public[NOISE_PUBLIC_KEY_LEN], struct noise_local *l)
+static void
+noise_precompute_ss(struct noise_local *l, struct noise_remote *r)
+{
+	rw_wlock(&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_wunlock(&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])
 {
-	bzero(r, sizeof(*r));
+	struct noise_remote *r, *ri;
+	uint64_t idx;
+
+	if ((r = malloc(sizeof(*r), M_NOISE, 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);
+
+	refcount_init(&r->r_refcnt, 1);
+	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;
 
-	SLIST_INSERT_HEAD(&r->r_unused_keypairs, &r->r_keypair[0], kp_entry);
-	SLIST_INSERT_HEAD(&r->r_unused_keypairs, &r->r_keypair[1], kp_entry);
-	SLIST_INSERT_HEAD(&r->r_unused_keypairs, &r->r_keypair[2], kp_entry);
+	bzero(&r->r_smr, sizeof(r->r_smr));
 
-	KASSERT(l != NULL, ("must provide local"));
-	r->r_local = l;
+	/* Insert to hashtable */
+	idx = siphash24(&l->l_hash_key, public, NOISE_PUBLIC_KEY_LEN) & HT_REMOTE_MASK;
 
-	rw_enter_write(&l->l_identity_lock);
-	noise_remote_precompute(r);
-	rw_exit_write(&l->l_identity_lock);
+	rw_wlock(&l->l_remote_lock);
+	CK_LIST_FOREACH(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++;
+		CK_LIST_INSERT_HEAD(&l->l_remote_hash[idx], r, r_entry);
+	} else {
+free:
+		free(r, M_NOISE);
+		noise_local_put(l);
+		r = NULL;
+	}
+	rw_wunlock(&l->l_remote_lock);
+
+	return (r);
 }
 
-int
+struct noise_remote *
+noise_remote_lookup(struct noise_local *l, const uint8_t public[NOISE_PUBLIC_KEY_LEN])
+{
+	struct epoch_tracker et;
+	struct noise_remote *r, *ret = NULL;
+	uint64_t idx;
+
+	idx = siphash24(&l->l_hash_key, public, NOISE_PUBLIC_KEY_LEN) & HT_REMOTE_MASK;
+
+	NET_EPOCH_ENTER(et);
+	CK_LIST_FOREACH(r, &l->l_remote_hash[idx], r_entry) {
+		if (timingsafe_bcmp(r->r_public, public, NOISE_PUBLIC_KEY_LEN) == 0) {
+			if (refcount_acquire_if_not_zero(&r->r_refcnt))
+				ret = r;
+			break;
+		}
+	}
+	NET_EPOCH_EXIT(et);
+	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_wlock(&l->l_index_lock);
+assign_id:
+	r_i->i_local_index = arc4random();
+	idx = r_i->i_local_index & HT_INDEX_MASK;
+	CK_LIST_FOREACH(i, &l->l_index_hash[idx], i_entry)
+		if (i->i_local_index == r_i->i_local_index)
+			goto assign_id;
+
+	CK_LIST_INSERT_HEAD(&l->l_index_hash[idx], r_i, i_entry);
+	rw_wunlock(&l->l_index_lock);
+
+	r->r_handshake_alive = 1;
+}
+
+struct noise_remote *
+noise_remote_index_lookup(struct noise_local *l, uint32_t idx0)
+{
+	struct epoch_tracker et;
+	struct noise_index *i;
+	struct noise_remote *r, *ret = NULL;
+	uint32_t idx = idx0 & HT_INDEX_MASK;
+
+	NET_EPOCH_ENTER(et);
+	CK_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 (refcount_acquire_if_not_zero(&r->r_refcnt))
+				ret = r;
+			break;
+		}
+	}
+	NET_EPOCH_EXIT(et);
+	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_wlock(&l->l_index_lock);
+		CK_LIST_REMOVE(&r->r_index, i_entry);
+		rw_wunlock(&l->l_index_lock);
+		r->r_handshake_alive = 0;
+		return (1);
+	}
+	return (0);
+}
+
+struct noise_remote *
+noise_remote_ref(struct noise_remote *r)
+{
+	refcount_acquire(&r->r_refcnt);
+	return (r);
+}
+
+static void
+noise_remote_smr_free(struct epoch_context *smr)
+{
+	struct noise_remote *r;
+	r = __containerof(smr, struct noise_remote, r_smr);
+	if (r->r_cleanup != NULL)
+		r->r_cleanup(r);
+	noise_local_put(r->r_local);
+	explicit_bzero(r, sizeof(*r));
+	free(r, M_NOISE);
+}
+
+void
+noise_remote_put(struct noise_remote *r)
+{
+	if (refcount_release(&r->r_refcnt))
+		NET_EPOCH_CALL(noise_remote_smr_free, &r->r_smr);
+}
+
+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_wlock(&l->l_remote_lock);
+	CK_LIST_REMOVE(r, r_entry);
+	l->l_remote_num--;
+	rw_wunlock(&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])
 {
-	int same;
-	rw_enter_write(&r->r_handshake_lock);
-	same = !timingsafe_bcmp(r->r_psk, psk, NOISE_SYMMETRIC_KEY_LEN);
-	if (!same) {
+	rw_wlock(&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);
-	return same ? EEXIST : 0;
+	rw_wunlock(&r->r_handshake_lock);
 }
 
 int
@@ -162,35 +489,406 @@ noise_remote_keys(struct noise_remote *r, uint8_t public[NOISE_PUBLIC_KEY_LEN],
 	if (public != NULL)
 		memcpy(public, r->r_public, NOISE_PUBLIC_KEY_LEN);
 
-	rw_enter_read(&r->r_handshake_lock);
+	rw_rlock(&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);
+	rw_runlock(&r->r_handshake_lock);
 
-	/* If r_psk != null_psk return 0, else ENOENT (no psk) */
-	return ret ? 0 : ENOENT;
+	return (ret ? 0 : ENOENT);
+}
+
+int
+noise_remote_initiation_expired(struct noise_remote *r)
+{
+	int expired;
+	rw_rlock(&r->r_handshake_lock);
+	expired = noise_timer_expired(&r->r_last_sent, REKEY_TIMEOUT, 0);
+	rw_runlock(&r->r_handshake_lock);
+	return (expired);
 }
 
 void
-noise_remote_precompute(struct noise_remote *r)
+noise_remote_handshake_clear(struct noise_remote *r)
 {
-	struct noise_local *l = r->r_local;
-	rw_assert_wrlock(&l->l_identity_lock);
-	if (!l->l_has_identity)
-		bzero(r->r_ss, NOISE_PUBLIC_KEY_LEN);
-	else if (!curve25519(r->r_ss, l->l_private, r->r_public))
-		bzero(r->r_ss, NOISE_PUBLIC_KEY_LEN);
+	rw_wlock(&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_wunlock(&r->r_handshake_lock);
+}
+
+void
+noise_remote_keypairs_clear(struct noise_remote *r)
+{
+	struct noise_keypair *kp;
+
+	rw_wlock(&r->r_keypair_lock);
+	kp = epoch_ptr_read(&r->r_next);
+	epoch_ptr_write(&r->r_next, NULL);
+	noise_keypair_drop(kp);
+
+	kp = epoch_ptr_read(&r->r_current);
+	epoch_ptr_write(&r->r_current, NULL);
+	noise_keypair_drop(kp);
+
+	kp = epoch_ptr_read(&r->r_previous);
+	epoch_ptr_write(&r->r_previous, NULL);
+	noise_keypair_drop(kp);
+	rw_wunlock(&r->r_keypair_lock);
+}
+
+static void
+noise_remote_expire_current(struct noise_remote *r)
+{
+	struct epoch_tracker et;
+	struct noise_keypair *kp;
+
+	noise_remote_handshake_clear(r);
+
+	NET_EPOCH_ENTER(et);
+	kp = epoch_ptr_read(&r->r_next);
+	if (kp != NULL) WRITE_ONCE(kp->kp_can_send, 0);
+	kp = epoch_ptr_read(&r->r_current);
+	if (kp != NULL) WRITE_ONCE(kp->kp_can_send, 0);
+	NET_EPOCH_EXIT(et);
+}
+
+/* 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_wlock(&r->r_keypair_lock);
+	next = epoch_ptr_read(&r->r_next);
+	current = epoch_ptr_read(&r->r_current);
+	previous = epoch_ptr_read(&r->r_previous);
+
+	if (kp->kp_is_initiator) {
+		if (next != NULL) {
+			epoch_ptr_write(&r->r_next, NULL);
+			epoch_ptr_write(&r->r_previous, next);
+			noise_keypair_drop(current);
+		} else {
+			epoch_ptr_write(&r->r_previous, current);
+		}
+		noise_keypair_drop(previous);
+		epoch_ptr_write(&r->r_current, kp);
+	} else {
+		epoch_ptr_write(&r->r_next, kp);
+		noise_keypair_drop(next);
+		epoch_ptr_write(&r->r_previous, NULL);
+		noise_keypair_drop(previous);
+
+	}
+	rw_wunlock(&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_wlock(&l->l_index_lock);
+	CK_LIST_INSERT_BEFORE(r_i, &kp->kp_index, i_entry);
+	CK_LIST_REMOVE(r_i, i_entry);
+	rw_wunlock(&l->l_index_lock);
 
-	rw_enter_write(&r->r_handshake_lock);
-	noise_remote_handshake_index_drop(r);
 	explicit_bzero(&r->r_handshake, sizeof(r->r_handshake));
-	rw_exit_write(&r->r_handshake_lock);
 }
 
+static int
+noise_received_with(struct noise_keypair *kp)
+{
+	struct epoch_tracker et;
+	struct noise_keypair *old;
+	struct noise_remote *r = kp->kp_remote;
+
+	NET_EPOCH_ENTER(et);
+	if (kp != epoch_ptr_read(&r->r_next)) {
+		NET_EPOCH_EXIT(et);
+		return (0);
+	}
+	NET_EPOCH_EXIT(et);
+
+	rw_wlock(&r->r_keypair_lock);
+	if (kp != epoch_ptr_read(&r->r_next)) {
+		rw_wunlock(&r->r_keypair_lock);
+		return (0);
+	}
+
+	old = epoch_ptr_read(&r->r_previous);
+	epoch_ptr_write(&r->r_previous, epoch_ptr_read(&r->r_current));
+	noise_keypair_drop(old);
+	epoch_ptr_write(&r->r_current, kp);
+	epoch_ptr_write(&r->r_next, NULL);
+	rw_wunlock(&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_NOISE, M_NOWAIT)) == NULL)
+		return (ENOSPC);
+
+	refcount_init(&kp->kp_refcnt, 1);
+	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));
+	bzero(&kp->kp_smr, sizeof(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 epoch_tracker et;
+	struct noise_index *i;
+	struct noise_keypair *kp, *ret = NULL;
+	uint32_t idx = idx0 & HT_INDEX_MASK;
+
+	NET_EPOCH_ENTER(et);
+	CK_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 (refcount_acquire_if_not_zero(&kp->kp_refcnt))
+				ret = kp;
+			break;
+		}
+	}
+	NET_EPOCH_EXIT(et);
+	return (ret);
+}
+
+struct noise_keypair *
+noise_keypair_current(struct noise_remote *r)
+{
+	struct epoch_tracker et;
+	struct noise_keypair *kp, *ret = NULL;
+
+	NET_EPOCH_ENTER(et);
+	kp = epoch_ptr_read(&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 (refcount_acquire_if_not_zero(&kp->kp_refcnt))
+			ret = kp;
+	}
+	NET_EPOCH_EXIT(et);
+	return (ret);
+}
+
+struct noise_keypair *
+noise_keypair_ref(struct noise_keypair *kp)
+{
+	refcount_acquire(&kp->kp_refcnt);
+	return (kp);
+}
+
+static void
+noise_keypair_smr_free(struct epoch_context *smr)
+{
+	struct noise_keypair *kp;
+	kp = __containerof(smr, struct noise_keypair, kp_smr);
+	noise_remote_put(kp->kp_remote);
+	explicit_bzero(kp, sizeof(*kp));
+	free(kp, M_NOISE);
+}
+
+
+void
+noise_keypair_put(struct noise_keypair *kp)
+{
+	if (refcount_release(&kp->kp_refcnt))
+		NET_EPOCH_CALL(noise_keypair_smr_free, &kp->kp_smr);
+}
+
+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_wlock(&l->l_index_lock);
+	CK_LIST_REMOVE(&kp->kp_index, i_entry);
+	rw_wunlock(&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));
+}
+
+void *
+noise_keypair_remote_arg(struct noise_keypair *kp)
+{
+	return kp->kp_remote->r_arg;
+}
+
+
+
+int
+noise_keypair_nonce_next(struct noise_keypair *kp, uint64_t *send)
+{
+#ifdef __LP64__
+	*send = atomic_fetchadd_64(&kp->kp_nonce_send, 1);
+#else
+	rw_wlock(&kp->kp_nonce_lock);
+	*send = ctr->c_send++;
+	rw_wunlock(&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_wlock(&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_wunlock(&kp->kp_nonce_lock);
+	return (ret);
+}
+
+int
+noise_keep_key_fresh_send(struct noise_remote *r)
+{
+	struct epoch_tracker et;
+	struct noise_keypair *current;
+	int keep_key_fresh;
+
+	NET_EPOCH_ENTER(et);
+	current = epoch_ptr_read(&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)));
+	NET_EPOCH_EXIT(et);
+
+	return (keep_key_fresh ? ESTALE : 0);
+}
+
+int
+noise_keep_key_fresh_recv(struct noise_remote *r)
+{
+	struct epoch_tracker et;
+	struct noise_keypair *current;
+	int keep_key_fresh;
+
+	NET_EPOCH_ENTER(et);
+	current = epoch_ptr_read(&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);
+	NET_EPOCH_EXIT(et);
+
+	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,
+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])
@@ -200,10 +898,12 @@ noise_create_initiation(struct noise_remote *r, uint32_t *s_idx,
 	uint8_t key[NOISE_SYMMETRIC_KEY_LEN];
 	int ret = EINVAL;
 
-	rw_enter_read(&l->l_identity_lock);
-	rw_enter_write(&r->r_handshake_lock);
+	rw_rlock(&l->l_identity_lock);
+	rw_wlock(&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 */
@@ -229,21 +929,22 @@ noise_create_initiation(struct noise_remote *r, uint32_t *s_idx,
 	noise_msg_encrypt(ets, ets,
 	    NOISE_TIMESTAMP_LEN, key, hs->hs_hash);
 
-	noise_remote_handshake_index_drop(r);
-	hs->hs_state = CREATED_INITIATION;
-	hs->hs_local_index = noise_remote_handshake_index_get(r);
-	*s_idx = hs->hs_local_index;
+	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);
+	rw_wunlock(&r->r_handshake_lock);
+	rw_runlock(&l->l_identity_lock);
 	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
-	return ret;
+	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],
+    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])
 {
@@ -254,7 +955,7 @@ noise_consume_initiation(struct noise_local *l, struct noise_remote **rp,
 	uint8_t	timestamp[NOISE_TIMESTAMP_LEN];
 	int ret = EINVAL;
 
-	rw_enter_read(&l->l_identity_lock);
+	rw_rlock(&l->l_identity_lock);
 	if (!l->l_has_identity)
 		goto error;
 	noise_param_init(hs.hs_ck, hs.hs_hash, l->l_public);
@@ -272,23 +973,23 @@ noise_consume_initiation(struct noise_local *l, struct noise_remote **rp,
 		goto error;
 
 	/* Lookup the remote we received from */
-	if ((r = l->l_upcall.u_remote_get(l->l_upcall.u_arg, r_public)) == NULL)
+	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;
+		goto error_put;
 
 	/* {t} */
 	if (noise_msg_decrypt(timestamp, ets,
 	    NOISE_TIMESTAMP_LEN + NOISE_AUTHTAG_LEN, key, hs.hs_hash) != 0)
-		goto error;
+		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(&r->r_handshake_lock);
+	rw_wlock(&r->r_handshake_lock);
 
 	/* Replay */
 	if (memcmp(timestamp, r->r_timestamp, NOISE_TIMESTAMP_LEN) > 0)
@@ -296,41 +997,45 @@ noise_consume_initiation(struct noise_local *l, struct noise_remote **rp,
 	else
 		goto error_set;
 	/* Flood attack */
-	if (noise_timer_expired(&r->r_last_init, 0, REJECT_INTERVAL))
-		getnanouptime(&r->r_last_init);
+	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_handshake_index_drop(r);
-	hs.hs_state = CONSUMED_INITIATION;
-	hs.hs_local_index = noise_remote_handshake_index_get(r);
-	hs.hs_remote_index = s_idx;
+	noise_remote_index_insert(l, r);
+	r->r_index.i_remote_index = s_idx;
+	r->r_handshake_initiator = 0;
 	r->r_handshake = hs;
-	*rp = r;
+	*rp = noise_remote_ref(r);
 	ret = 0;
 error_set:
-	rw_exit_write(&r->r_handshake_lock);
+	rw_wunlock(&r->r_handshake_lock);
+error_put:
+	noise_remote_put(r);
 error:
-	rw_exit_read(&l->l_identity_lock);
+	rw_runlock(&l->l_identity_lock);
 	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
 	explicit_bzero(&hs, sizeof(hs));
-	return ret;
+	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])
+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(&r->r_local->l_identity_lock);
-	rw_enter_write(&r->r_handshake_lock);
+	rw_rlock(&l->l_identity_lock);
+	rw_wlock(&r->r_handshake_lock);
 
-	if (hs->hs_state != CONSUMED_INITIATION)
+	if (!r->r_handshake_alive || r->r_handshake_initiator)
 		goto error;
 
 	/* e */
@@ -353,51 +1058,57 @@ noise_create_response(struct noise_remote *r, uint32_t *s_idx, uint32_t *r_idx,
 	/* {} */
 	noise_msg_encrypt(en, NULL, 0, key, hs->hs_hash);
 
-	hs->hs_state = CREATED_RESPONSE;
-	*r_idx = hs->hs_remote_index;
-	*s_idx = hs->hs_local_index;
-	ret = 0;
+	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(&r->r_local->l_identity_lock);
+	rw_wunlock(&r->r_handshake_lock);
+	rw_runlock(&l->l_identity_lock);
 	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
 	explicit_bzero(e, NOISE_PUBLIC_KEY_LEN);
-	return ret;
+	return (ret);
 }
 
 int
-noise_consume_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])
+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])
 {
-	struct noise_local *l = r->r_local;
-	struct noise_handshake hs;
+	uint8_t preshared_key[NOISE_SYMMETRIC_KEY_LEN];
 	uint8_t key[NOISE_SYMMETRIC_KEY_LEN];
-	uint8_t preshared_key[NOISE_PUBLIC_KEY_LEN];
+	struct noise_handshake hs;
+	struct noise_remote *r = NULL;
 	int ret = EINVAL;
 
-	rw_enter_read(&l->l_identity_lock);
+	if ((r = noise_remote_index_lookup(l, r_idx)) == NULL)
+		return (ret);
+
+	rw_rlock(&l->l_identity_lock);
 	if (!l->l_has_identity)
 		goto error;
 
-	rw_enter_read(&r->r_handshake_lock);
-	hs = r->r_handshake;
-	memcpy(preshared_key, r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
-	rw_exit_read(&r->r_handshake_lock);
-
-	if (hs.hs_state != CREATED_INITIATION ||
-	    hs.hs_local_index != r_idx)
+	rw_rlock(&r->r_handshake_lock);
+	if (!r->r_handshake_alive || !r->r_handshake_initiator) {
+		rw_runlock(&r->r_handshake_lock);
 		goto error;
+	}
+	memcpy(preshared_key, r->r_psk, NOISE_SYMMETRIC_KEY_LEN);
+	hs = r->r_handshake;
+	rw_runlock(&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;
+		goto error_zero;
 
 	/* se */
 	if (noise_mix_dh(hs.hs_ck, NULL, l->l_private, ue) != 0)
-		goto error;
+		goto error_zero;
 
 	/* psk */
 	noise_mix_psk(hs.hs_ck, hs.hs_hash, key, preshared_key);
@@ -405,365 +1116,28 @@ noise_consume_response(struct noise_remote *r, uint32_t s_idx, uint32_t r_idx,
 	/* {} */
 	if (noise_msg_decrypt(NULL, en,
 	    0 + NOISE_AUTHTAG_LEN, key, hs.hs_hash) != 0)
-		goto error;
+		goto error_zero;
 
-	hs.hs_remote_index = s_idx;
-
-	rw_enter_write(&r->r_handshake_lock);
-	if (r->r_handshake.hs_state == hs.hs_state &&
-	    r->r_handshake.hs_local_index == hs.hs_local_index) {
+	rw_wlock(&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_handshake.hs_state = CONSUMED_RESPONSE;
-		ret = 0;
+		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:
-	rw_exit_read(&l->l_identity_lock);
-	explicit_bzero(&hs, sizeof(hs));
+	rw_wunlock(&r->r_handshake_lock);
+error_zero:
+	explicit_bzero(preshared_key, NOISE_SYMMETRIC_KEY_LEN);
 	explicit_bzero(key, NOISE_SYMMETRIC_KEY_LEN);
-	return ret;
-}
-
-int
-noise_remote_begin_session(struct noise_remote *r)
-{
-	struct noise_handshake *hs = &r->r_handshake;
-	struct noise_keypair kp, *next, *current, *previous;
-
-	rw_enter_write(&r->r_handshake_lock);
-
-	/* We now derive the keypair from the handshake */
-	if (hs->hs_state == CONSUMED_RESPONSE) {
-		kp.kp_is_initiator = 1;
-		noise_kdf(kp.kp_send, kp.kp_recv, NULL, NULL,
-		    NOISE_SYMMETRIC_KEY_LEN, NOISE_SYMMETRIC_KEY_LEN, 0, 0,
-		    hs->hs_ck);
-	} else if (hs->hs_state == CREATED_RESPONSE) {
-		kp.kp_is_initiator = 0;
-		noise_kdf(kp.kp_recv, kp.kp_send, NULL, NULL,
-		    NOISE_SYMMETRIC_KEY_LEN, NOISE_SYMMETRIC_KEY_LEN, 0, 0,
-		    hs->hs_ck);
-	} else {
-		rw_exit_write(&r->r_handshake_lock);
-		return EINVAL;
-	}
-
-	kp.kp_valid = 1;
-	kp.kp_local_index = hs->hs_local_index;
-	kp.kp_remote_index = hs->hs_remote_index;
-	getnanouptime(&kp.kp_birthdate);
-	bzero(&kp.kp_ctr, sizeof(kp.kp_ctr));
-	rw_init(&kp.kp_ctr.c_lock, "noise_counter");
-
-	/* Now we need to add_new_keypair */
-	rw_enter_write(&r->r_keypair_lock);
-	next = r->r_next;
-	current = r->r_current;
-	previous = r->r_previous;
-
-	if (kp.kp_is_initiator) {
-		if (next != NULL) {
-			r->r_next = NULL;
-			r->r_previous = next;
-			noise_remote_keypair_free(r, current);
-		} else {
-			r->r_previous = current;
-		}
-
-		noise_remote_keypair_free(r, previous);
-
-		r->r_current = noise_remote_keypair_allocate(r);
-		*r->r_current = kp;
-	} else {
-		noise_remote_keypair_free(r, next);
-		r->r_previous = NULL;
-		noise_remote_keypair_free(r, previous);
-
-		r->r_next = noise_remote_keypair_allocate(r);
-		*r->r_next = kp;
-	}
-	rw_exit_write(&r->r_keypair_lock);
-
-	explicit_bzero(&r->r_handshake, sizeof(r->r_handshake));
-	rw_exit_write(&r->r_handshake_lock);
-
-	explicit_bzero(&kp, sizeof(kp));
-	return 0;
-}
-
-void
-noise_remote_clear(struct noise_remote *r)
-{
-	rw_enter_write(&r->r_handshake_lock);
-	noise_remote_handshake_index_drop(r);
-	explicit_bzero(&r->r_handshake, sizeof(r->r_handshake));
-	rw_exit_write(&r->r_handshake_lock);
-
-	rw_enter_write(&r->r_keypair_lock);
-	noise_remote_keypair_free(r, r->r_next);
-	noise_remote_keypair_free(r, r->r_current);
-	noise_remote_keypair_free(r, r->r_previous);
-	r->r_next = NULL;
-	r->r_current = NULL;
-	r->r_previous = NULL;
-	rw_exit_write(&r->r_keypair_lock);
-}
-
-void
-noise_remote_expire_current(struct noise_remote *r)
-{
-	rw_enter_write(&r->r_keypair_lock);
-	if (r->r_next != NULL)
-		r->r_next->kp_valid = 0;
-	if (r->r_current != NULL)
-		r->r_current->kp_valid = 0;
-	rw_exit_write(&r->r_keypair_lock);
-}
-
-int
-noise_remote_ready(struct noise_remote *r)
-{
-	struct noise_keypair *kp;
-	int ret;
-
-	rw_enter_read(&r->r_keypair_lock);
-	/* kp_ctr isn't locked here, we're happy to accept a racy read. */
-	if ((kp = r->r_current) == NULL ||
-	    !kp->kp_valid ||
-	    noise_timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME, 0) ||
-	    kp->kp_ctr.c_recv >= REJECT_AFTER_MESSAGES ||
-	    kp->kp_ctr.c_send >= REJECT_AFTER_MESSAGES)
-		ret = EINVAL;
-	else
-		ret = 0;
-	rw_exit_read(&r->r_keypair_lock);
-	return ret;
-}
-
-int
-noise_remote_encrypt(struct noise_remote *r, uint32_t *r_idx, uint64_t *nonce,
-    uint8_t *buf, size_t buflen)
-{
-	struct noise_keypair *kp;
-	int ret = EINVAL;
-
-	rw_enter_read(&r->r_keypair_lock);
-	if ((kp = r->r_current) == NULL)
-		goto error;
-
-	/* We confirm that our values are within our tolerances. We want:
-	 *  - a valid keypair
-	 *  - our keypair to be less than REJECT_AFTER_TIME seconds old
-	 *  - our receive counter to be less than REJECT_AFTER_MESSAGES
-	 *  - our send counter to be less than REJECT_AFTER_MESSAGES
-	 *
-	 * kp_ctr isn't locked here, we're happy to accept a racy read. */
-	if (!kp->kp_valid ||
-	    noise_timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME, 0) ||
-	    kp->kp_ctr.c_recv >= REJECT_AFTER_MESSAGES ||
-	    ((*nonce = noise_counter_send(&kp->kp_ctr)) > REJECT_AFTER_MESSAGES))
-		goto error;
-
-	/* We encrypt into the same buffer, so the caller must ensure that buf
-	 * has NOISE_AUTHTAG_LEN bytes to store the MAC. The nonce and index
-	 * are passed back out to the caller through the provided data pointer. */
-	*r_idx = kp->kp_remote_index;
-	chacha20poly1305_encrypt(buf, buf, buflen,
-	    NULL, 0, *nonce, kp->kp_send);
-
-	/* If our values are still within tolerances, but we are approaching
-	 * the tolerances, we notify the caller with ESTALE that they should
-	 * establish a new keypair. The current keypair can continue to be used
-	 * until the tolerances are hit. We notify if:
-	 *  - our send counter is valid and not less than REKEY_AFTER_MESSAGES
-	 *  - we're the initiator and our keypair is older than
-	 *    REKEY_AFTER_TIME seconds */
-	ret = ESTALE;
-	if ((kp->kp_valid && *nonce >= REKEY_AFTER_MESSAGES) ||
-	    (kp->kp_is_initiator &&
-	    noise_timer_expired(&kp->kp_birthdate, REKEY_AFTER_TIME, 0)))
-		goto error;
-
-	ret = 0;
-error:
-	rw_exit_read(&r->r_keypair_lock);
-	return ret;
-}
-
-int
-noise_remote_decrypt(struct noise_remote *r, uint32_t r_idx, uint64_t nonce,
-    uint8_t *buf, size_t buflen)
-{
-	struct noise_keypair *kp;
-	int ret = EINVAL;
-
-	/* We retrieve the keypair corresponding to the provided index. We
-	 * attempt the current keypair first as that is most likely. We also
-	 * want to make sure that the keypair is valid as it would be
-	 * catastrophic to decrypt against a zero'ed keypair. */
-	rw_enter_read(&r->r_keypair_lock);
-
-	if (r->r_current != NULL && r->r_current->kp_local_index == r_idx) {
-		kp = r->r_current;
-	} else if (r->r_previous != NULL && r->r_previous->kp_local_index == r_idx) {
-		kp = r->r_previous;
-	} else if (r->r_next != NULL && r->r_next->kp_local_index == r_idx) {
-		kp = r->r_next;
-	} else {
-		goto error;
-	}
-
-	/* We confirm that our values are within our tolerances. These values
-	 * are the same as the encrypt routine.
-	 *
-	 * kp_ctr isn't locked here, we're happy to accept a racy read. */
-	if (noise_timer_expired(&kp->kp_birthdate, REJECT_AFTER_TIME, 0) ||
-	    kp->kp_ctr.c_recv >= REJECT_AFTER_MESSAGES)
-		goto error;
-
-	/* Decrypt, then validate the counter. We don't want to validate the
-	 * counter before decrypting as we do not know the message is authentic
-	 * prior to decryption. */
-	if (chacha20poly1305_decrypt(buf, buf, buflen,
-	    NULL, 0, nonce, kp->kp_recv) == 0)
-		goto error;
-
-	if (noise_counter_recv(&kp->kp_ctr, nonce) != 0)
-		goto error;
-
-	/* If we've received the handshake confirming data packet then move the
-	 * next keypair into current. If we do slide the next keypair in, then
-	 * we skip the REKEY_AFTER_TIME_RECV check. This is safe to do as a
-	 * data packet can't confirm a session that we are an INITIATOR of. */
-	if (kp == r->r_next) {
-		rw_exit_read(&r->r_keypair_lock);
-		rw_enter_write(&r->r_keypair_lock);
-		if (kp == r->r_next && kp->kp_local_index == r_idx) {
-			noise_remote_keypair_free(r, r->r_previous);
-			r->r_previous = r->r_current;
-			r->r_current = r->r_next;
-			r->r_next = NULL;
-
-			ret = ECONNRESET;
-			goto error;
-		}
-		rw_enter(&r->r_keypair_lock, RW_DOWNGRADE);
-	}
-
-	/* Similar to when we encrypt, we want to notify the caller when we
-	 * are approaching our tolerances. We notify if:
-	 *  - we're the initiator and the current keypair is older than
-	 *    REKEY_AFTER_TIME_RECV seconds. */
-	ret = ESTALE;
-	kp = r->r_current;
-	if (kp != NULL &&
-	    kp->kp_valid &&
-	    kp->kp_is_initiator &&
-	    noise_timer_expired(&kp->kp_birthdate, REKEY_AFTER_TIME_RECV, 0))
-		goto error;
-
-	ret = 0;
-
-error:
-	rw_exit(&r->r_keypair_lock);
-	return ret;
-}
-
-/* Private functions - these should not be called outside this file under any
- * circumstances. */
-static struct noise_keypair *
-noise_remote_keypair_allocate(struct noise_remote *r)
-{
-	struct noise_keypair *kp;
-	kp = SLIST_FIRST(&r->r_unused_keypairs);
-	SLIST_REMOVE_HEAD(&r->r_unused_keypairs, kp_entry);
-	return kp;
-}
-
-static void
-noise_remote_keypair_free(struct noise_remote *r, struct noise_keypair *kp)
-{
-	struct noise_upcall *u = &r->r_local->l_upcall;
-	if (kp != NULL) {
-		SLIST_INSERT_HEAD(&r->r_unused_keypairs, kp, kp_entry);
-		u->u_index_drop(u->u_arg, kp->kp_local_index);
-		bzero(kp->kp_send, sizeof(kp->kp_send));
-		bzero(kp->kp_recv, sizeof(kp->kp_recv));
-	}
-}
-
-static uint32_t
-noise_remote_handshake_index_get(struct noise_remote *r)
-{
-	struct noise_upcall *u = &r->r_local->l_upcall;
-	return u->u_index_set(u->u_arg, r);
-}
-
-static void
-noise_remote_handshake_index_drop(struct noise_remote *r)
-{
-	struct noise_handshake *hs = &r->r_handshake;
-	struct noise_upcall *u = &r->r_local->l_upcall;
-	rw_assert_wrlock(&r->r_handshake_lock);
-	if (hs->hs_state != HS_ZEROED)
-		u->u_index_drop(u->u_arg, hs->hs_local_index);
-}
-
-static uint64_t
-noise_counter_send(struct noise_counter *ctr)
-{
-	uint64_t ret;
-	rw_enter_write(&ctr->c_lock);
-	ret = ctr->c_send++;
-	rw_exit_write(&ctr->c_lock);
-	return ret;
-}
-
-static int
-noise_counter_recv(struct noise_counter *ctr, uint64_t recv)
-{
-	uint64_t i, top, index_recv, index_ctr;
-	unsigned long bit;
-	int ret = EEXIST;
-
-	rw_enter_write(&ctr->c_lock);
-
-	/* Check that the recv counter is valid */
-	if (ctr->c_recv >= REJECT_AFTER_MESSAGES ||
-	    recv >= REJECT_AFTER_MESSAGES)
-		goto error;
-
-	/* If the packet is out of the window, invalid */
-	if (recv + COUNTER_WINDOW_SIZE < ctr->c_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 = ctr->c_recv / COUNTER_BITS;
-
-	if (recv > ctr->c_recv) {
-		top = MIN(index_recv - index_ctr, COUNTER_NUM);
-		for (i = 1; i <= top; i++)
-			ctr->c_backtrack[
-			    (i + index_ctr) & (COUNTER_NUM - 1)] = 0;
-		ctr->c_recv = recv;
-	}
-
-	index_recv %= COUNTER_NUM;
-	bit = 1ul << (recv % COUNTER_BITS);
-
-	if (ctr->c_backtrack[index_recv] & bit)
-		goto error;
-
-	ctr->c_backtrack[index_recv] |= bit;
-
-	ret = 0;
+	explicit_bzero(&hs, sizeof(hs));
 error:
-	rw_exit_write(&ctr->c_lock);
-	return ret;
+	rw_runlock(&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,
@@ -772,13 +1146,6 @@ noise_kdf(uint8_t *a, uint8_t *b, uint8_t *c, const uint8_t *x,
 	uint8_t out[BLAKE2S_HASH_SIZE + 1];
 	uint8_t sec[BLAKE2S_HASH_SIZE];
 
-#ifdef DIAGNOSTIC
-	MPASS(a_len <= BLAKE2S_HASH_SIZE && b_len <= BLAKE2S_HASH_SIZE &&
-			c_len <= BLAKE2S_HASH_SIZE);
-	MPASS(!(b || b_len || c || c_len) || (a && a_len));
-	MPASS(!(c || c_len) || (b && b_len));
-#endif
-
 	/* Extract entropy from "x" into sec */
 	blake2s_hmac(sec, x, ck, BLAKE2S_HASH_SIZE, x_len, NOISE_HASH_LEN);
 
@@ -822,11 +1189,11 @@ noise_mix_dh(uint8_t ck[NOISE_HASH_LEN], uint8_t key[NOISE_SYMMETRIC_KEY_LEN],
 	uint8_t dh[NOISE_PUBLIC_KEY_LEN];
 
 	if (!curve25519(dh, private, public))
-		return EINVAL;
+		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;
+	return (0);
 }
 
 static int
@@ -835,10 +1202,10 @@ noise_mix_ss(uint8_t ck[NOISE_HASH_LEN], uint8_t key[NOISE_SYMMETRIC_KEY_LEN],
 {
 	static uint8_t null_point[NOISE_PUBLIC_KEY_LEN];
 	if (timingsafe_bcmp(ss, null_point, NOISE_PUBLIC_KEY_LEN) == 0)
-		return ENOENT;
+		return (ENOENT);
 	noise_kdf(ck, key, NULL, ss,
 	    NOISE_HASH_LEN, NOISE_SYMMETRIC_KEY_LEN, 0, NOISE_PUBLIC_KEY_LEN, ck);
-	return 0;
+	return (0);
 }
 
 static void
@@ -901,9 +1268,9 @@ noise_msg_decrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
 	/* Nonce always zero for Noise_IK */
 	if (!chacha20poly1305_decrypt(dst, src, src_len,
 	    hash, NOISE_HASH_LEN, 0, key))
-		return EINVAL;
+		return (EINVAL);
 	noise_mix_hash(hash, src, src_len);
-	return 0;
+	return (0);
 }
 
 static void
@@ -942,11 +1309,7 @@ noise_timer_expired(struct timespec *birthdate, time_t sec, long nsec)
 	struct timespec uptime;
 	struct timespec expire = { .tv_sec = sec, .tv_nsec = nsec };
 
-	/* We don't really worry about a zeroed birthdate, to avoid the extra
-	 * check on every encrypt/decrypt. This does mean that r_last_init
-	 * check may fail if getnanouptime is < REJECT_INTERVAL from 0. */
-
 	getnanouptime(&uptime);
 	timespecadd(birthdate, &expire, &expire);
-	return timespeccmp(&uptime, &expire, >) ? ETIMEDOUT : 0;
+	return (timespeccmp(&uptime, &expire, >) ? ETIMEDOUT : 0);
 }
diff --git a/src/wg_noise.h b/src/wg_noise.h
index 198ddd1..50890a3 100644
--- a/src/wg_noise.h
+++ b/src/wg_noise.h
@@ -19,116 +19,86 @@
 #define NOISE_AUTHTAG_LEN	CHACHA20POLY1305_AUTHTAG_SIZE
 #define NOISE_HASH_LEN		BLAKE2S_HASH_SIZE
 
-/* 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_AFTER_TIME	180
-#define REJECT_INTERVAL		(1000000000 / 50) /* fifty times per sec */
-/* 24 = floor(log2(REJECT_INTERVAL)) */
-#define REJECT_INTERVAL_MASK	(~((1ull<<24)-1))
-
-enum noise_state_hs {
-	HS_ZEROED = 0,
-	CREATED_INITIATION,
-	CONSUMED_INITIATION,
-	CREATED_RESPONSE,
-	CONSUMED_RESPONSE,
-};
-
-struct noise_handshake {
-	enum noise_state_hs	 hs_state;
-	uint32_t		 hs_local_index;
-	uint32_t		 hs_remote_index;
-	uint8_t		 	 hs_e[NOISE_PUBLIC_KEY_LEN];
-	uint8_t		 	 hs_hash[NOISE_HASH_LEN];
-	uint8_t		 	 hs_ck[NOISE_HASH_LEN];
-};
-
-struct noise_counter {
-	struct rwlock		 c_lock;
-	uint64_t		 c_send;
-	uint64_t		 c_recv;
-	unsigned long		 c_backtrack[COUNTER_NUM];
-};
-
-struct noise_keypair {
-	SLIST_ENTRY(noise_keypair)	kp_entry;
-	int				kp_valid;
-	int				kp_is_initiator;
-	uint32_t			kp_local_index;
-	uint32_t			kp_remote_index;
-	uint8_t				kp_send[NOISE_SYMMETRIC_KEY_LEN];
-	uint8_t				kp_recv[NOISE_SYMMETRIC_KEY_LEN];
-	struct timespec			kp_birthdate; /* nanouptime */
-	struct noise_counter		kp_ctr;
-};
-
-struct noise_remote {
-	uint8_t				 r_public[NOISE_PUBLIC_KEY_LEN];
-	struct noise_local		*r_local;
-	uint8_t		 		 r_ss[NOISE_PUBLIC_KEY_LEN];
-
-	struct rwlock			 r_handshake_lock;
-	struct noise_handshake		 r_handshake;
-	uint8_t				 r_psk[NOISE_SYMMETRIC_KEY_LEN];
-	uint8_t				 r_timestamp[NOISE_TIMESTAMP_LEN];
-	struct timespec			 r_last_init; /* nanouptime */
-
-	struct rwlock			 r_keypair_lock;
-	SLIST_HEAD(,noise_keypair)	 r_unused_keypairs;
-	struct noise_keypair		*r_next, *r_current, *r_previous;
-	struct noise_keypair		 r_keypair[3]; /* 3: next, current, previous. */
-
-};
-
-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 noise_upcall {
-		void	 *u_arg;
-		struct noise_remote *
-			(*u_remote_get)(void *, uint8_t[NOISE_PUBLIC_KEY_LEN]);
-		uint32_t
-			(*u_index_set)(void *, struct noise_remote *);
-		void	(*u_index_drop)(void *, uint32_t);
-	}			l_upcall;
-};
-
-/* Set/Get noise parameters */
-void	noise_local_init(struct noise_local *, struct noise_upcall *);
-void	noise_local_lock_identity(struct noise_local *);
-void	noise_local_unlock_identity(struct noise_local *);
-int	noise_local_set_private(struct noise_local *,
+#define REKEY_TIMEOUT		5
+#define KEEPALIVE_TIMEOUT	10
+
+struct noise_local;
+struct noise_remote;
+struct noise_keypair;
+
+/* Local configuration */
+struct noise_local *
+	noise_local_alloc(void *);
+struct noise_local *
+	noise_local_ref(struct noise_local *);
+void	noise_local_put(struct noise_local *);
+void	noise_local_free(struct noise_local *, void (*)(struct noise_local *));
+void *	noise_local_arg(struct noise_local *);
+
+void	noise_local_private(struct noise_local *,
 	    const uint8_t[NOISE_PUBLIC_KEY_LEN]);
-int	noise_local_keys(struct noise_local *, uint8_t[NOISE_PUBLIC_KEY_LEN],
+int	noise_local_keys(struct noise_local *,
+	    uint8_t[NOISE_PUBLIC_KEY_LEN],
 	    uint8_t[NOISE_PUBLIC_KEY_LEN]);
 
-void	noise_remote_init(struct noise_remote *,
-	    const uint8_t[NOISE_PUBLIC_KEY_LEN], struct noise_local *);
-int	noise_remote_set_psk(struct noise_remote *,
+/* Remote configuration */
+struct noise_remote *
+	noise_remote_alloc(struct noise_local *, void *,
+	    const uint8_t[NOISE_PUBLIC_KEY_LEN],
 	    const uint8_t[NOISE_SYMMETRIC_KEY_LEN]);
-int	noise_remote_keys(struct noise_remote *, uint8_t[NOISE_PUBLIC_KEY_LEN],
+struct noise_remote *
+	noise_remote_lookup(struct noise_local *, const uint8_t[NOISE_PUBLIC_KEY_LEN]);
+struct noise_remote *
+	noise_remote_index_lookup(struct noise_local *, uint32_t);
+struct noise_remote *
+	noise_remote_ref(struct noise_remote *);
+void	noise_remote_put(struct noise_remote *);
+void	noise_remote_free(struct noise_remote *, void (*)(struct noise_remote *));
+struct noise_local *
+	noise_remote_local(struct noise_remote *);
+void *	noise_remote_arg(struct noise_remote *);
+
+void	noise_remote_set_psk(struct noise_remote *,
+	    const uint8_t[NOISE_SYMMETRIC_KEY_LEN]);
+int	noise_remote_keys(struct noise_remote *,
+	    uint8_t[NOISE_PUBLIC_KEY_LEN],
 	    uint8_t[NOISE_SYMMETRIC_KEY_LEN]);
+int	noise_remote_initiation_expired(struct noise_remote *);
+void	noise_remote_handshake_clear(struct noise_remote *);
+void	noise_remote_keypairs_clear(struct noise_remote *);
+
+/* Keypair functions */
+struct noise_keypair *
+	noise_keypair_lookup(struct noise_local *, uint32_t);
+struct noise_keypair *
+	noise_keypair_current(struct noise_remote *);
+struct noise_keypair *
+	noise_keypair_ref(struct noise_keypair *);
+void	noise_keypair_put(struct noise_keypair *);
+
+struct noise_remote *
+	noise_keypair_remote(struct noise_keypair *);
+void *	noise_keypair_remote_arg(struct noise_keypair *);
+
+int	noise_keypair_nonce_next(struct noise_keypair *, uint64_t *);
+int	noise_keypair_nonce_check(struct noise_keypair *, uint64_t);
+
+int	noise_keep_key_fresh_send(struct noise_remote *);
+int	noise_keep_key_fresh_recv(struct noise_remote *);
+void	noise_keypair_encrypt(
+	    struct noise_keypair *,
+	    uint32_t *r_idx,
+	    uint64_t nonce,
+	    uint8_t *buf,
+	    size_t buflen);
+int	noise_keypair_decrypt(
+	    struct noise_keypair *,
+	    uint64_t nonce,
+	    uint8_t *buf,
+	    size_t buflen);
 
-/* Should be called anytime noise_local_set_private is called */
-void	noise_remote_precompute(struct noise_remote *);
-
-/* Cryptographic functions */
+/* Handshake functions */
 int	noise_create_initiation(
 	    struct noise_remote *,
 	    uint32_t *s_idx,
@@ -152,29 +122,11 @@ int	noise_create_response(
 	    uint8_t en[0 + NOISE_AUTHTAG_LEN]);
 
 int	noise_consume_response(
-	    struct noise_remote *,
+	    struct noise_local *,
+	    struct noise_remote **,
 	    uint32_t s_idx,
 	    uint32_t r_idx,
 	    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
 	    uint8_t en[0 + NOISE_AUTHTAG_LEN]);
 
-int	noise_remote_begin_session(struct noise_remote *);
-void	noise_remote_clear(struct noise_remote *);
-void	noise_remote_expire_current(struct noise_remote *);
-
-int	noise_remote_ready(struct noise_remote *);
-
-int	noise_remote_encrypt(
-	    struct noise_remote *,
-	    uint32_t *r_idx,
-	    uint64_t *nonce,
-	    uint8_t *buf,
-	    size_t buflen);
-int	noise_remote_decrypt(
-	    struct noise_remote *,
-	    uint32_t r_idx,
-	    uint64_t nonce,
-	    uint8_t *buf,
-	    size_t buflen);
-
 #endif /* __NOISE_H__ */