Implement RLIMIT_SBSIZE in the kernel. This is a per-uid sockbuf total

usage limit.

10 files changed, 109 insertions, 23 deletions

diff --git a/sys/kern/kern_proc.c b/sys/kern/kern_proc.c
index 28e1f6d..0ead052 100644
--- a/sys/kern/kern_proc.c
+++ b/sys/kern/kern_proc.c
@@ -64,6 +64,7 @@ struct uidinfo {
 	LIST_ENTRY(uidinfo) ui_hash;
 	uid_t	ui_uid;
 	long	ui_proccnt;
+	rlim_t	ui_sbsize;
 };
 #define	UIHASH(uid)	(&uihashtbl[(uid) & uihash])
 static LIST_HEAD(uihashhead, uidinfo) *uihashtbl;
@@ -115,10 +116,10 @@ chgproccnt(uid, diff)
 			break;
 	if (uip) {
 		uip->ui_proccnt += diff;
-		if (uip->ui_proccnt > 0)
-			return (uip->ui_proccnt);
 		if (uip->ui_proccnt < 0)
 			panic("chgproccnt: procs < 0");
+		if (uip->ui_proccnt > 0 || uip->ui_sbsize > 0)
+			return (uip->ui_proccnt);
 		LIST_REMOVE(uip, ui_hash);
 		FREE(uip, M_PROC);
 		return (0);
@@ -132,6 +133,45 @@ chgproccnt(uid, diff)
 	LIST_INSERT_HEAD(uipp, uip, ui_hash);
 	uip->ui_uid = uid;
 	uip->ui_proccnt = diff;
+	uip->ui_sbsize = 0;
+	return (diff);
+}
+
+/*
+ * Change the total socket buffer size a user has used.
+ */
+rlim_t
+chgsbsize(uid, diff)
+	uid_t	uid;
+	rlim_t	diff;
+{
+	register struct uidinfo *uip;
+	register struct uihashhead *uipp;
+
+	uipp = UIHASH(uid);
+	for (uip = uipp->lh_first; uip != 0; uip = uip->ui_hash.le_next)
+		if (uip->ui_uid == uid)
+			break;
+	if (diff <= 0) {
+		if (diff == 0)
+			return (uip ? uip->ui_sbsize : 0);
+		KASSERT(uip != NULL, ("uidinfo (%d) gone", uid));
+	}
+	if (uip) {
+		uip->ui_sbsize += diff;
+		KASSERT(uip->ui_sbsize >= 0, ("ui_sbsize (%d) < 0", uid));
+		if (uip->ui_sbsize == 0 && uip->ui_proccnt == 0) {
+			LIST_REMOVE(uip, ui_hash);
+			FREE(uip, M_PROC);
+			return (0);
+		}
+		return (uip->ui_sbsize);
+	}
+	MALLOC(uip, struct uidinfo *, sizeof(*uip), M_PROC, M_WAITOK);
+	LIST_INSERT_HEAD(uipp, uip, ui_hash);
+	uip->ui_uid = uid;
+	uip->ui_proccnt = 0;
+	uip->ui_sbsize = diff;
 	return (diff);
 }
 
diff --git a/sys/kern/uipc_sockbuf.c b/sys/kern/uipc_sockbuf.c
index 9c1cd07..c99bd3b 100644
--- a/sys/kern/uipc_sockbuf.c
+++ b/sys/kern/uipc_sockbuf.c
@@ -376,10 +376,11 @@ soreserve(so, sndcc, rcvcc)
 	register struct socket *so;
 	u_long sndcc, rcvcc;
 {
+	struct proc *p = curproc;
 
-	if (sbreserve(&so->so_snd, sndcc) == 0)
+	if (sbreserve(&so->so_snd, sndcc, so, p) == 0)
 		goto bad;
-	if (sbreserve(&so->so_rcv, rcvcc) == 0)
+	if (sbreserve(&so->so_rcv, rcvcc, so, p) == 0)
 		goto bad2;
 	if (so->so_rcv.sb_lowat == 0)
 		so->so_rcv.sb_lowat = 1;
@@ -389,7 +390,7 @@ soreserve(so, sndcc, rcvcc)
 		so->so_snd.sb_lowat = so->so_snd.sb_hiwat;
 	return (0);
 bad2:
-	sbrelease(&so->so_snd);
+	sbrelease(&so->so_snd, so);
 bad:
 	return (ENOBUFS);
 }
