Fix up stdin handling

1 files changed, 27 insertions, 18 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);