Extend coverage of SOCK_LOCK(so) to include so_count, the socket

reference count:

- Assert SOCK_LOCK(so) macros that directly manipulate so_count:
  soref(), sorele().

- Assert SOCK_LOCK(so) in macros/functions that rely on the state of
  so_count: sofree(), sotryfree().

- Acquire SOCK_LOCK(so) before calling these functions or macros in
  various contexts in the stack, both at the socket and protocol
  layers.

- In some cases, perform soisdisconnected() before sotryfree(), as
  this could result in frobbing of a non-present socket if
  sotryfree() actually frees the socket.

- Note that sofree()/sotryfree() will release the socket lock even if
  they don't free the socket.

Submitted by:	sam
Sponsored by:	FreeBSD Foundation
Obtained from:	BSD/OS

sys/kern/kern_descrip.c

@@ -2024,7 +2024,9 @@ fgetsock(struct thread *td, int fd, struct socket **spp, u_int *fflagp)
 		*spp = fp->f_data;
 		if (fflagp)
 			*fflagp = fp->f_flag;
+		SOCK_LOCK(*spp);
 		soref(*spp);
+		SOCK_UNLOCK(*spp);
 	}
 	FILEDESC_UNLOCK(td->td_proc->p_fd);
 	return (error);
@@ -2039,6 +2041,7 @@ fputsock(struct socket *so)
 {
 
 	NET_ASSERT_GIANT();
+	SOCK_LOCK(so);
 	sorele(so);
 }
 

sys/kern/uipc_socket.c

@@ -201,9 +201,12 @@ socreate(dom, aso, type, proto, cred, td)
 #ifdef MAC
 	mac_create_socket(cred, so);
 #endif
+	SOCK_LOCK(so);
 	soref(so);
+	SOCK_UNLOCK(so);
 	error = (*prp->pr_usrreqs->pru_attach)(so, proto, td);
 	if (error) {
+		SOCK_LOCK(so);
 		so->so_state |= SS_NOFDREF;
 		sorele(so);
 		return (error);
@@ -292,10 +295,14 @@ sofree(so)
 	int s;
 
 	KASSERT(so->so_count == 0, ("socket %p so_count not 0", so));
+	SOCK_LOCK_ASSERT(so);
 
-	if (so->so_pcb != NULL || (so->so_state & SS_NOFDREF) == 0)
+	if (so->so_pcb != NULL || (so->so_state & SS_NOFDREF) == 0) {
+		SOCK_UNLOCK(so);
 		return;
+	}
 
+	SOCK_UNLOCK(so);
 	ACCEPT_LOCK();
 	head = so->so_head;
 	if (head != NULL) {
@@ -409,6 +416,7 @@ soclose(so)
 			error = error2;
 	}
 discard:
+	SOCK_LOCK(so);
 	if (so->so_state & SS_NOFDREF)
 		panic("soclose: NOFDREF");
 	so->so_state |= SS_NOFDREF;
@@ -428,6 +436,7 @@ soabort(so)
 
 	error = (*so->so_proto->pr_usrreqs->pru_abort)(so);
 	if (error) {
+		SOCK_LOCK(so);
 		sotryfree(so);	/* note: does not decrement the ref count */
 		return error;
 	}

sys/kern/uipc_syscalls.c

@@ -309,7 +309,14 @@ accept1(td, uap, compat)
 	KASSERT(!(so->so_qstate & SQ_INCOMP), ("accept1: so SQ_INCOMP"));
 	KASSERT(so->so_qstate & SQ_COMP, ("accept1: so not SQ_COMP"));
 
+        /*
+         * Before changing the flags on the socket, we have to bump the
+	 * reference count.  Otherwise, if the protocol calls sofree(),
+         * the socket will be released due to a zero refcount.
+         */
+	SOCK_LOCK(so);
 	soref(so);			/* file descriptor reference */
+	SOCK_UNLOCK(so);
 	
 	TAILQ_REMOVE(&head->so_comp, so, so_list);
 	head->so_qlen--;

sys/kern/uipc_usrreq.c

@@ -114,6 +114,7 @@ uipc_abort(struct socket *so)
 	UNP_LOCK();
 	unp_drop(unp, ECONNABORTED);
 	unp_detach(unp);	/* NB: unlocks */
+	SOCK_LOCK(so);
 	sotryfree(so);
 	return (0);
 }

sys/net/raw_cb.c

@@ -32,7 +32,9 @@
 
 #include <sys/param.h>
 #include <sys/domain.h>
+#include <sys/lock.h>
 #include <sys/malloc.h>
+#include <sys/mutex.h>
 #include <sys/protosw.h>
 #include <sys/socket.h>
 #include <sys/socketvar.h>
