Get commands running over one socket

Added a new packet to limit data sent from the other side.

1 files changed, 203 insertions, 51 deletions

diff --git a/qarshd.c b/qarshd.c
index 91e2501..6cb52e0 100644
--- a/qarshd.c
+++ b/qarshd.c
@@ -41,6 +41,7 @@
 #include "sockutil.h"
 #include "qarsh_packet.h"
 
+#define QARSHD_BUFSIZE 4096
 /*
  * QA Remote Shell Daemon
  */
@@ -49,9 +50,12 @@ int debug = 0;
 int dopause = 0;
 
 /* Globals */
-struct sockaddr_storage peername;
+const int qinfd = 0; /* qarshd in file descriptor */
+const int qoutfd = 1; /* qarshd out file descriptor */
 int child_exitted = 0;
+pid_t child_pid;
 sigset_t orig_sigmask;
+int childfds[3] = { -1, -1, -1 }; /* pipes to child for stdin/stdout/stderr */
 
 /* A mini cache for rstat so we can check it in pushfile */
 char *saved_path = NULL;
@@ -98,10 +102,18 @@ sig_handler(int sig)
 }
 
 pid_t 
-run_cmd(const char *cmd, int p_in, int p_out, int p_err)
+run_cmd(const char *cmd)
 {
 	pid_t pid;
-	int new_in, new_out, new_err;
+	int pipefds[2], parentfds[3];
+
+	pipe(pipefds);
+	childfds[0] = pipefds[1]; parentfds[0] = pipefds[0]; /* stdin */
+	pipe(pipefds);
+	childfds[1] = pipefds[0]; parentfds[1] = pipefds[1]; /* stdout */
+	pipe(pipefds);
+	childfds[2] = pipefds[0]; parentfds[2] = pipefds[1]; /* stderr */
+
 
 	syslog(LOG_INFO, "Running cmdline: %s\n", cmd);
 	if ((pid = fork()) < 0) {
@@ -114,21 +126,22 @@ run_cmd(const char *cmd, int p_in, int p_out, int p_err)
 		setpgrp();
 		sigprocmask(SIG_SETMASK, &orig_sigmask, NULL);
 
-		/* Connect stdin, stdout, and stderr to qarsh */
-		new_in = connect_to_peer(&peername, p_in);
-		if (new_in == -1) syslog(LOG_WARNING, "connect to new_in failed");
-		dup2(new_in, fileno(stdin));
-		new_out = connect_to_peer(&peername, p_out);
-		if (new_out == -1) syslog(LOG_WARNING, "connect to new_out failed");
-		dup2(new_out, fileno(stdout));
-		new_err = connect_to_peer(&peername, p_err);
-		if (new_err == -1) syslog(LOG_WARNING, "connect to new_err failed");
-		dup2(new_err, fileno(stderr));
+		/* Connect stdin, stdout, and stderr to parent's pipes */
+		dup2(parentfds[0], fileno(stdin));
+		dup2(parentfds[1], fileno(stdout));
+		dup2(parentfds[2], fileno(stderr));
+		/* close end of pipes we're not using */
+		close(childfds[0]);
+		close(childfds[1]);
+		close(childfds[2]);
 
 		execlp("sh", "sh", "-c", cmd, NULL);
 		fprintf(stderr, "exec of %s failed: %d, %s\n", cmd, errno, strerror(errno));
 		exit(127);
 	} 
+	close(parentfds[0]);
+	close(parentfds[1]);
+	close(parentfds[2]);
 	return pid;
 }
 
@@ -225,7 +238,7 @@ pushfile(const char *path, int remfd)
 		} else {
 			qp = make_qp_data(remfd, offset, nbytes, buf);
 		}
-		send_packet(fileno(stdout), qp);
+		send_packet(qoutfd, qp);
 		offset += nbytes;
 		qpfree(qp);
 	} while (nbytes > 0);
@@ -259,19 +272,22 @@ rstat(const char *path)
 	return rp;
 }
 
+/* Handle only qarsh related packets */
 void
