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 e6dbb1aa2..a4aaa422f 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 */