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

4 files changed, 21 insertions, 1 deletions

diff --git a/sys/kern/kern_descrip.c b/sys/kern/kern_descrip.c
index e684efd..8e955ca 100644
--- a/sys/kern/kern_descrip.c
+++ b/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);
 }
 
diff --git a/sys/kern/uipc_socket.c b/sys/kern/uipc_socket.c
index 66a10d9f..145a5a6 100644
--- a/sys/kern/uipc_socket.c
+++ b/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 @@ drop:
 			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;
 	}
diff --git a/sys/kern/uipc_syscalls.c b/sys/kern/uipc_syscalls.c
index 9ea94f4..c99fc89 100644
--- a/sys/kern/uipc_syscalls.c
+++ b/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--;
diff --git a/sys/kern/uipc_usrreq.c b/sys/kern/uipc_usrreq.c
index 76a34e1..6660d7b 100644
--- a/sys/kern/uipc_usrreq.c
+++ b/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);
 }