@@ -93,6 +95,7 @@ raw_detach(rp)
 {
 	struct socket *so = rp->rcb_socket;
 
+	SOCK_LOCK(so);
 	so->so_pcb = 0;
 	sotryfree(so);
 	LIST_REMOVE(rp, list);

sys/net/raw_usrreq.c

@@ -34,6 +34,7 @@
 #include <sys/lock.h>
 #include <sys/malloc.h>
 #include <sys/mbuf.h>
+#include <sys/mutex.h>
 #include <sys/protosw.h>
 #include <sys/signalvar.h>
 #include <sys/socket.h>
@@ -139,8 +140,9 @@ raw_uabort(struct socket *so)
 	if (rp == 0)
 		return EINVAL;
 	raw_disconnect(rp);
+	soisdisconnected(so);
+	SOCK_LOCK(so);
 	sotryfree(so);
-	soisdisconnected(so);	/* XXX huh? called after the sofree()? */
 	return 0;
 }
 

sys/netatalk/ddp_pcb.c

@@ -253,6 +253,7 @@ void
 at_pcbdetach(struct socket *so, struct ddpcb *ddp)
 {
     soisdisconnected(so);
+    SOCK_LOCK(so);
     so->so_pcb = NULL;
     sotryfree(so);
 

sys/netatm/atm_socket.c

@@ -173,6 +173,7 @@ atm_sock_detach(so)
 	/*
 	 * Break links and free control blocks
 	 */
+	SOCK_LOCK(so);
 	so->so_pcb = NULL;
 	sotryfree(so);
 

sys/netgraph/bluetooth/socket/ng_btsocket_hci_raw.c

@@ -1417,6 +1417,7 @@ ng_btsocket_hci_raw_detach(struct socket *so)
 	bzero(pcb, sizeof(*pcb));
 	FREE(pcb, M_NETGRAPH_BTSOCKET_HCI_RAW);
 
+	SOCK_LOCK(so);
 	so->so_pcb = NULL;
 	sotryfree(so);
 

sys/netgraph/bluetooth/socket/ng_btsocket_l2cap.c

@@ -1804,6 +1804,7 @@ ng_btsocket_l2cap_rtclean(void *context, int pending)
 				FREE(pcb, M_NETGRAPH_BTSOCKET_L2CAP);
 
 				soisdisconnected(so);
+				SOCK_LOCK(so);
 				so->so_pcb = NULL;
 				sotryfree(so);
 
@@ -2346,6 +2347,7 @@ ng_btsocket_l2cap_detach(struct socket *so)
 	FREE(pcb, M_NETGRAPH_BTSOCKET_L2CAP);
 
 	soisdisconnected(so);
+	SOCK_LOCK(so);
 	so->so_pcb = NULL;
 	sotryfree(so);
 

sys/netgraph/bluetooth/socket/ng_btsocket_l2cap_raw.c

@@ -1129,6 +1129,7 @@ ng_btsocket_l2cap_raw_detach(struct socket *so)
 	bzero(pcb, sizeof(*pcb));
 	FREE(pcb, M_NETGRAPH_BTSOCKET_L2CAP_RAW);
 
+	SOCK_LOCK(so);
 	so->so_pcb = NULL;
 	sotryfree(so);
 

sys/netgraph/bluetooth/socket/ng_btsocket_rfcomm.c

@@ -724,6 +724,7 @@ ng_btsocket_rfcomm_detach(struct socket *so)
 	FREE(pcb, M_NETGRAPH_BTSOCKET_RFCOMM);
 
 	soisdisconnected(so);
+	SOCK_LOCK(so);
 	so->so_pcb = NULL;
 	sotryfree(so);
 
@@ -1370,8 +1371,10 @@ ng_btsocket_rfcomm_session_accept(ng_btsocket_rfcomm_session_p s0)
 	s0->l2so->so_qlen --;
 	l2so->so_qstate &= ~SQ_COMP;
 	l2so->so_head = NULL;
+	SOCK_LOCK(l2so);
 	soref(l2so);
 	l2so->so_state |= SS_NBIO;
+	SOCK_UNLOCK(l2so);
 	ACCEPT_UNLOCK();
 
 	error = soaccept(l2so, (struct sockaddr **) &l2sa);

sys/netgraph/ng_ksocket.c

@@ -1205,8 +1205,10 @@ ng_ksocket_finish_accept(priv_p priv)
 	head->so_qlen--;
 	so->so_qstate &= ~SQ_COMP;
 	so->so_head = NULL;
+	SOCK_LOCK(so);
 	soref(so);
 	so->so_state |= SS_NBIO;
+	SOCK_UNLOCK(so);
 	ACCEPT_UNLOCK();
 
 	/* XXX KNOTE(&head->so_rcv.sb_sel.si_note, 0); */

sys/netinet/in_pcb.c

@@ -690,6 +690,7 @@ in_pcbdetach(inp)
 	inp->inp_gencnt = ++ipi->ipi_gencnt;
 	in_pcbremlists(inp);
 	if (so) {
+		SOCK_LOCK(so);
 		so->so_pcb = 0;
 		sotryfree(so);
 	}

sys/netinet/tcp_subr.c

@@ -1657,13 +1657,14 @@ tcp_twstart(tp)
 	}
 	tcp_discardcb(tp);
 	so = inp->inp_socket;