@@ -400,12 +401,25 @@ bad:
  * if buffering efficiency is near the normal case.
  */
 int
-sbreserve(sb, cc)
+sbreserve(sb, cc, so, p)
 	struct sockbuf *sb;
 	u_long cc;
+	struct socket *so;
+	struct proc *p;
 {
+	rlim_t delta;
+
+	/*
+	 * p will only be NULL when we're in an interrupt
+	 * (e.g. in tcp_input())
+	 */
 	if ((u_quad_t)cc > (u_quad_t)sb_max * MCLBYTES / (MSIZE + MCLBYTES))
 		return (0);
+	delta = (rlim_t)cc - sb->sb_hiwat;
+	if (p && delta >= 0 && chgsbsize(so->so_cred->cr_uid, 0) + delta >
+	    p->p_rlimit[RLIMIT_SBSIZE].rlim_cur)
+		return (0);
+	(void)chgsbsize(so->so_cred->cr_uid, delta);
 	sb->sb_hiwat = cc;
 	sb->sb_mbmax = min(cc * sb_efficiency, sb_max);
 	if (sb->sb_lowat > sb->sb_hiwat)
@@ -417,11 +431,13 @@ sbreserve(sb, cc)
  * Free mbufs held by a socket, and reserved mbuf space.
  */
 void
-sbrelease(sb)
+sbrelease(sb, so)
 	struct sockbuf *sb;
+	struct socket *so;
 {
 
 	sbflush(sb);
+	(void)chgsbsize(so->so_cred->cr_uid, -(rlim_t)sb->sb_hiwat);
 	sb->sb_hiwat = sb->sb_mbmax = 0;
 }
 
diff --git a/sys/kern/uipc_socket.c b/sys/kern/uipc_socket.c
index 9719b79..3d15c27 100644
--- a/sys/kern/uipc_socket.c
+++ b/sys/kern/uipc_socket.c
@@ -158,6 +158,12 @@ sodealloc(so)
 {
 
 	so->so_gencnt = ++so_gencnt;
+	if (so->so_rcv.sb_hiwat)
+		(void)chgsbsize(so->so_cred->cr_uid,
+		    -(rlim_t)so->so_rcv.sb_hiwat);
+	if (so->so_snd.sb_hiwat)
+		(void)chgsbsize(so->so_cred->cr_uid,
+		    -(rlim_t)so->so_snd.sb_hiwat);
 	crfree(so->so_cred);
 	zfreei(so->so_zone, so);
 }
@@ -212,7 +218,7 @@ sofree(so)
 		so->so_state &= ~SS_INCOMP;
 		so->so_head = NULL;
 	}
-	sbrelease(&so->so_snd);
+	sbrelease(&so->so_snd, so);
 	sorflush(so);
 	sodealloc(so);
 }
@@ -922,7 +928,7 @@ sorflush(so)
 	splx(s);
 	if (pr->pr_flags & PR_RIGHTS && pr->pr_domain->dom_dispose)
 		(*pr->pr_domain->dom_dispose)(asb.sb_mb);
-	sbrelease(&asb);
+	sbrelease(&asb, so);
 }
 
 /*
@@ -1030,8 +1036,8 @@ sosetopt(so, sopt)
 			case SO_SNDBUF:
 			case SO_RCVBUF:
 				if (sbreserve(sopt->sopt_name == SO_SNDBUF ?
-					      &so->so_snd : &so->so_rcv,
-					      (u_long) optval) == 0) {
+				    &so->so_snd : &so->so_rcv, (u_long)optval,
+				    so, curproc) == 0) {
 					error = ENOBUFS;
 					goto bad;
 				}
diff --git a/sys/kern/uipc_socket2.c b/sys/kern/uipc_socket2.c
index 9c1cd07..c99bd3b 100644
--- a/sys/kern/uipc_socket2.c
+++ b/sys/kern/uipc_socket2.c
@@ -376,10 +376,11 @@ soreserve(so, sndcc, rcvcc)
 	register struct socket *so;
 	u_long sndcc, rcvcc;
 {
+	struct proc *p = curproc;
 
-	if (sbreserve(&so->so_snd, sndcc) == 0)
+	if (sbreserve(&so->so_snd, sndcc, so, p) == 0)
 		goto bad;
-	if (sbreserve(&so->so_rcv, rcvcc) == 0)
+	if (sbreserve(&so->so_rcv, rcvcc, so, p) == 0)
 		goto bad2;
 	if (so->so_rcv.sb_lowat == 0)
 		so->so_rcv.sb_lowat = 1;
@@ -389,7 +390,7 @@ soreserve(so, sndcc, rcvcc)
 		so->so_snd.sb_lowat = so->so_snd.sb_hiwat;
 	return (0);
 bad2:
-	sbrelease(&so->so_snd);
+	sbrelease(&so->so_snd, so);
 bad:
 	return (ENOBUFS);
 }
@@ -400,12 +401,25 @@ bad:
  * if buffering efficiency is near the normal case.
  */
 int
