wg_noise: check keypair recvwith after nonce

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

3 files changed, 31 insertions, 38 deletions

diff --git a/src/if_wg.c b/src/if_wg.c
index a64403b..7b74348 100644
--- a/src/if_wg.c
+++ b/src/if_wg.c
@@ -1459,7 +1459,7 @@ wg_decrypt(struct wg_softc *sc, struct wg_packet *pkt)
 	struct mbuf		*m;
 	struct ip		*ip;
 	struct ip6_hdr		*ip6;
-	int			 res, len;
+	int			 len;
 
 	peer = noise_keypair_remote_arg(pkt->p_keypair);
 	m = pkt->p_mbuf;
@@ -1470,15 +1470,8 @@ wg_decrypt(struct wg_softc *sc, struct wg_packet *pkt)
 	m_adj(m, sizeof(struct wg_pkt_data));
 
 	pkt->p_nonce = le64toh(data.nonce);
-	res = noise_keypair_decrypt(pkt->p_keypair, pkt->p_nonce, m);
-
-	if (__predict_false(res == EINVAL)) {
+	if (noise_keypair_decrypt(pkt->p_keypair, pkt->p_nonce, m) != 0)
 		goto error;
-	} else if (__predict_false(res == ECONNRESET)) {
-		wg_timers_event_handshake_complete(peer);
-	} else if (__predict_false(res != 0)) {
-		panic("unexpected response: %d\n", res);
-	}
 
 	/* A packet with length 0 is a keepalive packet */
 	if (__predict_false(m->m_pkthdr.len == 0)) {
@@ -1631,6 +1624,9 @@ wg_deliver_in(struct wg_peer *peer)
 			if (noise_keypair_nonce_check(pkt->p_keypair, pkt->p_nonce) != 0)
 				goto error;
 
+			if (noise_keypair_received_with(pkt->p_keypair) == ECONNRESET)
+				wg_timers_event_handshake_complete(peer);
+
 			wg_timers_event_any_authenticated_packet_received(peer);
 			wg_timers_event_any_authenticated_packet_traversal(peer);
 			wg_peer_set_endpoint(peer, &pkt->p_endpoint);
diff --git a/src/wg_noise.c b/src/wg_noise.c
index f12c557..4ecc368 100644
--- a/src/wg_noise.c
+++ b/src/wg_noise.c
@@ -130,7 +130,6 @@ 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 *);
 
@@ -608,31 +607,6 @@ noise_add_new_keypair(struct noise_local *l, struct noise_remote *r,
 }
 
 static int
-noise_received_with(struct noise_keypair *kp)
-{
-	struct noise_keypair *old;
-	struct noise_remote *r = kp->kp_remote;
-
-	if (kp != epoch_ptr_read(&r->r_next))
-		return (0);
-
-	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;
@@ -713,6 +687,31 @@ noise_keypair_ref(struct noise_keypair *kp)
 	return (kp);
 }
 
+int
+noise_keypair_received_with(struct noise_keypair *kp)
+{
+	struct noise_keypair *old;
+	struct noise_remote *r = kp->kp_remote;
+
+	if (kp != epoch_ptr_read(&r->r_next))
+		return (0);
+
+	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 void
 noise_keypair_smr_free(struct epoch_context *smr)
 {
@@ -876,9 +875,6 @@ noise_keypair_decrypt(struct noise_keypair *kp, uint64_t nonce, struct mbuf *m)
 	if (chacha20poly1305_decrypt_mbuf(m, nonce, kp->kp_recv) == 0)
 		return (EINVAL);
 
-	if (noise_received_with(kp) != 0)
-		return (ECONNRESET);
-
 	return (0);
 }
 
diff --git a/src/wg_noise.h b/src/wg_noise.h
index a6de34b..d70d0be 100644
--- a/src/wg_noise.h
+++ b/src/wg_noise.h
@@ -75,6 +75,7 @@ struct noise_keypair *
 	noise_keypair_current(struct noise_remote *);
 struct noise_keypair *
 	noise_keypair_ref(struct noise_keypair *);
+int	noise_keypair_received_with(struct noise_keypair *);
 void	noise_keypair_put(struct noise_keypair *);
 
 struct noise_remote *