5 files changed, 100 insertions, 42 deletions

diff --git a/sys/kern/kern_descrip.c b/sys/kern/kern_descrip.c
index 18c70c36a06c..80229a75ed2b 100644
--- a/sys/kern/kern_descrip.c
+++ b/sys/kern/kern_descrip.c
@@ -78,6 +78,8 @@ __FBSDID("$FreeBSD$");
 #include <sys/ktrace.h>
 #endif
 
+#include <net/vnet.h>
+
 #include <security/audit/audit.h>
 
 #include <vm/uma.h>
@@ -2339,7 +2341,9 @@ fputsock(struct socket *so)
 
 	ACCEPT_LOCK();
 	SOCK_LOCK(so);
+	CURVNET_SET(so->so_vnet);
 	sorele(so);
+	CURVNET_RESTORE();
 }
 
 /*
diff --git a/sys/kern/sys_socket.c b/sys/kern/sys_socket.c
index 564e8c6b7488..4cbae7440771 100644
--- a/sys/kern/sys_socket.c
+++ b/sys/kern/sys_socket.c
@@ -80,9 +80,7 @@ soo_read(struct file *fp, struct uio *uio, struct ucred *active_cred,
 	if (error)
 		return (error);
 #endif
-	CURVNET_SET(so->so_vnet);
 	error = soreceive(so, 0, uio, 0, 0, 0);
-	CURVNET_RESTORE();
 	return (error);
 }
 
@@ -123,7 +121,6 @@ soo_ioctl(struct file *fp, u_long cmd, void *data, struct ucred *active_cred,
 	struct socket *so = fp->f_data;
 	int error = 0;
 
-	CURVNET_SET(so->so_vnet);
 	switch (cmd) {
 	case FIONBIO:
 		SOCK_LOCK(so);
@@ -210,14 +207,18 @@ soo_ioctl(struct file *fp, u_long cmd, void *data, struct ucred *active_cred,
 		 */
 		if (IOCGROUP(cmd) == 'i')
 			error = ifioctl(so, cmd, data, td);
-		else if (IOCGROUP(cmd) == 'r')
+		else if (IOCGROUP(cmd) == 'r') {
+			CURVNET_SET(so->so_vnet);
 			error = rtioctl_fib(cmd, data, so->so_fibnum);
-		else
+			CURVNET_RESTORE();
+		} else {
+			CURVNET_SET(so->so_vnet);
 			error = ((*so->so_proto->pr_usrreqs->pru_control)
 			    (so, cmd, data, 0, td));
+			CURVNET_RESTORE();
+		}
 		break;
 	}
-	CURVNET_RESTORE();
 	return (error);
 }
 
diff --git a/sys/kern/uipc_socket.c b/sys/kern/uipc_socket.c
index 83a1ae2c8c6a..8f876e4b2d24 100644
--- a/sys/kern/uipc_socket.c
+++ b/sys/kern/uipc_socket.c
@@ -92,6 +92,12 @@
  * from a listen queue to a file descriptor, in order to prevent garbage
  * collection of the socket at an untimely moment.  For a number of reasons,
  * these interfaces are not preferred, and should be avoided.
+ * 
+ * NOTE: With regard to VNETs the general rule is that callers do not set
+ * curvnet. Exceptions to this rule include soabort(), sodisconnect(),
+ * sofree() (and with that sorele(), sotryfree()), as well as sonewconn()
+ * and sorflush(), which are usually called from a pre-set VNET context.
+ * sopoll() currently does not need a VNET context to be set.
  */
 
 #include <sys/cdefs.h>
@@ -166,6 +172,10 @@ int	maxsockets;
 MALLOC_DEFINE(M_SONAME, "soname", "socket name");
 MALLOC_DEFINE(M_PCB, "pcb", "protocol control block");
 
+#define	VNET_SO_ASSERT(so)						\
+	VNET_ASSERT(curvnet != NULL,					\
+	    ("%s:%d curvnet is NULL, so=%p", __func__, __LINE__, (so)));
+
 static int somaxconn = SOMAXCONN;
 static int sysctl_somaxconn(SYSCTL_HANDLER_ARGS);
 /* XXX: we dont have SYSCTL_USHORT */
@@ -286,6 +296,8 @@ soalloc(struct vnet *vnet)
 	so->so_gencnt = ++so_gencnt;
 	++numopensockets;
 #ifdef VIMAGE
+	VNET_ASSERT(vnet != NULL, ("%s:%d vnet is NULL, so=%p",
+	    __func__, __LINE__, so));
 	vnet->vnet_sockcnt++;
 	so->so_vnet = vnet;
 #endif
