wg_noise: insert/remove peer independent of alloc/destroy

This is needed, to remove the peer from the public key hashtable before
calling noise_remote_destroy. This will prevent any incoming handshakes
from starting in that time. It also cleans up the insert path to make it
more like it was before the wg_noise EPOCH changes.

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

3 files changed, 49 insertions, 34 deletions

diff --git a/src/if_wg.c b/src/if_wg.c
index aab70bd..850d16a 100644
--- a/src/if_wg.c
+++ b/src/if_wg.c
@@ -326,7 +326,7 @@ static void wg_timers_run_persistent_keepalive(void *);
 static int wg_aip_add(struct wg_softc *, struct wg_peer *, sa_family_t, const void *, uint8_t);
 static struct wg_peer *wg_aip_lookup(struct wg_softc *, sa_family_t, void *);
 static void wg_aip_remove_all(struct wg_softc *, struct wg_peer *);
-static struct wg_peer *wg_peer_alloc(struct wg_softc *);
+static struct wg_peer *wg_peer_alloc(struct wg_softc *, const uint8_t [WG_KEY_SIZE]);
 static void wg_peer_free_deferred(struct noise_remote *);
 static void wg_peer_destroy(struct wg_peer *);
 static void wg_peer_destroy_all(struct wg_softc *);
@@ -387,7 +387,7 @@ static void wg_module_deinit(void);
 
 /* TODO Peer */
 static struct wg_peer *
