Fix up stdin handling

2 files changed, 45 insertions, 24 deletions

diff --git a/qarsh.c b/qarsh.c
index b5c6e58..87240a7 100644
--- a/qarsh.c
+++ b/qarsh.c
@@ -201,12 +201,6 @@ run_remote_cmd(char *cmdline)
 	/* Setup signal handling stuff so we can propogate signals */
 	setup_signals();
 
-	if (fcntl(fileno(stdin), F_SETFL, O_NONBLOCK) != 0) {
-		fprintf(stderr, 
-			"fcntl stdin O_NONBLOCK failed, %d: %s\n", 
-			errno, strerror(errno));
-	}
-	
 	/* Tell qarshd how much data it can send on stdout and stderr */
 	qp = make_qp_data_allow(1, QARSH_BUFSIZE);
 	send_packet(qarsh_fd, qp);
@@ -228,13 +222,25 @@ run_remote_cmd(char *cmdline)
 		FD_ZERO(&rfds);
 		FD_ZERO(&wfds);
 		FD_SET(qarsh_fd, &rfds);
-		/* Ignore stdin if it's a tty, we don't want to deal with
-		 * tty i/o if we're a background process */
-		if (!isatty(fileno(stdin)) && allowed_in) { FD_SET(fileno(stdin), &rfds); }
+		/* allowed_in is set to -1 when stdin is closed */
+		if (allowed_in != -1) {
+			/* close stdin if it is a tty */
+			if (isatty(fileno(stdin))) {
+				qp = make_qp_data(0, 0, 0, NULL);
+				send_packet(qarsh_fd, qp);
+				qpfree(qp);
+				close(fileno(stdin));
+				allowed_in = -1;
+			} else {
+				/* Only set stdin if we're allowed to send some data, otherwise we
+				 * could read 0 bytes and close stdin before we actually get eof */
+				if (allowed_in > 0) FD_SET(fileno(stdin), &rfds);
+			}
+		}
 		if (z_out) { FD_SET(fileno(stdout), &wfds); }
 		if (z_err) { FD_SET(fileno(stderr), &wfds); }
 
-		nset = pselect(FD_SETSIZE, &rfds, &wfds, NULL, &timeout, 
+		nset = pselect(qarsh_fd+1, &rfds, &wfds, NULL, &timeout, 
 				&pselect_sigmask);
 
 		/* Timeout hit, send a heartbeat */
@@ -264,13 +270,15 @@ run_remote_cmd(char *cmdline)
 			
 			if (nset && FD_ISSET(fileno(stdin), &rfds)) {
 				nbytes = read(fileno(stdin), buf, allowed_in);
-				qp = make_qp_data(0, 0, nbytes, buf);
-				send_packet(qarsh_fd, qp);
-				qpfree(qp);
-				allowed_in -= nbytes;
-				if (nbytes == 0) {
-					close(fileno(stdin));
-					allowed_in = 0;
+				if (nbytes >= 0) {
+					qp = make_qp_data(0, 0, nbytes, buf);
+					send_packet(qarsh_fd, qp);
+					qpfree(qp);
+					allowed_in -= nbytes;
+					if (nbytes == 0) {
+						close(fileno(stdin));
+						allowed_in = -1;
+					}
 				}
 				nset--;
 			}
@@ -334,7 +342,8 @@ run_remote_cmd(char *cmdline)
 					}
 				} else if (qp->qp_type == QP_DALLOW) {
 					if (qp->qp_dallow.qp_remfd == 0) {
-						allowed_in += qp->qp_dallow.qp_count;
+						/* If we already closed stdin, don't change allowed_in */
+						if (allowed_in != -1) allowed_in += qp->qp_dallow.qp_count;
 					} else {
 						fprintf(stderr, "ERROR: Received data allow for fd %d\n",
 								qp->qp_dallow.qp_remfd);
diff --git a/qarshd.c b/qarshd.c
index 915098c..7f0f944 100644
--- a/qarshd.c
+++ b/qarshd.c
@@ -283,9 +283,10 @@ handle_qarsh()
 	int child_status;
 	struct qa_packet *qp = NULL, *rp = NULL;
 	int allowed_out = 0, allowed_err = 0; /* number of bytes we can send to client */
-	char buf[4096], buf_in[4096];
-	int eof_in = 0;
-	int z_in = 0;
+	char buf[4096];		/* short term buffer for stdout and stderr */
+	char buf_in[4096];	/* long term buffer for stdin */
+	int z_in = 0;		/* Number of bytes in stdin buffer */
+	int eof_in = 0;		/* Have we seen EOF on stdin yet? */
 	int nbytes;
 
 	sigemptyset(&sigmask);
@@ -367,9 +368,20 @@ handle_qarsh()
 								qp->qp_data.qp_remfd, qp->qp_data.qp_count, QARSHD_BUFSIZE - z_in);
 						break;
 					}
-					if (qp->qp_data.qp_count == 0) eof_in = 1;
-					memcpy(buf_in+z_in, qp->qp_data.qp_blob, qp->qp_data.qp_count);
-					z_in += qp->qp_data.qp_count;
+					if (eof_in) {
+						syslog(LOG_ERR, "Received data on stdin after EOF\n");
+						break;
+					}
+					if (qp->qp_data.qp_count == 0) {
+						eof_in = 1;
+						if (z_in == 0) {
+							close(childfds[0]);
+							childfds[0] = -1;
+						}
+					} else {
+						memcpy(buf_in+z_in, qp->qp_data.qp_blob, qp->qp_data.qp_count);
+						z_in += qp->qp_data.qp_count;
+					}
 					break;
 				case QP_DALLOW:
 					if (qp->qp_dallow.qp_remfd == 1) {