 @@ -1,14 +1,14 @@
-/*	$NetBSD: if_wg.c,v 1.42 2020/08/31 20:26:21 riastradh Exp $	*/
+/*	$NetBSD: if_wg.c,v 1.43 2020/08/31 20:26:46 riastradh Exp $	*/
 /*
  * Copyright (C) Ryota Ozaki <ozaki.ryota@gmail.com>
  * All rights reserved.
+ *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions
  * are met:
  * 1. Redistributions of source code must retain the above copyright
  *    notice, this list of conditions and the following disclaimer.
  * 2. Redistributions in binary form must reproduce the above copyright
  *    notice, this list of conditions and the following disclaimer in the
  *    documentation and/or other materials provided with the distribution.
 @@ -31,27 +31,27 @@
 /*
  * This network interface aims to implement the WireGuard protocol.
  * The implementation is based on the paper of WireGuard as of
  * 2018-06-30 [1].  The paper is referred in the source code with label
  * [W].  Also the specification of the Noise protocol framework as of
  * 2018-07-11 [2] is referred with label [N].
+ *
  * [1] https://www.wireguard.com/papers/wireguard.pdf
  * [2] http://noiseprotocol.org/noise.pdf
  */
 #include <sys/cdefs.h>
-__KERNEL_RCSID(0, "$NetBSD: if_wg.c,v 1.42 2020/08/31 20:26:21 riastradh Exp $");
+__KERNEL_RCSID(0, "$NetBSD: if_wg.c,v 1.43 2020/08/31 20:26:46 riastradh Exp $");
 #ifdef _KERNEL_OPT
 #include "opt_inet.h"
 #endif
 #include <sys/param.h>
 #include <sys/types.h>
 #include <sys/atomic.h>
 #include <sys/callout.h>
 #include <sys/cprng.h>
 #include <sys/cpu.h>
 #include <sys/device.h>
 @@ -1319,27 +1319,26 @@ wg_fill_msg_init(struct wg_softc *wg, st
 	WG_DLOG("%s: sender=%x\n", __func__, wgs->wgs_sender_index);
+}
 static void
 wg_handle_msg_init(struct wg_softc *wg, const struct wg_msg_init *wgmi,
     const struct sockaddr *src)
+{
 	uint8_t ckey[WG_CHAINING_KEY_LEN]; /* [W] 5.4.2: Ci */
 	uint8_t hash[WG_HASH_LEN]; /* [W] 5.4.2: Hi */
 	uint8_t cipher_key[WG_CIPHER_KEY_LEN];
 	uint8_t peer_pubkey[WG_STATIC_KEY_LEN];
 	struct wg_peer *wgp;
 	struct wg_session *wgs;
 	bool reset_state_on_error = false;
 	int error, ret;
 	struct psref psref_peer;
 	struct psref psref_session;
 	uint8_t mac1[WG_MAC_LEN];
 	WG_TRACE("init msg received");
 	/*
 	 * [W] 5.4.2: First Message: Initiator to Responder
 	 * "When the responder receives this message, it does the same
 	 *  operations so that its final state variables are identical,
 	 *  replacing the operands of the DH function to produce equivalent
 	 *  values."
 @@ -1399,27 +1398,26 @@ wg_handle_msg_init(struct wg_softc *wg,
 		wg_clear_states(wgs);
 		wgs->wgs_state = WGS_STATE_UNKNOWN;
+	}
 	if (wgs->wgs_state == WGS_STATE_INIT_ACTIVE) {
 		WG_TRACE("Sesssion already initializing, ignoring the message");
 		mutex_exit(wgs->wgs_lock);
 		goto out_wgp;
+	}
 	if (wgs->wgs_state == WGS_STATE_INIT_PASSIVE) {
 		WG_TRACE("Sesssion already initializing, destroying old states");
 		wg_clear_states(wgs);
+	}
 	wgs->wgs_state = WGS_STATE_INIT_PASSIVE;
 	reset_state_on_error = true;
 	wg_get_session(wgs, &psref_session);
 	mutex_exit(wgs->wgs_lock);
 	wg_algo_mac_mac1(mac1, sizeof(mac1),
 	    wg->wg_pubkey, sizeof(wg->wg_pubkey),
 	    (const uint8_t *)wgmi, offsetof(struct wg_msg_init, wgmi_mac1));
 	/*
 	 * [W] 5.3: Denial of Service Mitigation & Cookies
 	 * "the responder, ..., must always reject messages with an invalid
 	 *  msg.mac1"
 	 */
 	if (!consttime_memequal(mac1, wgmi->wgmi_mac1, sizeof(mac1))) {
 @@ -1499,32 +1497,30 @@ wg_handle_msg_init(struct wg_softc *wg,
 	wg_update_endpoint_if_necessary(wgp, src);
 	(void)wg_send_handshake_msg_resp(wg, wgp, wgmi);
 	wg_calculate_keys(wgs, false);
 	wg_clear_states(wgs);
 	wg_put_session(wgs, &psref_session);
 	wg_put_peer(wgp, &psref_peer);
 	return;
 out:
-	if (reset_state_on_error) {
+	mutex_enter(wgs->wgs_lock);
-		mutex_enter(wgs->wgs_lock);
+	KASSERT(wgs->wgs_state == WGS_STATE_INIT_PASSIVE);
-		KASSERT(wgs->wgs_state == WGS_STATE_INIT_PASSIVE);
+	wgs->wgs_state = WGS_STATE_UNKNOWN;
-		wgs->wgs_state = WGS_STATE_UNKNOWN;
+	mutex_exit(wgs->wgs_lock);
 		mutex_exit(wgs->wgs_lock);
 	wg_put_session(wgs, &psref_session);
 out_wgp:
 	wg_put_peer(wgp, &psref_peer);
+}
 static void
 wg_schedule_handshake_timeout_timer(struct wg_peer *wgp)
+{
 	mutex_enter(wgp->wgp_lock);
 	if (__predict_true(wgp->wgp_state != WGP_STATE_DESTROYING)) {
 		callout_schedule(&wgp->wgp_handshake_timeout_timer,
 		    MIN(wg_rekey_timeout, INT_MAX/hz) * hz);