in imsg_read() avoid calling recvmsg() if we detect that we will be short

on descriptors, this can be achieved thanks to the new getdtablecount()
system call. application may provide a reserve count to ensure that the
recvmsg() call is not called when they don't have enough descriptors to
work properly.

change the API so that transient errors that can be retried immediately
are retried within the function right away, whereas transient errors for
which the application may want to take action will set errno to EAGAIN.

ok deraadt@ and henning@

2 files changed, 60 insertions, 28 deletions

diff --git a/lib/libutil/imsg-buffer.c b/lib/libutil/imsg-buffer.c
index dec27ffefd3..9f0475774dd 100644
--- a/lib/libutil/imsg-buffer.c
+++ b/lib/libutil/imsg-buffer.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: imsg-buffer.c,v 1.1 2010/05/26 16:44:32 nicm Exp $	*/
+/*	$OpenBSD: imsg-buffer.c,v 1.2 2012/06/02 21:46:53 gilles Exp $	*/
 
 /*
  * Copyright (c) 2003, 2004 Henning Brauer <henning@openbsd.org>
@@ -157,22 +157,23 @@ ibuf_write(struct msgbuf *msgbuf)
 		i++;
 	}
 
+again:
 	if ((n = writev(msgbuf->fd, iov, i)) == -1) {
-		if (errno == EAGAIN || errno == ENOBUFS ||
-		    errno == EINTR)	/* try later */
-			return (0);
-		else
-			return (-1);
+		if (errno == EAGAIN || errno == EINTR)
+			goto again;
+		if (errno == ENOBUFS)
+			errno = EAGAIN;
+		return (-1);
 	}
 
 	if (n == 0) {			/* connection closed */
 		errno = 0;
-		return (-2);
+		return (0);
 	}
 
 	msgbuf_drain(msgbuf, n);
 
-	return (0);
+	return (1);
 }
 
 void
@@ -256,17 +257,18 @@ msgbuf_write(struct msgbuf *msgbuf)
 		*(int *)CMSG_DATA(cmsg) = buf->fd;
 	}
 
+again:
 	if ((n = sendmsg(msgbuf->fd, &msg, 0)) == -1) {
-		if (errno == EAGAIN || errno == ENOBUFS ||
-		    errno == EINTR)	/* try later */
-			return (0);
-		else
-			return (-1);
+		if (errno == EAGAIN || errno == EINTR)
+			goto again;
+		if (errno == ENOBUFS)
+			errno = EAGAIN;
+		return (-1);
 	}
 
 	if (n == 0) {			/* connection closed */
 		errno = 0;
-		return (-2);
+		return (0);
 	}
 
 	/*
@@ -280,7 +282,7 @@ msgbuf_write(struct msgbuf *msgbuf)
 
 	msgbuf_drain(msgbuf, n);
 
-	return (0);
+	return (1);
 }
 
 void
diff --git a/lib/libutil/imsg.c b/lib/libutil/imsg.c
index a0be894f688..05e57c7e2d4 100644
--- a/lib/libutil/imsg.c
+++ b/lib/libutil/imsg.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: imsg.c,v 1.1 2010/05/26 16:44:32 nicm Exp $	*/
+/*	$OpenBSD: imsg.c,v 1.2 2012/06/02 21:46:53 gilles Exp $	*/
 
 /*
  * Copyright (c) 2003, 2004 Henning Brauer <henning@openbsd.org>
@@ -28,6 +28,8 @@
 
 #include "imsg.h"
 
+int	 imsg_fd_overhead = 0;
+
 int	 imsg_get_fd(struct imsgbuf *);
 
 void
@@ -48,10 +50,10 @@ imsg_read(struct imsgbuf *ibuf)
 	struct cmsghdr		*cmsg;
 	union {
 		struct cmsghdr hdr;
-		char	buf[CMSG_SPACE(sizeof(int) * 16)];
+		char	buf[CMSG_SPACE(sizeof(int) * 1)];
 	} cmsgbuf;
 	struct iovec		 iov;
-	ssize_t			 n;
+	ssize_t			 n = -1;
 	int			 fd;
 	struct imsg_fd		*ifd;
 
@@ -64,11 +66,23 @@ imsg_read(struct imsgbuf *ibuf)
 	msg.msg_control = &cmsgbuf.buf;
 	msg.msg_controllen = sizeof(cmsgbuf.buf);
 
+	if ((ifd = calloc(1, sizeof(struct imsg_fd))) == NULL)
+		return (-1);
+
+again:
+	if (getdtablecount() + imsg_fd_overhead +
+	    (CMSG_SPACE(sizeof(int))-CMSG_SPACE(0))/sizeof(int)
+	    >= getdtablesize()) {
+		errno = EAGAIN;
+		return (-1);
+	}
+	
 	if ((n = recvmsg(ibuf->fd, &msg, 0)) == -1) {
-		if (errno != EINTR && errno != EAGAIN) {
-			return (-1);
-		}
-		return (-2);
+		if (errno == EMSGSIZE)
+			goto fail;
+		if (errno != EINTR && errno != EAGAIN)
+			goto fail;
+		goto again;
 	}
 
 	ibuf->r.wpos += n;
@@ -77,17 +91,33 @@ imsg_read(struct imsgbuf *ibuf)
 	    cmsg = CMSG_NXTHDR(&msg, cmsg)) {
 		if (cmsg->cmsg_level == SOL_SOCKET &&
 		    cmsg->cmsg_type == SCM_RIGHTS) {
-			fd = (*(int *)CMSG_DATA(cmsg));
-			if ((ifd = calloc(1, sizeof(struct imsg_fd))) == NULL) {
-				close(fd);
-				return (-1);
+			int i;
+			int j;
+
+			/*
+			 * We only accept one file descriptor.  Due to C
+			 * padding rules, our control buffer might contain
+			 * more than one fd, and we must close them.
+			 */
+			j = ((char *)cmsg + cmsg->cmsg_len -
+			    (char *)CMSG_DATA(cmsg)) / sizeof(int);
+			for (i = 0; i < j; i++) {
+				fd = ((int *)CMSG_DATA(cmsg))[i];
+				if (i == 0) {
+					ifd->fd = fd;
+					TAILQ_INSERT_TAIL(&ibuf->fds, ifd,
+					    entry);
+					ifd = NULL;
+				} else
+					close(fd);
 			}
-			ifd->fd = fd;
-			TAILQ_INSERT_TAIL(&ibuf->fds, ifd, entry);
 		}
 		/* we do not handle other ctl data level */
 	}
 
+fail:
+	if (ifd)
+		free(ifd);
 	return (n);
 }