-handle_packets(int infd)
+handle_qarsh()
 {
-	fd_set rfds;
-	int nfd;
+	fd_set rfds, wfds;
+	int nfd, maxfd;
 	struct timespec timeout;
-	struct qa_packet *qp = NULL, *rp = NULL;
 	sigset_t sigmask;
 	struct sigaction sa;
-
-	pid_t child_pid = 0;
 	int child_status;
-	off_t nbytes;	
+	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;
+	int nbytes;
 
 	sigemptyset(&sigmask);
 	sigaddset(&sigmask, SIGCHLD);
@@ -280,28 +296,171 @@ handle_packets(int infd)
 	sa.sa_mask = sigmask;
 	sa.sa_flags = SA_RESTART;
 	sigaction(SIGCHLD, &sa, NULL);
-	if (dopause) {
-		signal(SIGALRM, sig_handler);
-		pause();
-		signal(SIGALRM, SIG_DFL);
-	}
 
+	qp = make_qp_data_allow(0, 4096);
+	send_packet(qoutfd, qp);
+	qpfree(qp);
+	
 	for (;;) {
-		FD_SET(infd, &rfds);
-		timeout.tv_sec = 3;
-		timeout.tv_nsec = 0;
 
-		nfd = pselect(infd+1, &rfds, NULL, NULL, &timeout, &orig_sigmask);
+		FD_ZERO(&rfds); FD_ZERO(&wfds);
+		FD_SET(qinfd, &rfds);
+		maxfd = qinfd;
+		if (childfds[0] > -1 && z_in) {
+			FD_SET(childfds[0], &wfds);
+			maxfd = childfds[0] > maxfd ? childfds[0] : maxfd;
+		}
+		if (childfds[1] > -1 && allowed_out) {
+			FD_SET(childfds[1], &rfds);
+			maxfd = childfds[1] > maxfd ? childfds[1] : maxfd;
+		}
+		if (childfds[2] > -1 && allowed_err) {
+			FD_SET(childfds[2], &rfds);
+			maxfd = childfds[2] > maxfd ? childfds[2] : maxfd;
+		}
+		nfd = pselect(maxfd+1, &rfds, &wfds, NULL, &timeout, &orig_sigmask);
 
 		if (child_exitted) {
 			waitpid(child_pid, &child_status, 0);
 			child_exitted--;
 			child_pid = 0;
 			rp = make_qp_cmdexit(child_pid, child_status);
-			send_packet(fileno(stdout), rp);
+			send_packet(qoutfd, rp);
 			qpfree(rp);
 		}
+		if (nfd < 0) {
+			if (errno == EINTR) {
+				/* signals handled above here */
+				continue;
+			} else {
+				syslog(LOG_ERR, "select errno %d, %s\n", 
+				       errno, strerror(errno));
+			}
+		} else if (nfd > 0) {
+			if (nfd && FD_ISSET(qinfd, &rfds)) {
+				qp = recv_packet(qinfd);
+				if (qp == NULL) {
+					if (debug) syslog(LOG_DEBUG, "That's enough\n");
+					break;
+				}
+				if (debug) dump_qp(qp);
+				switch(qp->qp_type) {
+				case QP_KILL:
+					if (child_pid) {
+						syslog(LOG_INFO, "Sending child %d signal %d", 
+							child_pid, qp->qp_kill.qp_sig);
+						kill(child_pid * -1, qp->qp_kill.qp_sig);
+					}
+					break;
+				case QP_RUNCMD:
+					syslog(LOG_ERR, "Only one command allowed per connection");
+					break;
+				case QP_DATA:
+					if (qp->qp_data.qp_remfd != 0) {
+						syslog(LOG_ERR, "Received data for fd %d\n", qp->qp_data.qp_remfd);
+						break;
+					}
+					if (qp->qp_data.qp_count > (QARSHD_BUFSIZE - z_in)) {
+						syslog(LOG_ERR, "Received too much data for fd %d, %d > %d\n",
+								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;
+					break;
+				case QP_DALLOW:
+					if (qp->qp_dallow.qp_remfd == 1) {
+						allowed_out += qp->qp_dallow.qp_count;
+					} else if (qp->qp_dallow.qp_remfd == 2) {
+						allowed_err += qp->qp_dallow.qp_count;
+					} else {
+						syslog(LOG_ERR, "Received data allow for fd %d\n", qp->qp_dallow.qp_remfd);
+					}
+					break;
+				default:
+					syslog(LOG_ERR, "Expected message type %s", qp_packet_type(qp->qp_type));
+					break;
+				}
+				nfd--;
+			}
+			if (nfd && FD_ISSET(childfds[0], &wfds)) {
+				/* Child is ready for data on stdin */
+				nbytes = write(childfds[0], buf_in, z_in);
+				if (nbytes == z_in) {
+					z_in = 0;
+					if (eof_in) {
+						close(childfds[0]);
+						childfds[0] = -1;
+					}
+				} else {
+					memmove(buf_in, buf_in+nbytes, z_in - nbytes);
+					z_in -= nbytes;
+				}
+				if (!eof_in) {
+					qp = make_qp_data_allow(0, nbytes);
+					send_packet(qoutfd, qp);
+					qpfree(qp);
+				}
+				nfd--;
+			}
+			if (nfd && FD_ISSET(childfds[1], &rfds)) {
+				/* Child has something to send to stdout */
+				nbytes = read(childfds[1], buf, allowed_out);
+				qp = make_qp_data(1, 0, nbytes, buf);
+				send_packet(qoutfd, qp);
+				qpfree(qp);
+				allowed_out -= nbytes;
+				if (nbytes == 0) {
+					close(childfds[1]);
+					childfds[1] = -1;
+					allowed_out = 0;
+				}
+				nfd--;
+			}
+			if (nfd && FD_ISSET(childfds[2], &rfds)) {
+				/* Child has something to send to stderr */
+				nbytes = read(childfds[2], buf, allowed_err);
+				qp = make_qp_data(2, 0, nbytes, buf);
+				send_packet(qoutfd, qp);
+				qpfree(qp);
+				allowed_err -= nbytes;
+				if (nbytes == 0) {
+					close(childfds[2]);
+					childfds[2] = -1;
+					allowed_err = 0;
+				}
+				nfd--;
+			}
+
+		} else {
+			if (debug) syslog(LOG_DEBUG, "Nothing to do in handle_qarsh\n");
+		}
+	}
+}
 
+void
+handle_packets()
+{
+	fd_set rfds;
+	int nfd;
+	struct timeval timeout;
+	struct qa_packet *qp = NULL, *rp = NULL;
+
+	off_t nbytes;	
+
+	if (dopause) {
+		signal(SIGALRM, sig_handler);
+		pause();
+		signal(SIGALRM, SIG_DFL);
+	}
+
+	for (;;) {
+		FD_SET(qinfd, &rfds);
+		timeout.tv_sec = 3;
+		timeout.tv_usec = 0;
+
+		nfd = select(qinfd+1, &rfds, NULL, NULL, &timeout);
 		if (nfd < 0) {
 			if (errno == EINTR) {
 				/* signals handled above here */
@@ -311,21 +470,13 @@ handle_packets(int infd)
 				       errno, strerror(errno));
 			}
 		} else if (nfd > 0) {
-			qp = recv_packet(infd);
+			qp = recv_packet(qinfd);
 			if (qp == NULL) {
 				if (debug) syslog(LOG_DEBUG, "That's enough\n");
 				break;
 			}
 			if (debug) dump_qp(qp);
 			switch (qp->qp_type) {
-			case QP_KILL:
-				if (child_pid) {
-
-					syslog(LOG_INFO, "Sending child %d signal %d", 
-						child_pid, qp->qp_kill.qp_sig);
-					kill(child_pid * -1, qp->qp_kill.qp_sig);
-				}
-				break;
 			case QP_SETUSER:
 				if (setup_user(qp->qp_setuser.qp_user,
 						qp->qp_setuser.qp_group) == 0) {
@@ -333,25 +484,24 @@ handle_packets(int infd)
 				} else {
 					rp = make_qp_ack(QP_SETUSER, 1);
 				}
-				send_packet(fileno(stdout), rp);
+				send_packet(qoutfd, rp);
 				qpfree(rp);
 				break;
 			case QP_RUNCMD:
-				child_pid = run_cmd(qp->qp_runcmd.qp_cmdline,
-						qp->qp_runcmd.qp_stdin_port,
-						qp->qp_runcmd.qp_stdout_port,
-						qp->qp_runcmd.qp_stderr_port);
+				child_pid = run_cmd(qp->qp_runcmd.qp_cmdline);
+				qpfree(qp);
+				handle_qarsh(); /* continue processing qarsh separately */
 				break;
 			case QP_RECVFILE: /* Setup file descriptors to handle incoming data */
 				rp = prepare_recvfile(qp);
-				send_packet(fileno(stdout), rp);
+				send_packet(qoutfd, rp);
 				qpfree(rp);
 				break;
 			case QP_DATA:
 				assert(qp->qp_data.qp_remfd == remotefds[0].fd);
 				rp = receive_data(&remotefds[0], qp);
 				if (rp) {
-					send_packet(fileno(stdout), rp);
+					send_packet(qoutfd, rp);
 					qpfree(rp);
 				}
 				break;
@@ -366,15 +516,16 @@ handle_packets(int infd)
 				} else {
 					rp = make_qp_returncode(0, 0, "Transfer Complete");
 				}
-				send_packet(fileno(stdout), rp);
+				send_packet(qoutfd, rp);
 				qpfree(rp);
 				break;
 			case QP_RSTAT:
 				syslog(LOG_INFO, "Got a QP_RSTAT with path = %s\n",
 				       qp->qp_rstat.qp_path);
 				rp = rstat(qp->qp_rstat.qp_path);
-				send_packet(fileno(stdout), rp);
+				send_packet(qoutfd, rp);
 				qpfree(rp);
+				qpfree(qp);
 				break;
 			default:
 				syslog(LOG_WARNING, 
@@ -383,7 +534,7 @@ handle_packets(int infd)
 			}
 			qpfree(qp);
 		} else {
-			if (debug) syslog(LOG_DEBUG, "Nothing to do\n");
+			if (debug) syslog(LOG_DEBUG, "Nothing to do in handle_packets\n");
 		}
 		
 	}
@@ -394,6 +545,7 @@ main(int argc, char *argv[])
 {
 	int ch;
 	socklen_t peerlen;
+	struct sockaddr_storage peername;
 	char peer_hoststr[NI_MAXHOST];
 	char peer_portstr[NI_MAXSERV];
 
@@ -424,7 +576,7 @@ main(int argc, char *argv[])
 		peername.ss_family == AF_INET ? "IPv4" : "IPv6");
 
 	/* Start reading packets from stdin */
-	handle_packets(0);
+	handle_packets();
 
 	return 0;
 }