script: correctly handle buffered "tail" of output. +35 bytes.

1 files changed, 52 insertions, 60 deletions

diff --git a/util-linux/script.c b/util-linux/script.c
index e6dbb1a..a4aaa42 100644
--- a/util-linux/script.c
+++ b/util-linux/script.c
@@ -13,67 +13,38 @@
 
 #include "libbb.h"
 
-struct globals {
-	int child_pid;
-	int attr_ok; /* NB: 0: ok */
-	struct termios tt;
-	const char *fname;
-};
-#define G (*ptr_to_globals)
-#define child_pid (G.child_pid)
-#define attr_ok   (G.attr_ok  )
-#define tt        (G.tt       )
-#define fname     (G.fname    )
-#define INIT_G() do { \
-	SET_PTR_TO_GLOBALS(xzalloc(sizeof(G))); \
-	fname = "typescript"; \
-} while (0)
-
-static void done(void)
-{
-	if (child_pid) { /* we are parent */
-		if (attr_ok == 0)
-			tcsetattr(0, TCSAFLUSH, &tt);
-		if (!(option_mask32 & 8)) /* not -q */
-			printf("Script done, file is %s\n", fname);
-	}
-	exit(0);
-}
+static smallint fd_count = 2;
 
-#ifdef UNUSED
 static void handle_sigchld(int sig)
 {
-	/* wait for the exited child and exit */
-	while (wait_any_nohang(&sig) > 0)
-		continue;
-	done();
+	fd_count = 0;
 }
-#endif
-
-#if ENABLE_GETOPT_LONG
-static const char getopt_longopts[] ALIGN1 =
-	"append\0"  No_argument       "a"
-	"command\0" Required_argument "c"
-	"flush\0"   No_argument       "f"
-	"quiet\0"   No_argument       "q"
-	;
-#endif
 
 int script_main(int argc, char *argv[]) MAIN_EXTERNALLY_VISIBLE;
 int script_main(int argc, char *argv[])
 {
-	int opt, pty;
-	int winsz_ok;
+	int opt;
 	int mode;
-	struct termios rtt;
+	int child_pid;
+	int attr_ok; /* NB: 0: ok */
+	int winsz_ok;
+	int pty;
+	char pty_line[32];
+	struct termios tt, rtt;
 	struct winsize win;
-	char line[32];
+	const char *fname = "typescript";
 	const char *shell;
 	char shell_opt[] = "-i";
 	char *shell_arg = NULL;
 
-	INIT_G();
 #if ENABLE_GETOPT_LONG
+	static const char getopt_longopts[] ALIGN1 =
+		"append\0"  No_argument       "a"
+		"command\0" Required_argument "c"
+		"flush\0"   No_argument       "f"
+		"quiet\0"   No_argument       "q"
+		;
+
 	applet_long_options = getopt_longopts;
 #endif
 	opt_complementary = "?1"; /* max one arg */
@@ -95,10 +66,10 @@ int script_main(int argc, char *argv[])
 	}
 	shell = getenv("SHELL");
 	if (shell == NULL) {
-		shell = _PATH_BSHELL;
+		shell = DEFAULT_SHELL;
 	}
 
-	pty = getpty(line, sizeof(line));
+	pty = getpty(pty_line, sizeof(pty_line));
 	if (pty < 0) {
 		bb_perror_msg_and_die("can't get pty");
 	}
@@ -112,9 +83,12 @@ int script_main(int argc, char *argv[])
 	rtt.c_lflag &= ~ECHO;
 	tcsetattr(0, TCSAFLUSH, &rtt);
 
-	/* We exit as soon as child exits */
-	//signal(SIGCHLD, handle_sigchld);
-	signal(SIGCHLD, (void (*)(int)) done);
+	/* "script" from util-linux exits when child exits,
+	 * we wouldn't wait for EOF from slave pty
+	 * (output may be produced by grandchildren of child) */
+	signal(SIGCHLD, handle_sigchld);
+
+	/* TODO: SIGWINCH? pass window size changes down to slave? */
 
 	child_pid = vfork();
 	if (child_pid < 0) {
@@ -123,11 +97,10 @@ int script_main(int argc, char *argv[])
 
 	if (child_pid) {
 		/* parent */
-		char buf[256];
+#define buf bb_common_bufsiz1
 		struct pollfd pfd[2];
-		int outfd;
-		int fd_count = 2;
 		struct pollfd *ppfd = pfd;
+		int outfd, count, loop;
 
 		outfd = xopen(fname, mode);
 		pfd[0].fd = 0;
@@ -135,14 +108,16 @@ int script_main(int argc, char *argv[])
 		pfd[1].fd = pty;
 		pfd[1].events = POLLIN;
 		ndelay_on(pty); /* this descriptor is not shared, can do this */
-		/* ndelay_on(0); - NO, stdin can be shared! */
+		/* ndelay_on(0); - NO, stdin can be shared! Pity :( */
 
 		/* copy stdin to pty master input,
 		 * copy pty master output to stdout and file */
 		/* TODO: don't use full_write's, use proper write buffering */
-		while (fd_count && safe_poll(ppfd, fd_count, -1) > 0) {
+		while (fd_count) {
+			/* not safe_poll! we want SIGCHLD to EINTR poll */
+			poll(ppfd, fd_count, -1);
 			if (pfd[0].revents) {
-				int count = safe_read(0, buf, sizeof(buf));
+				count = safe_read(0, buf, sizeof(buf));
 				if (count <= 0) {
 					/* err/eof: don't read anymore */
 					pfd[0].revents = 0;
@@ -153,7 +128,6 @@ int script_main(int argc, char *argv[])
 				}
 			}
 			if (pfd[1].revents) {
-				int count;
 				errno = 0;
 				count = safe_read(pty, buf, sizeof(buf));
 				if (count <= 0 && errno != EAGAIN) {
@@ -170,14 +144,32 @@ int script_main(int argc, char *argv[])
 				}
 			}
 		}
-		done(); /* does not return */
+		/* If loop was exited because SIGCHLD handler set fd_count to 0,
+		 * there still can be some buffered output. But not loop forever:
+		 * we won't pump orphaned grandchildren's output indefinitely.
+		 * Testcase: running this in script:
+		 *      exec dd if=/dev/zero bs=1M count=1
+		 * must have "1+0 records in, 1+0 records out" captured too.
+		 * (util-linux's script doesn't do this. buggy :) */
+		loop = 999;
+		/* pty is in O_NONBLOCK mode, we exit as soon as buffer is empty */
+		while (--loop && (count = safe_read(pty, buf, sizeof(buf))) > 0) {
+			full_write(1, buf, count);
+			full_write(outfd, buf, count);
+		}
+
+		if (attr_ok == 0)
+			tcsetattr(0, TCSAFLUSH, &tt);
+		if (!(opt & 8)) /* not -q */
+			printf("Script done, file is %s\n", fname);
+		return EXIT_SUCCESS;
 	}
 
 	/* child: make pty slave to be input, output, error; run shell */
 	close(pty); /* close pty master */
 	/* open pty slave to fd 0,1,2 */
 	close(0);               
-	xopen(line, O_RDWR); /* uses fd 0 */
+	xopen(pty_line, O_RDWR); /* uses fd 0 */
 	xdup2(0, 1);
 	xdup2(0, 2);
 	/* copy our original stdin tty's parameters to pty */