script: correctly handle buffered "tail" of output. +35 bytes.
authorDenis Vlasenko <vda.linux@googlemail.com>
Thu, 28 Feb 2008 10:10:10 +0000 (10:10 -0000)
committerDenis Vlasenko <vda.linux@googlemail.com>
Thu, 28 Feb 2008 10:10:10 +0000 (10:10 -0000)
util-linux/script.c

index e6dbb1aa2d5665fabba34441210a1a36659077c7..a4aaa422f7f10e932e273bb3596e366f7b7ac1c2 100644 (file)
 
 #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 */