From 47517268b302171d3b3e5100a61f63c8cdbd5a97 Mon Sep 17 00:00:00 2001
From: mvs <mvs@openbsd.org>
Date: Wed, 27 Mar 2024 22:47:53 +0000
Subject: [PATCH] Introduce SB_OWNLOCK to mark sockets which `so_rcv' buffer
 modified outside socket lock.

`sb_mtx' mutex(9) used for this case and it should not be released between
`so_rcv' usage check and corresponding sbwait() sleep. Otherwise wakeup()
could be lost sometimes.

ok bluhm
---
 sys/kern/uipc_socket.c  | 25 ++++++++++++++++++-------
 sys/kern/uipc_socket2.c | 31 +++++++++++++++++++++++--------
 sys/sys/socketvar.h     | 21 ++++++++++++---------
 3 files changed, 53 insertions(+), 24 deletions(-)

diff --git a/sys/kern/uipc_socket.c b/sys/kern/uipc_socket.c
index 06313d40cbf..1d8a101f107 100644
--- a/sys/kern/uipc_socket.c
+++ b/sys/kern/uipc_socket.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: uipc_socket.c,v 1.322 2024/03/26 09:46:47 mvs Exp $	*/
+/*	$OpenBSD: uipc_socket.c,v 1.323 2024/03/27 22:47:53 mvs Exp $	*/
 /*	$NetBSD: uipc_socket.c,v 1.21 1996/02/04 02:17:52 christos Exp $	*/
 
 /*
@@ -161,7 +161,7 @@ soalloc(const struct protosw *prp, int wait)
 		}
 		break;
 	case AF_UNIX:
-		so->so_rcv.sb_flags |= SB_MTXLOCK;
+		so->so_rcv.sb_flags |= SB_MTXLOCK | SB_OWNLOCK;
 		break;
 	}
 
@@ -903,12 +903,23 @@ restart:
 		}
 		SBLASTRECORDCHK(&so->so_rcv, "soreceive sbwait 1");
 		SBLASTMBUFCHK(&so->so_rcv, "soreceive sbwait 1");
-		sb_mtx_unlock(&so->so_rcv);
-		sbunlock(so, &so->so_rcv);
-		error = sbwait(so, &so->so_rcv);
-		if (error) {
+
+		if (so->so_rcv.sb_flags & SB_OWNLOCK) {
+			sbunlock_locked(so, &so->so_rcv);
 			sounlock_shared(so);
-			return (error);
+			error = sbwait_locked(so, &so->so_rcv);
+			sb_mtx_unlock(&so->so_rcv);
+			if (error)
+				return (error);
+			solock_shared(so);
+		} else {
+			sb_mtx_unlock(&so->so_rcv);
+			sbunlock(so, &so->so_rcv);
+			error = sbwait(so, &so->so_rcv);
+			if (error) {
+				sounlock_shared(so);
+				return (error);
+			}
 		}
 		goto restart;
 	}
diff --git a/sys/kern/uipc_socket2.c b/sys/kern/uipc_socket2.c
index 153010a59cc..4a9eaae346d 100644
--- a/sys/kern/uipc_socket2.c
+++ b/sys/kern/uipc_socket2.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: uipc_socket2.c,v 1.145 2024/03/26 09:46:47 mvs Exp $	*/
+/*	$OpenBSD: uipc_socket2.c,v 1.146 2024/03/27 22:47:53 mvs Exp $	*/
 /*	$NetBSD: uipc_socket2.c,v 1.11 1996/02/04 02:17:55 christos Exp $	*/
 
 /*
@@ -522,6 +522,18 @@ sbmtxassertlocked(struct socket *so, struct sockbuf *sb)
 /*
  * Wait for data to arrive at/drain from a socket buffer.
  */
+int
+sbwait_locked(struct socket *so, struct sockbuf *sb)
+{
+	int prio = (sb->sb_flags & SB_NOINTR) ? PSOCK : PSOCK | PCATCH;
+
+	MUTEX_ASSERT_LOCKED(&sb->sb_mtx);
+
+	sb->sb_flags |= SB_WAIT;
+	return msleep_nsec(&sb->sb_cc, &sb->sb_mtx, prio, "sbwait",
+	    sb->sb_timeo_nsecs);
+}
+
 int
 sbwait(struct socket *so, struct sockbuf *sb)
 {
@@ -573,20 +585,23 @@ out:
 }
 
 void
-sbunlock(struct socket *so, struct sockbuf *sb)
+sbunlock_locked(struct socket *so, struct sockbuf *sb)
 {
-	int dowakeup = 0;
+	MUTEX_ASSERT_LOCKED(&sb->sb_mtx);
 
-	mtx_enter(&sb->sb_mtx);
 	sb->sb_flags &= ~SB_LOCK;
 	if (sb->sb_flags & SB_WANT) {
 		sb->sb_flags &= ~SB_WANT;
-		dowakeup = 1;
+		wakeup(&sb->sb_flags);
 	}
-	mtx_leave(&sb->sb_mtx);
+}
 
-	if (dowakeup)
-		wakeup(&sb->sb_flags);
+void
+sbunlock(struct socket *so, struct sockbuf *sb)
+{
+	mtx_enter(&sb->sb_mtx);
+	sbunlock_locked(so, sb);
+	mtx_leave(&sb->sb_mtx);
 }
 
 /*
diff --git a/sys/sys/socketvar.h b/sys/sys/socketvar.h
index 782ebce75fd..0169c6273d6 100644
--- a/sys/sys/socketvar.h
+++ b/sys/sys/socketvar.h
@@ -1,4 +1,4 @@
-/*	$OpenBSD: socketvar.h,v 1.126 2024/03/26 09:46:47 mvs Exp $	*/
+/*	$OpenBSD: socketvar.h,v 1.127 2024/03/27 22:47:53 mvs Exp $	*/
 /*	$NetBSD: socketvar.h,v 1.18 1996/02/09 18:25:38 christos Exp $	*/
 
 /*-
@@ -127,14 +127,15 @@ struct socket {
 		uint64_t sb_timeo_nsecs;/* timeout for read/write */
 		struct klist sb_klist;	/* process selecting read/write */
 	} so_rcv, so_snd;
-#define	SB_MAX		(2*1024*1024)	/* default for max chars in sockbuf */
-#define	SB_LOCK		0x01		/* lock on data queue */
-#define	SB_WANT		0x02		/* someone is waiting to lock */
-#define	SB_WAIT		0x04		/* someone is waiting for data/space */
-#define	SB_ASYNC	0x10		/* ASYNC I/O, need signals */
-#define	SB_SPLICE	0x20		/* buffer is splice source or drain */
-#define	SB_NOINTR	0x40		/* operations not interruptible */
-#define SB_MTXLOCK	0x80		/* use sb_mtx for sockbuf protection */
+#define SB_MAX		(2*1024*1024)	/* default for max chars in sockbuf */
+#define SB_LOCK		0x0001		/* lock on data queue */
+#define SB_WANT		0x0002		/* someone is waiting to lock */
+#define SB_WAIT		0x0004		/* someone is waiting for data/space */
+#define SB_ASYNC	0x0010		/* ASYNC I/O, need signals */
+#define SB_SPLICE	0x0020		/* buffer is splice source or drain */
+#define SB_NOINTR	0x0040		/* operations not interruptible */
+#define SB_MTXLOCK	0x0080		/* use sb_mtx for sockbuf protection */
+#define SB_OWNLOCK	0x0100		/* sb_mtx used standalone */
 
 	void	(*so_upcall)(struct socket *so, caddr_t arg, int waitf);
 	caddr_t	so_upcallarg;		/* Arg for above */
@@ -320,6 +321,7 @@ int sblock(struct socket *, struct sockbuf *, int);
 
 /* release lock on sockbuf sb */
 void sbunlock(struct socket *, struct sockbuf *);
+void sbunlock_locked(struct socket *, struct sockbuf *);
 
 #define	SB_EMPTY_FIXUP(sb) do {						\
 	if ((sb)->sb_mb == NULL) {					\
@@ -367,6 +369,7 @@ int	sbcheckreserve(u_long, u_long);
 int	sbchecklowmem(void);
 int	sbreserve(struct socket *, struct sockbuf *, u_long);
 int	sbwait(struct socket *, struct sockbuf *);
+int	sbwait_locked(struct socket *, struct sockbuf *);
 void	soinit(void);
 void	soabort(struct socket *);
 int	soaccept(struct socket *, struct mbuf *);
-- 
2.20.1