-sbreserve(sb, cc)
+sbreserve(sb, cc, so, p)
 	struct sockbuf *sb;
 	u_long cc;
+	struct socket *so;
+	struct proc *p;
 {
+	rlim_t delta;
+
+	/*
+	 * p will only be NULL when we're in an interrupt
+	 * (e.g. in tcp_input())
+	 */
 	if ((u_quad_t)cc > (u_quad_t)sb_max * MCLBYTES / (MSIZE + MCLBYTES))
 		return (0);
+	delta = (rlim_t)cc - sb->sb_hiwat;
+	if (p && delta >= 0 && chgsbsize(so->so_cred->cr_uid, 0) + delta >
+	    p->p_rlimit[RLIMIT_SBSIZE].rlim_cur)
+		return (0);
+	(void)chgsbsize(so->so_cred->cr_uid, delta);
 	sb->sb_hiwat = cc;
 	sb->sb_mbmax = min(cc * sb_efficiency, sb_max);
 	if (sb->sb_lowat > sb->sb_hiwat)
@@ -417,11 +431,13 @@ sbreserve(sb, cc)
  * Free mbufs held by a socket, and reserved mbuf space.
  */
 void
-sbrelease(sb)
+sbrelease(sb, so)
 	struct sockbuf *sb;
+	struct socket *so;
 {
 
 	sbflush(sb);
+	(void)chgsbsize(so->so_cred->cr_uid, -(rlim_t)sb->sb_hiwat);
 	sb->sb_hiwat = sb->sb_mbmax = 0;
 }
 
diff --git a/sys/kern/uipc_usrreq.c b/sys/kern/uipc_usrreq.c
index 7f9304a..83a7831 100644
--- a/sys/kern/uipc_usrreq.c
+++ b/sys/kern/uipc_usrreq.c
@@ -236,6 +236,8 @@ uipc_rcvd(struct socket *so, int flags)
 		so2->so_snd.sb_mbmax += unp->unp_mbcnt - so->so_rcv.sb_mbcnt;
 		unp->unp_mbcnt = so->so_rcv.sb_mbcnt;
 		so2->so_snd.sb_hiwat += unp->unp_cc - so->so_rcv.sb_cc;
+		(void)chgsbsize(so2->so_cred->cr_uid,
+		    (rlim_t)unp->unp_cc - so->so_rcv.sb_cc);
 		unp->unp_cc = so->so_rcv.sb_cc;
 		sowwakeup(so2);
 		break;
@@ -342,6 +344,8 @@ uipc_send(struct socket *so, int flags, struct mbuf *m, struct sockaddr *nam,
 		unp->unp_conn->unp_mbcnt = so2->so_rcv.sb_mbcnt;
 		so->so_snd.sb_hiwat -=
 		    so2->so_rcv.sb_cc - unp->unp_conn->unp_cc;
+		(void)chgsbsize(so->so_cred->cr_uid,
+		    (rlim_t)unp->unp_conn->unp_cc - so2->so_rcv.sb_cc);
 		unp->unp_conn->unp_cc = so2->so_rcv.sb_cc;
 		sorwakeup(so2);
 		m = 0;
diff --git a/sys/netinet/tcp_input.c b/sys/netinet/tcp_input.c
index e45c64b..5f61ff3 100644
--- a/sys/netinet/tcp_input.c
+++ b/sys/netinet/tcp_input.c
@@ -2298,7 +2298,7 @@ tcp_mss(tp, offer)
 		bufsize = roundup(bufsize, mss);
 		if (bufsize > sb_max)
 			bufsize = sb_max;
-		(void)sbreserve(&so->so_snd, bufsize);
+		(void)sbreserve(&so->so_snd, bufsize, so, NULL);
 	}
 	tp->t_maxseg = mss;
 
@@ -2310,7 +2310,7 @@ tcp_mss(tp, offer)
 		bufsize = roundup(bufsize, mss);
 		if (bufsize > sb_max)
 			bufsize = sb_max;
-		(void)sbreserve(&so->so_rcv, bufsize);
+		(void)sbreserve(&so->so_rcv, bufsize, so, NULL);
 	}
 
 	/*
diff --git a/sys/netinet/tcp_reass.c b/sys/netinet/tcp_reass.c
index e45c64b..5f61ff3 100644
--- a/sys/netinet/tcp_reass.c
+++ b/sys/netinet/tcp_reass.c
@@ -2298,7 +2298,7 @@ tcp_mss(tp, offer)
 		bufsize = roundup(bufsize, mss);
 		if (bufsize > sb_max)
 			bufsize = sb_max;
-		(void)sbreserve(&so->so_snd, bufsize);
+		(void)sbreserve(&so->so_snd, bufsize, so, NULL);
 	}
 	tp->t_maxseg = mss;
 
@@ -2310,7 +2310,7 @@ tcp_mss(tp, offer)
 		bufsize = roundup(bufsize, mss);
 		if (bufsize > sb_max)
 			bufsize = sb_max;
-		(void)sbreserve(&so->so_rcv, bufsize);
+		(void)sbreserve(&so->so_rcv, bufsize, so, NULL);
 	}
 
 	/*
diff --git a/sys/sys/proc.h b/sys/sys/proc.h
index 99e429f..f3cee66 100644
--- a/sys/sys/proc.h
+++ b/sys/sys/proc.h
@@ -375,6 +375,7 @@ struct vm_zone;
 extern struct vm_zone *proc_zone;
 
 int	chgproccnt __P((uid_t uid, int diff));
+rlim_t	chgsbsize __P((uid_t uid, rlim_t diff));
 int	enterpgrp __P((struct proc *p, pid_t pgid, int mksess));
 void	fixjobc __P((struct proc *p, struct pgrp *pgrp, int entering));
 int	inferior __P((struct proc *p));
diff --git a/sys/sys/resource.h b/sys/sys/resource.h
index 3dc2182..5b8491e 100644
--- a/sys/sys/resource.h
+++ b/sys/sys/resource.h
@@ -87,8 +87,9 @@ struct	rusage {
 #define	RLIMIT_MEMLOCK	6		/* locked-in-memory address space */
 #define	RLIMIT_NPROC	7		/* number of processes */
 #define	RLIMIT_NOFILE	8		/* number of open files */
