 @@ -1,14 +1,14 @@
-/*	$NetBSD: if_wg.c,v 1.25 2020/08/27 02:52:33 riastradh Exp $	*/
+/*	$NetBSD: if_wg.c,v 1.26 2020/08/27 02:53:47 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.25 2020/08/27 02:52:33 riastradh Exp $");
+__KERNEL_RCSID(0, "$NetBSD: if_wg.c,v 1.26 2020/08/27 02:53:47 riastradh Exp $");
 #ifdef _KERNEL_OPT
 #include "opt_inet.h"
 #endif
 #include <sys/param.h>
 #include <sys/systm.h>
 #include <sys/kernel.h>
 #include <sys/mbuf.h>
 #include <sys/socket.h>
 #include <sys/sockio.h>
 #include <sys/errno.h>
 #include <sys/ioctl.h>
 @@ -2333,31 +2333,27 @@ wg_handle_msg_data(struct wg_softc *wg,
     const struct sockaddr *src)
+{
 	struct wg_msg_data *wgmd;
 	char *encrypted_buf = NULL, *decrypted_buf;
 	size_t encrypted_len, decrypted_len;
 	struct wg_session *wgs;
 	struct wg_peer *wgp;
 	size_t mlen;
 	struct psref psref;
 	int error, af;
 	bool success, free_encrypted_buf = false, ok;
 	struct mbuf *n;
-	if (m->m_len < sizeof(struct wg_msg_data)) {
+	KASSERT(m->m_len >= sizeof(struct wg_msg_data));
 		m = m_pullup(m, sizeof(struct wg_msg_data));
 		if (m == NULL)
 			return;
 	wgmd = mtod(m, struct wg_msg_data *);
 	KASSERT(wgmd->wgmd_type == WG_MSG_TYPE_DATA);
 	WG_TRACE("data");
 	wgs = wg_lookup_session_by_index(wg, wgmd->wgmd_receiver, &psref);
 	if (wgs == NULL) {
 		WG_TRACE("No session found");
 		m_freem(m);
 		return;
+	}
 	wgp = wgs->wgs_peer;
 @@ -2563,93 +2559,112 @@ wg_handle_msg_cookie(struct wg_softc *wg
+	}
 	/*
 	 * [W] 6.6: Interaction with Cookie Reply System
 	 * "it should simply store the decrypted cookie value from the cookie
 	 *  reply message, and wait for the expiration of the REKEY-TIMEOUT
 	 *  timer for retrying a handshake initiation message."
 	 */
 	wgp->wgp_latest_cookie_time = time_uptime;
 	memcpy(wgp->wgp_latest_cookie, cookie, sizeof(wgp->wgp_latest_cookie));
 out:
 	wg_put_session(wgs, &psref);
+}
-static bool
+static struct mbuf *
-wg_validate_msg_length(struct wg_softc *wg, const struct mbuf *m)
+wg_validate_msg_header(struct wg_softc *wg, struct mbuf *m)
+{
-	struct wg_msg *wgm;
+	struct wg_msg wgm;
-	size_t mlen;
+	size_t mbuflen;
 	size_t msglen;
-	mlen = m_length(m);
+	/*
-	if (__predict_false(mlen < sizeof(struct wg_msg)))
+	 * Get the mbuf chain length.  It is already guaranteed, by
-		return false;
+	 * wg_overudp_cb, to be large enough for a struct wg_msg.
 	 */
 	mbuflen = m_length(m);
 	KASSERT(mbuflen >= sizeof(struct wg_msg));
-	wgm = mtod(m, struct wg_msg *);
+	/*
-	switch (wgm->wgm_type) {
+	 * Copy the message header (32-bit message type) out -- we'll
 	 * worry about contiguity and alignment later.
 	 */
 	m_copydata(m, 0, sizeof(wgm), &wgm);
 	switch (wgm.wgm_type) {
 	case WG_MSG_TYPE_INIT:
-		if (__predict_true(mlen >= sizeof(struct wg_msg_init)))
+		msglen = sizeof(struct wg_msg_init);
 			return true;
 		break;
 	case WG_MSG_TYPE_RESP:
-		if (__predict_true(mlen >= sizeof(struct wg_msg_resp)))
+		msglen = sizeof(struct wg_msg_resp);
 			return true;
 		break;
 	case WG_MSG_TYPE_COOKIE:
-		if (__predict_true(mlen >= sizeof(struct wg_msg_cookie)))
+		msglen = sizeof(struct wg_msg_cookie);
 			return true;
 		break;
 	case WG_MSG_TYPE_DATA:
-		if (__predict_true(mlen >= sizeof(struct wg_msg_data)))
+		msglen = sizeof(struct wg_msg_data);
 			return true;
 		break;
 	default:
 		WG_LOG_RATECHECK(&wg->wg_ppsratecheck, LOG_DEBUG,
-		    "Unexpected msg type: %u\n", wgm->wgm_type);
+		    "Unexpected msg type: %u\n", wgm.wgm_type);
-		return false;
+		goto error;
+	}
 	WG_DLOG("Invalid msg size: mlen=%lu type=%u\n", mlen, wgm->wgm_type);
-	return false;
+	/* Verify the mbuf chain is long enough for this type of message.  */
 	if (__predict_false(mbuflen < msglen)) {
 		WG_DLOG("Invalid msg size: mbuflen=%lu type=%u\n", mbuflen,
 		    wgm.wgm_type);
 		goto error;
+	}
 	/* Make the message header contiguous if necessary.  */
 	if (__predict_false(m->m_len < msglen)) {
 		m = m_pullup(m, msglen);
 		if (m == NULL)
 			return NULL;
+	}
 	return m;
 error:
 	m_freem(m);
 	return NULL;
+}
 static void
 wg_handle_packet(struct wg_softc *wg, struct mbuf *m,
     const struct sockaddr *src)
+{
 	struct wg_msg *wgm;
 	bool valid;
-	valid = wg_validate_msg_length(wg, m);
+	m = wg_validate_msg_header(wg, m);
-	if (!valid) {
+	if (__predict_false(m == NULL))
 		m_freem(m);
 		return;
 	KASSERT(m->m_len >= sizeof(struct wg_msg));
 	wgm = mtod(m, struct wg_msg *);
 	switch (wgm->wgm_type) {
 	case WG_MSG_TYPE_INIT:
 		wg_handle_msg_init(wg, (struct wg_msg_init *)wgm, src);
 		break;
 	case WG_MSG_TYPE_RESP:
 		wg_handle_msg_resp(wg, (struct wg_msg_resp *)wgm, src);
 		break;
 	case WG_MSG_TYPE_COOKIE:
 		wg_handle_msg_cookie(wg, (struct wg_msg_cookie *)wgm);
 		break;
 	case WG_MSG_TYPE_DATA:
 		wg_handle_msg_data(wg, m, src);
 		break;
 	default:
-		/* wg_validate_msg_length should already reject this case */
+		/* wg_validate_msg_header should already reject this case */
 		break;
+	}
+}
 static void
 wg_receive_packets(struct wg_softc *wg, const int af)
+{
 	for (;;) {
 		int error, flags;
 		struct socket *so;
 		struct mbuf *m = NULL;
 		struct uio dummy_uio;