@@ -309,6 +321,8 @@ sodealloc(struct socket *so)
 	so->so_gencnt = ++so_gencnt;
 	--numopensockets;	/* Could be below, but faster here. */
 #ifdef VIMAGE
+	VNET_ASSERT(so->so_vnet != NULL, ("%s:%d so_vnet is NULL, so=%p",
+	    __func__, __LINE__, so));
 	so->so_vnet->vnet_sockcnt--;
 #endif
 	mtx_unlock(&so_global_mtx);
@@ -448,6 +462,7 @@ sonewconn(struct socket *head, int connstatus)
 #endif
 	knlist_init_mtx(&so->so_rcv.sb_sel.si_note, SOCKBUF_MTX(&so->so_rcv));
 	knlist_init_mtx(&so->so_snd.sb_sel.si_note, SOCKBUF_MTX(&so->so_snd));
+	VNET_SO_ASSERT(head);
 	if (soreserve(so, head->so_snd.sb_hiwat, head->so_rcv.sb_hiwat) ||
 	    (*so->so_proto->pr_usrreqs->pru_attach)(so, 0, NULL)) {
 		sodealloc(so);
@@ -522,8 +537,12 @@ sobind(struct socket *so, struct sockaddr *nam, struct thread *td)
 int
 solisten(struct socket *so, int backlog, struct thread *td)
 {
+	int error;
 
-	return ((*so->so_proto->pr_usrreqs->pru_listen)(so, backlog, td));
+	CURVNET_SET(so->so_vnet);
+	error = (*so->so_proto->pr_usrreqs->pru_listen)(so, backlog, td);
+	CURVNET_RESTORE();
+	return error;
 }
 
 int
@@ -611,6 +630,7 @@ sofree(struct socket *so)
 	SOCK_UNLOCK(so);
 	ACCEPT_UNLOCK();
 
+	VNET_SO_ASSERT(so);
 	if (pr->pr_flags & PR_RIGHTS && pr->pr_domain->dom_dispose != NULL)
 		(*pr->pr_domain->dom_dispose)(so->so_rcv.sb_mb);
 	if (pr->pr_usrreqs->pru_detach != NULL)
@@ -740,6 +760,7 @@ soabort(struct socket *so)
 	KASSERT(so->so_state & SS_NOFDREF, ("soabort: !SS_NOFDREF"));
 	KASSERT((so->so_state & SQ_COMP) == 0, ("soabort: SQ_COMP"));
 	KASSERT((so->so_state & SQ_INCOMP) == 0, ("soabort: SQ_INCOMP"));
+	VNET_SO_ASSERT(so);
 
 	if (so->so_proto->pr_usrreqs->pru_abort != NULL)
 		(*so->so_proto->pr_usrreqs->pru_abort)(so);
@@ -757,7 +778,10 @@ soaccept(struct socket *so, struct sockaddr **nam)
 	KASSERT((so->so_state & SS_NOFDREF) != 0, ("soaccept: !NOFDREF"));
 	so->so_state &= ~SS_NOFDREF;
 	SOCK_UNLOCK(so);
+
+	CURVNET_SET(so->so_vnet);
 	error = (*so->so_proto->pr_usrreqs->pru_accept)(so, nam);
+	CURVNET_RESTORE();
 	return (error);
 }
 
@@ -795,8 +819,12 @@ soconnect(struct socket *so, struct sockaddr *nam, struct thread *td)
 int
 soconnect2(struct socket *so1, struct socket *so2)
 {
+	int error;
 
-	return ((*so1->so_proto->pr_usrreqs->pru_connect2)(so1, so2));
+	CURVNET_SET(so1->so_vnet);
+	error = (*so1->so_proto->pr_usrreqs->pru_connect2)(so1, so2);
+	CURVNET_RESTORE();
+	return (error);
 }
 
 int
@@ -808,6 +836,7 @@ sodisconnect(struct socket *so)
 		return (ENOTCONN);
 	if (so->so_state & SS_ISDISCONNECTING)
 		return (EALREADY);
+	VNET_SO_ASSERT(so);
 	error = (*so->so_proto->pr_usrreqs->pru_disconnect)(so);
 	return (error);
 }
