Always check the return value of the write() system call

dtach was assuming that writes with a small byte count could never fail,
and was not handling partial writes for larger byte counts in a few
places.

This should also suppress unused result warnings from gcc on systems
that define _FORTIFY_SOURCE by default.

4 files changed, 64 insertions, 14 deletions

diff --git a/attach.c b/attach.c
index 8ac3bf3..217838e 100644
--- a/attach.c
+++ b/attach.c
@@ -111,7 +111,7 @@ process_kbd(int s, struct packet *pkt)
 	{
 		/* Tell the master that we are suspending. */
 		pkt->type = MSG_DETACH;
-		write(s, pkt, sizeof(struct packet));
+		write_packet_or_fail(s, pkt);
 
 		/* And suspend... */
 		tcsetattr(0, TCSADRAIN, &orig_term);
@@ -121,13 +121,13 @@ process_kbd(int s, struct packet *pkt)
 
 		/* Tell the master that we are returning. */
 		pkt->type = MSG_ATTACH;
-		write(s, pkt, sizeof(struct packet));
+		write_packet_or_fail(s, pkt);
 
 		/* We would like a redraw, too. */
 		pkt->type = MSG_REDRAW;
 		pkt->len = redraw_method;
 		ioctl(0, TIOCGWINSZ, &pkt->u.ws);
-		write(s, pkt, sizeof(struct packet));
+		write_packet_or_fail(s, pkt);
 		return;
 	}
 	/* Detach char? */
@@ -141,7 +141,7 @@ process_kbd(int s, struct packet *pkt)
 		win_changed = 1;
 
 	/* Push it out */
-	write(s, pkt, sizeof(struct packet));
+	write_packet_or_fail(s, pkt);
 }
 
 int
@@ -218,18 +218,18 @@ attach_main(int noerror)
 	tcsetattr(0, TCSADRAIN, &cur_term);
 
 	/* Clear the screen. This assumes VT100. */
-	write(1, "\33[H\33[J", 6);
+	write_buf_or_fail(1, "\33[H\33[J", 6);
 
 	/* Tell the master that we want to attach. */
 	memset(&pkt, 0, sizeof(struct packet));
 	pkt.type = MSG_ATTACH;
-	write(s, &pkt, sizeof(struct packet));
+	write_packet_or_fail(s, &pkt);
 
 	/* We would like a redraw, too. */
 	pkt.type = MSG_REDRAW;
 	pkt.len = redraw_method;
 	ioctl(0, TIOCGWINSZ, &pkt.u.ws);
-	write(s, &pkt, sizeof(struct packet));
+	write_packet_or_fail(s, &pkt);
 
 	/* Wait for things to happen */
 	while (1)
@@ -263,7 +263,7 @@ attach_main(int noerror)
 				exit(1);
 			}
 			/* Send the data to the terminal. */
-			write(1, buf, len);
+			write_buf_or_fail(1, buf, len);
 			n--;
 		}
 		/* stdin activity */
@@ -290,7 +290,7 @@ attach_main(int noerror)
 
 			pkt.type = MSG_WINCH;
 			ioctl(0, TIOCGWINSZ, &pkt.u.ws);
-			write(s, &pkt, sizeof(pkt));
+			write_packet_or_fail(s, &pkt);
 		}
 	}
 	return 0;
@@ -358,8 +358,12 @@ push_main()
 		}
 
 		pkt.len = len;
-		if (write(s, &pkt, sizeof(struct packet)) < 0)
+		len = write(s, &pkt, sizeof(struct packet));
+		if (len != sizeof(struct packet))
 		{
+			if (len >= 0)
+				errno = EPIPE;
+
 			printf("%s: %s: %s\n", progname, sockname,
 			       strerror(errno));
 			return 1;
diff --git a/dtach.h b/dtach.h
index 56867c8..fd6b59c 100644
--- a/dtach.h
+++ b/dtach.h
@@ -139,6 +139,9 @@ struct packet
 /* This hopefully moves to the bottom of the screen */
 #define EOS "\033[999H"
 
+void write_buf_or_fail(int fd, const void *buf, size_t count);
+void write_packet_or_fail(int fd, const struct packet *pkt);
+
 int attach_main(int noerror);
 int master_main(char **argv, int waitattach, int dontfork);
 int push_main(void);
diff --git a/main.c b/main.c
index 89e22ce..5d991fa 100644
--- a/main.c
+++ b/main.c
@@ -45,6 +45,49 @@ int redraw_method = REDRAW_UNSPEC;
 struct termios orig_term;
 int dont_have_tty;
 
+/* Write buf to fd handling partial writes. Exit on failure. */
+void
+write_buf_or_fail(int fd, const void *buf, size_t count)
+{
+	while (count != 0)
+	{
+		ssize_t ret = write(fd, buf, count);
+
+		if (ret >= 0)
+		{
+			buf = (const char *)buf + ret;
+			count -= ret;
+		}
+		else if (ret < 0 && errno == EINTR)
+			continue;
+		else
+		{
+			printf(EOS "\r\n[write failed]\r\n");
+			exit(1);
+		}
+	}
+}
+
+/* Write pkt to fd. Exit on failure. */
+void
+write_packet_or_fail(int fd, const struct packet *pkt)
+{
+	while (1)
+	{
+		ssize_t ret = write(fd, pkt, sizeof(struct packet));
+
+		if (ret == sizeof(struct packet))
+			return;
+		else if (ret < 0 && errno == EINTR)
+			continue;
+		else
+		{
+			printf(EOS "\r\n[write failed]\r\n");
+			exit(1);
+		}
+	}
+}
+
 static void
 usage()
 {
diff --git a/master.c b/master.c
index 81aa6a9..75f4926 100644
--- a/master.c
+++ b/master.c
@@ -379,7 +379,7 @@ client_activity(struct client *p)
 		return;
 
 	/* Close the client on an error. */
-	if (len <= 0)
+	if (len != sizeof(struct packet))
 	{
 		close(p->fd);
 		if (p->next)
@@ -393,7 +393,7 @@ client_activity(struct client *p)
 	if (pkt.type == MSG_PUSH)
 	{
 		if (pkt.len <= sizeof(pkt.u.buf))
-			write(the_pty.fd, pkt.u.buf, pkt.len);
+			write_buf_or_fail(the_pty.fd, pkt.u.buf, pkt.len);
 	}
 
 	/* Attach or detach from the program. */
@@ -434,7 +434,7 @@ client_activity(struct client *p)
                 	if (((the_pty.term.c_lflag & (ECHO|ICANON)) == 0) &&
                         	(the_pty.term.c_cc[VMIN] == 1))
 			{
-				write(the_pty.fd, &c, 1);
+				write_buf_or_fail(the_pty.fd, &c, 1);
 			}
 		}
 		/* Send a WINCH signal to the program. */
@@ -672,7 +672,7 @@ master_main(char **argv, int waitattach, int dontfork)
 		len = read(fd[0], buf, sizeof(buf));
 		if (len > 0)
 		{
-			write(2, buf, len);
+			write_buf_or_fail(2, buf, len);
 			kill(pid, SIGTERM);
 			return 1;
 		}