+#define	RLIMIT_SBSIZE	9		/* maximum size of all socket buffers */
 
-#define	RLIM_NLIMITS	9		/* number of resource limits */
+#define	RLIM_NLIMITS	10		/* number of resource limits */
 
 #define	RLIM_INFINITY	((rlim_t)(((u_quad_t)1 << 63) - 1))
 
@@ -108,6 +109,7 @@ static char *rlimit_ident[] = {
 	"memlock",
 	"nproc",
 	"nofile",
+	"sbsize",
 };
 #endif
 
diff --git a/sys/sys/socketvar.h b/sys/sys/socketvar.h
index b9539ca..471ddbe 100644
--- a/sys/sys/socketvar.h
+++ b/sys/sys/socketvar.h
@@ -319,8 +319,9 @@ void	sbdrop __P((struct sockbuf *sb, int len));
 void	sbdroprecord __P((struct sockbuf *sb));
 void	sbflush __P((struct sockbuf *sb));
 void	sbinsertoob __P((struct sockbuf *sb, struct mbuf *m0));
-void	sbrelease __P((struct sockbuf *sb));
-int	sbreserve __P((struct sockbuf *sb, u_long cc));
+void	sbrelease __P((struct sockbuf *sb, struct socket *so));
+int	sbreserve __P((struct sockbuf *sb, u_long cc, struct socket *so,
+		       struct proc *p));
 void	sbtoxsockbuf __P((struct sockbuf *sb, struct xsockbuf *xsb));
 int	sbwait __P((struct sockbuf *sb));
 int	sb_lock __P((struct sockbuf *sb));