@@ -1071,6 +1100,7 @@ sosend_dgram(struct socket *so, struct sockaddr *addr, struct uio *uio,
 	 * there are probably other places that this also happens.  We must
 	 * rethink this.
 	 */
+	VNET_SO_ASSERT(so);
 	error = (*so->so_proto->pr_usrreqs->pru_send)(so,
 	    (flags & MSG_OOB) ? PRUS_OOB :
 	/*
@@ -1258,6 +1288,7 @@ restart:
 			 * places that this also happens.  We must rethink
 			 * this.
 			 */
+			VNET_SO_ASSERT(so);
 			error = (*so->so_proto->pr_usrreqs->pru_send)(so,
 			    (flags & MSG_OOB) ? PRUS_OOB :
 			/*
@@ -1324,6 +1355,7 @@ soreceive_rcvoob(struct socket *so, struct uio *uio, int flags)
 	int error;
 
 	KASSERT(flags & MSG_OOB, ("soreceive_rcvoob: (flags & MSG_OOB) == 0"));
+	VNET_SO_ASSERT(so);
 
 	m = m_get(M_WAIT, MT_DATA);
 	error = (*pr->pr_usrreqs->pru_rcvoob)(so, m, flags & MSG_PEEK);
@@ -1432,8 +1464,10 @@ soreceive_generic(struct socket *so, struct sockaddr **psa, struct uio *uio,
 	if (mp != NULL)
 		*mp = NULL;
 	if ((pr->pr_flags & PR_WANTRCVD) && (so->so_state & SS_ISCONFIRMING)
-	    && uio->uio_resid)
+	    && uio->uio_resid) {
+		VNET_SO_ASSERT(so);
 		(*pr->pr_usrreqs->pru_rcvd)(so, 0);
+	}
 
 	error = sblock(&so->so_rcv, SBLOCKWAIT(flags));
 	if (error)
@@ -1580,6 +1614,7 @@ dontblock:
 			cm->m_next = NULL;
 			if (pr->pr_domain->dom_externalize != NULL) {
 				SOCKBUF_UNLOCK(&so->so_rcv);
+				VNET_SO_ASSERT(so);
 				error = (*pr->pr_domain->dom_externalize)
 				    (cm, controlp);
 				SOCKBUF_LOCK(&so->so_rcv);
@@ -1795,6 +1830,7 @@ dontblock:
 			 */
 			if (pr->pr_flags & PR_WANTRCVD) {
 				SOCKBUF_UNLOCK(&so->so_rcv);
+				VNET_SO_ASSERT(so);
 				(*pr->pr_usrreqs->pru_rcvd)(so, flags);
 				SOCKBUF_LOCK(&so->so_rcv);
 			}
@@ -1841,6 +1877,7 @@ dontblock:
 		if (!(flags & MSG_SOCALLBCK) &&
 		    (pr->pr_flags & PR_WANTRCVD)) {
 			SOCKBUF_UNLOCK(&so->so_rcv);
+			VNET_SO_ASSERT(so);
 			(*pr->pr_usrreqs->pru_rcvd)(so, flags);
 			SOCKBUF_LOCK(&so->so_rcv);
 		}
@@ -2036,6 +2073,7 @@ deliver:
 		    (((flags & MSG_WAITALL) && uio->uio_resid > 0) ||
 		     !(flags & MSG_SOCALLBCK))) {
 			SOCKBUF_UNLOCK(sb);
+			VNET_SO_ASSERT(so);
 			(*so->so_proto->pr_usrreqs->pru_rcvd)(so, flags);
 			SOCKBUF_LOCK(sb);
 		}
@@ -2245,9 +2283,13 @@ int
 soreceive(struct socket *so, struct sockaddr **psa, struct uio *uio,
     struct mbuf **mp0, struct mbuf **controlp, int *flagsp)
 {
+	int error;
 
-	return (so->so_proto->pr_usrreqs->pru_soreceive(so, psa, uio, mp0,
+	CURVNET_SET(so->so_vnet);
+	error = (so->so_proto->pr_usrreqs->pru_soreceive(so, psa, uio, mp0,
 	    controlp, flagsp));
+	CURVNET_RESTORE();
+	return (error);
 }
 
 int
@@ -2258,17 +2300,19 @@ soshutdown(struct socket *so, int how)
 
 	if (!(how == SHUT_RD || how == SHUT_WR || how == SHUT_RDWR))
 		return (EINVAL);
+
+	CURVNET_SET(so->so_vnet);
 	if (pr->pr_usrreqs->pru_flush != NULL) {
 	        (*pr->pr_usrreqs->pru_flush)(so, how);
 	}
 	if (how != SHUT_WR)
 		sorflush(so);
 	if (how != SHUT_RD) {
-		CURVNET_SET(so->so_vnet);
 		error = (*pr->pr_usrreqs->pru_shutdown)(so);
 		CURVNET_RESTORE();
 		return (error);
 	}
+	CURVNET_RESTORE();
 	return (0);
 }
 
@@ -2279,6 +2323,8 @@ sorflush(struct socket *so)
 	struct protosw *pr = so->so_proto;
 	struct sockbuf asb;
 
+	VNET_SO_ASSERT(so);
+
 	/*
 	 * In order to avoid calling dom_dispose with the socket buffer mutex
 	 * held, and in order to generally avoid holding the lock for a long
@@ -2292,7 +2338,6 @@ sorflush(struct socket *so)
 	 * socket buffer.  Don't let our acquire be interrupted by a signal
 	 * despite any existing socket disposition on interruptable waiting.
 	 */
-	CURVNET_SET(so->so_vnet);
 	socantrcvmore(so);
 	(void) sblock(sb, SBL_WAIT | SBL_NOINTR);
 
@@ -2316,7 +2361,6 @@ sorflush(struct socket *so)
 	if (pr->pr_flags & PR_RIGHTS && pr->pr_domain->dom_dispose != NULL)
 		(*pr->pr_domain->dom_dispose)(asb.sb_mb);
 	sbrelease_internal(&asb, so);
-	CURVNET_RESTORE();
 }
 
 /*
@@ -2380,11 +2424,14 @@ sosetopt(struct socket *so, struct sockopt *sopt)
 	struct mac extmac;
 #endif
 
+	CURVNET_SET(so->so_vnet);
 	error = 0;
 	if (sopt->sopt_level != SOL_SOCKET) {
-		if (so->so_proto && so->so_proto->pr_ctloutput)
-			return ((*so->so_proto->pr_ctloutput)
-				  (so, sopt));
+		if (so->so_proto && so->so_proto->pr_ctloutput) {
+			error = (*so->so_proto->pr_ctloutput)(so, sopt);
+			CURVNET_RESTORE();
+			return (error);
+		}
 		error = ENOPROTOOPT;
 	} else {
 		switch (sopt->sopt_name) {
@@ -2570,6 +2617,7 @@ sosetopt(struct socket *so, struct sockopt *sopt)
 		}
 	}
 bad:
+	CURVNET_RESTORE();
 	return (error);
 }
 
@@ -2613,13 +2661,15 @@ sogetopt(struct socket *so, struct sockopt *sopt)
 	struct mac extmac;
 #endif
 
+	CURVNET_SET(so->so_vnet);
 	error = 0;
 	if (sopt->sopt_level != SOL_SOCKET) {
-		if (so->so_proto && so->so_proto->pr_ctloutput) {
-			return ((*so->so_proto->pr_ctloutput)
-				  (so, sopt));
-		} else
-			return (ENOPROTOOPT);
+		if (so->so_proto && so->so_proto->pr_ctloutput)
+			error = (*so->so_proto->pr_ctloutput)(so, sopt);
+		else
+			error = ENOPROTOOPT;
+		CURVNET_RESTORE();
+		return (error);
 	} else {
 		switch (sopt->sopt_name) {
 #ifdef INET
@@ -2703,11 +2753,11 @@ integer:
 			error = sooptcopyin(sopt, &extmac, sizeof(extmac),
 			    sizeof(extmac));
 			if (error)
-				return (error);
+				goto bad;
 			error = mac_getsockopt_label(sopt->sopt_td->td_ucred,
 			    so, &extmac);
 			if (error)
-				return (error);
+				goto bad;
 			error = sooptcopyout(sopt, &extmac, sizeof extmac);
 #else
 			error = EOPNOTSUPP;
@@ -2719,11 +2769,11 @@ integer:
 			error = sooptcopyin(sopt, &extmac, sizeof(extmac),
 			    sizeof(extmac));
 			if (error)
-				return (error);
+				goto bad;
 			error = mac_getsockopt_peerlabel(
 			    sopt->sopt_td->td_ucred, so, &extmac);
 			if (error)
-				return (error);
+				goto bad;
 			error = sooptcopyout(sopt, &extmac, sizeof extmac);
 #else
 			error = EOPNOTSUPP;
@@ -2746,8 +2796,12 @@ integer:
 			error = ENOPROTOOPT;
 			break;
 		}
-		return (error);
 	}
+#ifdef MAC
+bad:
+#endif
+	CURVNET_RESTORE();
+	return (error);
 }
 
 /* XXX; prepare mbuf for (__FreeBSD__ < 3) routines. */
@@ -2881,6 +2935,10 @@ sopoll(struct socket *so, int events, struct ucred *active_cred,
     struct thread *td)
 {
 
+	/*
+	 * We do not need to set or assert curvnet as long as everyone uses
+	 * sopoll_generic().
+	 */
 	return (so->so_proto->pr_usrreqs->pru_sopoll(so, events, active_cred,
 	    td));
 }
diff --git a/sys/kern/uipc_syscalls.c b/sys/kern/uipc_syscalls.c
index 3270593f16ea..5ff469b8eb4c 100644
--- a/sys/kern/uipc_syscalls.c
+++ b/sys/kern/uipc_syscalls.c
@@ -260,14 +260,9 @@ listen(td, uap)
 		so = fp->f_data;
 #ifdef MAC
 		error = mac_socket_check_listen(td->td_ucred, so);
-		if (error == 0) {
+		if (error == 0)
 #endif
-			CURVNET_SET(so->so_vnet);
 			error = solisten(so, uap->backlog, td);
-			CURVNET_RESTORE();
-#ifdef MAC
-		}
-#endif
 		fdrop(fp, td);
 	}
 	return(error);
@@ -428,9 +423,7 @@ kern_accept(struct thread *td, int s, struct sockaddr **name,
 	tmp = fflag & FASYNC;
 	(void) fo_ioctl(nfp, FIOASYNC, &tmp, td->td_ucred, td);
 	sa = 0;
-	CURVNET_SET(so->so_vnet);
 	error = soaccept(so, &sa);
-	CURVNET_RESTORE();
 	if (error) {
 		/*
 		 * return a namelen of zero for older code which might
@@ -981,11 +974,9 @@ kern_recvit(td, s, mp, fromseg, controlp)
 		ktruio = cloneuio(&auio);
 #endif
 	len = auio.uio_resid;
-	CURVNET_SET(so->so_vnet);
 	error = soreceive(so, &fromsa, &auio, (struct mbuf **)0,
 	    (mp->msg_control || controlp) ? &control : (struct mbuf **)0,
 	    &mp->msg_flags);
-	CURVNET_RESTORE();
 	if (error) {
 		if (auio.uio_resid != (int)len && (error == ERESTART ||
 		    error == EINTR || error == EWOULDBLOCK))
@@ -1331,9 +1322,7 @@ kern_setsockopt(td, s, level, name, val, valseg, valsize)
 	error = getsock(td->td_proc->p_fd, s, &fp, NULL);
 	if (error == 0) {
 		so = fp->f_data;
-		CURVNET_SET(so->so_vnet);
 		error = sosetopt(so, &sopt);
-		CURVNET_RESTORE();
 		fdrop(fp, td);
 	}
 	return(error);
@@ -1412,9 +1401,7 @@ kern_getsockopt(td, s, level, name, val, valseg, valsize)
 	error = getsock(td->td_proc->p_fd, s, &fp, NULL);
 	if (error == 0) {
 		so = fp->f_data;
-		CURVNET_SET(so->so_vnet);
 		error = sogetopt(so, &sopt);
-		CURVNET_RESTORE();
 		*valsize = sopt.sopt_valsize;
 		fdrop(fp, td);
 	}
diff --git a/sys/kern/uipc_usrreq.c b/sys/kern/uipc_usrreq.c
index eb1fa7545453..551bab8e2557 100644
--- a/sys/kern/uipc_usrreq.c
+++ b/sys/kern/uipc_usrreq.c
@@ -1227,7 +1227,9 @@ unp_connect(struct socket *so, struct sockaddr *nam, struct thread *td)
 	}
 	if (so->so_proto->pr_flags & PR_CONNREQUIRED) {
 		if (so2->so_options & SO_ACCEPTCONN) {
+			CURVNET_SET(so2->so_vnet);
 			so3 = sonewconn(so2, 0);
+			CURVNET_RESTORE();
 		} else
 			so3 = NULL;
 		if (so3 == NULL) {
@@ -2107,8 +2109,14 @@ unp_gc(__unused void *arg, int pending)
 	 * struct files associated with these sockets but leave each socket
 	 * with one remaining ref.
 	 */
-	for (i = 0; i < total; i++)
-		sorflush(unref[i]->f_data);
+	for (i = 0; i < total; i++) {
+		struct socket *so;
+
+		so = unref[i]->f_data;
+		CURVNET_SET(so->so_vnet);
+		sorflush(so);
+		CURVNET_RESTORE();
+	}
 
 	/*
 	 * And finally release the sockets so they can be reclaimed.