+	SOCK_LOCK(so);
 	so->so_pcb = NULL;
 	tw->tw_cred = crhold(so->so_cred);
 	tw->tw_so_options = so->so_options;
-	if (acknow)
-		tcp_twrespond(tw, TH_ACK);
 	sotryfree(so);
 	inp->inp_socket = NULL;
+	if (acknow)
+		tcp_twrespond(tw, TH_ACK);
 	inp->inp_ppcb = (caddr_t)tw;
 	inp->inp_vflag |= INP_TIMEWAIT;
 	tcp_timer_2msl_reset(tw, tw_time);

sys/netinet/tcp_timewait.c

@@ -1657,13 +1657,14 @@ tcp_twstart(tp)
 	}
 	tcp_discardcb(tp);
 	so = inp->inp_socket;
+	SOCK_LOCK(so);
 	so->so_pcb = NULL;
 	tw->tw_cred = crhold(so->so_cred);
 	tw->tw_so_options = so->so_options;
-	if (acknow)
-		tcp_twrespond(tw, TH_ACK);
 	sotryfree(so);
 	inp->inp_socket = NULL;
+	if (acknow)
+		tcp_twrespond(tw, TH_ACK);
 	inp->inp_ppcb = (caddr_t)tw;
 	inp->inp_vflag |= INP_TIMEWAIT;
 	tcp_timer_2msl_reset(tw, tw_time);

sys/netipx/ipx_pcb.c

@@ -268,6 +268,7 @@ ipx_pcbdetach(ipxp)
 {
 	struct socket *so = ipxp->ipxp_socket;
 
+	SOCK_LOCK(so);
 	so->so_pcb = 0;
 	sotryfree(so);
 	if (ipxp->ipxp_route.ro_rt != NULL)

sys/netipx/ipx_usrreq.c

@@ -423,8 +423,9 @@ ipx_usr_abort(so)
 	s = splnet();
 	ipx_pcbdetach(ipxp);
 	splx(s);
-	sotryfree(so);
 	soisdisconnected(so);
+	SOCK_LOCK(so);
+	sotryfree(so);
 	return (0);
 }
 

sys/netnatm/natm.c

@@ -135,6 +135,7 @@ natm_usr_detach(struct socket *so)
      * we turn on 'drain' *before* we sofree.
      */
     npcb_free(npcb, NPCB_DESTROY);	/* drain */
+    SOCK_LOCK(so);
     so->so_pcb = NULL;
     sotryfree(so);
  out:
@@ -463,6 +464,7 @@ struct proc *p;
        */
 
       npcb_free(npcb, NPCB_DESTROY);	/* drain */
+      SOCK_LOCK(so);
       so->so_pcb = NULL;
       sotryfree(so);
 

sys/sys/socketvar.h

@@ -57,7 +57,7 @@ typedef	u_quad_t so_gen_t;
  * (g) used only as a sleep/wakeup address, no value.
  */
 struct socket {
-	int	so_count;		/* reference count */
+	int	so_count;		/* (b) reference count */
 	short	so_type;		/* (a) generic type, see socket.h */
 	short	so_options;		/* from socket call, see socket.h */
 	short	so_linger;		/* time to linger while closing */
@@ -316,19 +316,26 @@ struct xsocket {
  * the structure.
  */
 #define	soref(so) do {							\
+	SOCK_LOCK_ASSERT(so);						\
 	++(so)->so_count;						\
 } while (0)
 
 #define	sorele(so) do {							\
+	SOCK_LOCK_ASSERT(so);						\
 	if ((so)->so_count <= 0)					\
 		panic("sorele");					\
 	if (--(so)->so_count == 0)					\
 		sofree(so);						\
+	else								\
+		SOCK_UNLOCK(so);					\
 } while (0)
 
 #define	sotryfree(so) do {						\
+	SOCK_LOCK_ASSERT(so);						\
 	if ((so)->so_count == 0)					\
 		sofree(so);						\
+	else								\
+		SOCK_UNLOCK(so);					\
 } while(0)
 
 #define	sorwakeup(so) do {						\