-wg_peer_alloc(struct wg_softc *sc)
+wg_peer_alloc(struct wg_softc *sc, const uint8_t pub_key[WG_KEY_SIZE])
 {
 	struct wg_peer *peer;
 
@@ -399,6 +399,12 @@ wg_peer_alloc(struct wg_softc *sc)
 	LIST_INIT(&peer->p_aips);
 	peer->p_aips_num = 0;
 
+	if ((peer->p_remote = noise_remote_alloc(sc->sc_local, peer, pub_key)) == NULL) {
+		free(peer, M_WG);
+		return NULL;
+	}
+	cookie_maker_init(&peer->p_cookie, pub_key);
+
 	rw_init(&peer->p_endpoint_lock, "wg_peer_endpoint");
 	wg_queue_init(&peer->p_stage_queue, "stageq");
 	wg_queue_init(&peer->p_encrypt_serial, "txq");
@@ -443,9 +449,8 @@ static void
 wg_peer_destroy(struct wg_peer *peer)
 {
 	sx_assert(&peer->p_sc->sc_lock, SX_XLOCKED);
-	/* Callers should already have called:
-	 *    wg_hashtable_peer_remove(&sc->sc_hashtable, peer);
-	 */
+
+	noise_remote_disable(peer->p_remote);
 	wg_aip_remove_all(peer->p_sc, peer);
 
 	/* We disable all timers, so we can't call the following tasks. */
@@ -2084,10 +2089,11 @@ wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 		wg_aip_remove_all(sc, peer);
 	}
 	if (peer == NULL) {
+		if ((peer = wg_peer_alloc(sc, pub_key)) == NULL) {
+			err = ENOMEM;
+			goto out;
+		}
 		need_insert = true;
-		peer = wg_peer_alloc(sc);
-		MPASS(peer != NULL);
-		cookie_maker_init(&peer->p_cookie, pub_key);
 	}
 	if (nvlist_exists_binary(nvl, "endpoint")) {
 		endpoint = nvlist_get_binary(nvl, "endpoint", &size);
@@ -2103,8 +2109,7 @@ wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 			err = EINVAL;
 			goto out;
 		}
-		if (!need_insert)
-			noise_remote_set_psk(peer->p_remote, preshared_key);
+		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");
@@ -2147,9 +2152,7 @@ wg_peer_add(struct wg_softc *sc, const nvlist_t *nvl)
 		}
 	}
 	if (need_insert) {
-		if ((peer->p_remote = noise_remote_alloc(sc->sc_local, peer,
-		    pub_key, preshared_key)) == NULL)
-			goto out;
+		noise_remote_enable(peer->p_remote);
 		TAILQ_INSERT_TAIL(&sc->sc_peers, peer, p_entry);
 		sc->sc_peers_num++;
 		if (sc->sc_ifp->if_link_state == LINK_STATE_UP)
diff --git a/src/wg_noise.c b/src/wg_noise.c
index c92e44a..a321fc9 100644
--- a/src/wg_noise.c
+++ b/src/wg_noise.c
@@ -81,6 +81,7 @@ struct noise_remote {
 	struct noise_index		 r_index;
 
 	CK_LIST_ENTRY(noise_remote) 	 r_entry;
+	int				 r_entry_valid;
 	uint8_t				 r_public[NOISE_PUBLIC_KEY_LEN];
 
 	struct rwlock			 r_handshake_lock;
@@ -277,16 +278,15 @@ noise_precompute_ss(struct noise_local *l, struct noise_remote *r)
 /* 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])
+    const uint8_t public[NOISE_PUBLIC_KEY_LEN])
 {
 	struct noise_remote *r;
-	uint64_t idx;
 
 	if ((r = malloc(sizeof(*r), M_NOISE, M_NOWAIT)) == NULL)
 		return (NULL);
 
 	r->r_index.i_is_keypair = 0;
+	r->r_entry_valid = 0;
 
 	memcpy(r->r_public, public, NOISE_PUBLIC_KEY_LEN);
 
@@ -297,7 +297,7 @@ noise_remote_alloc(struct noise_local *l, void *arg,
 	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);
+	bzero(r->r_psk, sizeof(r->r_psk));
 	noise_precompute_ss(l, r);
 
 	refcount_init(&r->r_refcnt, 1);
@@ -309,21 +309,39 @@ noise_remote_alloc(struct noise_local *l, void *arg,
 
 	bzero(&r->r_smr, sizeof(r->r_smr));
 
+	return (r);
+}
+
+void
+noise_remote_enable(struct noise_remote *r)
+{
+	struct noise_local *l = r->r_local;
+	uint64_t idx;
+
 	/* Insert to hashtable */
-	idx = siphash24(&l->l_hash_key, public, NOISE_PUBLIC_KEY_LEN) & HT_REMOTE_MASK;
+	idx = siphash24(&l->l_hash_key, r->r_public, NOISE_PUBLIC_KEY_LEN) & HT_REMOTE_MASK;
 
 	rw_wlock(&l->l_remote_lock);
-	if (l->l_remote_num < MAX_REMOTE_PER_LOCAL) {
+	if (!r->r_entry_valid && l->l_remote_num < MAX_REMOTE_PER_LOCAL) {
+		r->r_entry_valid = 1;
 		l->l_remote_num++;
 		CK_LIST_INSERT_HEAD(&l->l_remote_hash[idx], r, r_entry);
-	} else {
-		free(r, M_NOISE);
-		noise_local_put(l);
-		r = NULL;
 	}
 	rw_wunlock(&l->l_remote_lock);
+}
 
-	return (r);
+void
+noise_remote_disable(struct noise_remote *r)
+{
+	struct noise_local *l = r->r_local;
+	/* remove from hashtable */
+	rw_wlock(&l->l_remote_lock);
+	if (r->r_entry_valid) {
+		r->r_entry_valid = 1;
+		CK_LIST_REMOVE(r, r_entry);
+		l->l_remote_num--;
+	};
+	rw_wunlock(&l->l_remote_lock);
 }
 
 struct noise_remote *
@@ -449,15 +467,8 @@ noise_remote_put(struct noise_remote *r)
 void
 noise_remote_free(struct noise_remote *r, void (*cleanup)(struct noise_remote *))
 {
-	struct noise_local *l = r->r_local;
-
 	r->r_cleanup = cleanup;
-
-	/* remove from hashtable */
-	rw_wlock(&l->l_remote_lock);
-	CK_LIST_REMOVE(r, r_entry);
-	l->l_remote_num--;
-	rw_wunlock(&l->l_remote_lock);
+	noise_remote_disable(r);
 
 	/* now clear all keypairs and handshakes, then put this reference */
 	noise_remote_handshake_clear(r);
diff --git a/src/wg_noise.h b/src/wg_noise.h
index d70d0be..893f62c 100644
--- a/src/wg_noise.h
+++ b/src/wg_noise.h
@@ -45,8 +45,9 @@ int	noise_local_keys(struct noise_local *,
 /* 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]);
+	    const uint8_t[NOISE_PUBLIC_KEY_LEN]);
+void	noise_remote_enable(struct noise_remote *);
+void	noise_remote_disable(struct noise_remote *);
 struct noise_remote *
 	noise_remote_lookup(struct noise_local *, const uint8_t[NOISE_PUBLIC_KEY_LEN]);
 struct noise_remote *