5 files changed, 98 insertions, 16 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;