pintos-os.org Git - pintos-anon/blob - src/utils/squish-pty.c

   1 #define _GNU_SOURCE 1
   2 #include <errno.h>
   3 #include <fcntl.h>
   4 #include <signal.h>
   5 #include <stdarg.h>
   6 #include <stdbool.h>
   7 #include <stdio.h>
   8 #include <stdlib.h>
   9 #include <string.h>
  10 #include <stropts.h>
  11 #include <sys/ioctl.h>
  12 #include <sys/stat.h>
  13 #include <sys/time.h>
  14 #include <sys/types.h>
  15 #include <sys/wait.h>
  16 #include <termios.h>
  17 #include <unistd.h>
  18
  19 static void
  20 fail_io (const char *msg, ...)
  21      __attribute__ ((noreturn))
  22      __attribute__ ((format (printf, 1, 2)));
  23
  24 /* Prints MSG, formatting as with printf(),
  25    plus an error message based on errno,
  26    and exits. */
  27 static void
  28 fail_io (const char *msg, ...)
  29 {
  30   va_list args;
  31
  32   va_start (args, msg);
  33   vfprintf (stderr, msg, args);
  34   va_end (args);
  35
  36   if (errno != 0)
  37     fprintf (stderr, ": %s", strerror (errno));
  38   putc ('\n', stderr);
  39   exit (EXIT_FAILURE);
  40 }
  41
  42 /* If FD is a terminal, configures it for noncanonical input mode
  43    with VMIN and VTIME set as indicated.
  44    If FD is not a terminal, has no effect. */
  45 static void
  46 make_noncanon (int fd, int vmin, int vtime)
  47 {
  48   if (isatty (fd))
  49     {
  50       struct termios termios;
  51       if (tcgetattr (fd, &termios) < 0)
  52         fail_io ("tcgetattr");
  53       termios.c_lflag &= ~(ICANON | ECHO);
  54       termios.c_cc[VMIN] = vmin;
  55       termios.c_cc[VTIME] = vtime;
  56       if (tcsetattr (fd, TCSANOW, &termios) < 0)
  57         fail_io ("tcsetattr");
  58     }
  59 }
  60
  61 /* Make FD non-blocking if NONBLOCKING is true,
  62    or blocking if NONBLOCKING is false. */
  63 static void
  64 make_nonblocking (int fd, bool nonblocking)
  65 {
  66   int flags = fcntl (fd, F_GETFL);
  67   if (flags < 0)
  68     fail_io ("fcntl");
  69   if (nonblocking)
  70     flags |= O_NONBLOCK;
  71   else
  72     flags &= ~O_NONBLOCK;
  73   if (fcntl (fd, F_SETFL, flags) < 0)
  74     fail_io ("fcntl");
  75 }
  76
  77 /* Handle a read or write on *FD, which is the pty if FD_IS_PTY
  78    is true, that returned end-of-file or error indication RETVAL.
  79    The system call is named CALL, for use in error messages.
  80    Sets *FD to -1 if the fd is no longer readable or writable. */
  81 static void
  82 handle_error (ssize_t retval, int *fd, bool fd_is_pty, const char *call)
  83 {
  84   if (fd_is_pty)
  85     {
  86       if (retval < 0)
  87         {
  88           if (errno == EIO)
  89             {
  90               /* Slave side of pty has been closed. */
  91               *fd = -1;
  92             }
  93           else
  94             fail_io (call);
  95         }
  96     }
  97   else
  98     {
  99       if (retval == 0)
 100         {
 101           close (*fd);
 102           *fd = -1;
 103         }
 104       else
 105         fail_io (call);
 106     }
 107 }
 108
 109 /* Copies data from stdin to PTY and from PTY to stdout until no
 110    more data can be read or written. */
 111 static void
 112 relay (int pty, int dead_child_fd)
 113 {
 114   struct pipe
 115     {
 116       int in, out;
 117       char buf[BUFSIZ];
 118       size_t size, ofs;
 119       bool active;
 120     };
 121   struct pipe pipes[2];
 122
 123   /* Make PTY, stdin, and stdout non-blocking. */
 124   make_nonblocking (pty, true);
 125   make_nonblocking (STDIN_FILENO, true);
 126   make_nonblocking (STDOUT_FILENO, true);
 127
 128   /* Configure noncanonical mode on PTY and stdin to avoid
 129      waiting for end-of-line.  We want to minimize context
 130      switching on PTY (for efficiency) and minimize latency on
 131      stdin to avoid a laggy user experience. */
 132   make_noncanon (pty, 16, 1);
 133   make_noncanon (STDIN_FILENO, 1, 0);
 134
 135   memset (pipes, 0, sizeof pipes);
 136   pipes[0].in = STDIN_FILENO;
 137   pipes[0].out = pty;
 138   pipes[1].in = pty;
 139   pipes[1].out = STDOUT_FILENO;
 140
 141   while (pipes[1].in != -1)
 142     {
 143       fd_set read_fds, write_fds;
 144       int retval;
 145       int i;
 146
 147       FD_ZERO (&read_fds);
 148       FD_ZERO (&write_fds);
 149       for (i = 0; i < 2; i++)
 150         {
 151           struct pipe *p = &pipes[i];
 152
 153           /* Don't do anything with the stdin->pty pipe until we
 154              have some data for the pty->stdout pipe.  If we get
 155              too eager, Bochs will throw away our input. */
 156           if (i == 0 && !pipes[1].active)
 157             continue;
 158
 159           if (p->in != -1 && p->size + p->ofs < sizeof p->buf)
 160             FD_SET (p->in, &read_fds);
 161           if (p->out != -1 && p->size > 0)
 162             FD_SET (p->out, &write_fds);
 163         }
 164       FD_SET (dead_child_fd, &read_fds);
 165
 166       do
 167         {
 168           retval = select (FD_SETSIZE, &read_fds, &write_fds, NULL, NULL);
 169         }
 170       while (retval < 0 && errno == EINTR);
 171       if (retval < 0)
 172         fail_io ("select");
 173
 174       if (FD_ISSET (dead_child_fd, &read_fds))
 175         break;
 176
 177       for (i = 0; i < 2; i++)
 178         {
 179           struct pipe *p = &pipes[i];
 180           if (p->in != -1 && FD_ISSET (p->in, &read_fds))
 181             {
 182               ssize_t n = read (p->in, p->buf + p->ofs + p->size,
 183                                 sizeof p->buf - p->ofs - p->size);
 184               if (n > 0)
 185                 {
 186                   p->active = true;
 187                   p->size += n;
 188                   if (p->size == BUFSIZ && p->ofs != 0)
 189                     {
 190                       memmove (p->buf, p->buf + p->ofs, p->size);
 191                       p->ofs = 0;
 192                     }
 193                 }
 194               else
 195                 handle_error (n, &p->in, p->in == pty, "read");
 196             }
 197           if (p->out != -1 && FD_ISSET (p->out, &write_fds))
 198             {
 199               ssize_t n = write (p->out, p->buf + p->ofs, p->size);
 200               if (n > 0)
 201                 {
 202                   p->ofs += n;
 203                   p->size -= n;
 204                   if (p->size == 0)
 205                     p->ofs = 0;
 206                 }
 207               else
 208                 handle_error (n, &p->out, p->out == pty, "write");
 209             }
 210         }
 211     }
 212
 213     if (pipes[1].out == -1)
 214       return;
 215
 216     make_nonblocking (STDOUT_FILENO, false);
 217     for (;;)
 218       {
 219         struct pipe *p = &pipes[1];
 220         ssize_t n;
 221
 222         /* Write buffer. */
 223         while (p->size > 0)
 224           {
 225             n = write (p->out, p->buf + p->ofs, p->size);
 226             if (n < 0)
 227               fail_io ("write");
 228             else if (n == 0)
 229               fail_io ("zero-length write");
 230             p->ofs += n;
 231             p->size -= n;
 232           }
 233         p->ofs = 0;
 234
 235         p->size = n = read (p->in, p->buf, sizeof p->buf);
 236         if (n <= 0)
 237           return;
 238       }
 239 }
 240
 241 static int dead_child_fd;
 242
 243 static void
 244 sigchld_handler (int signo __attribute__ ((unused)))
 245 {
 246   if (write (dead_child_fd, "", 1) < 0)
 247     _exit (1);
 248 }
 249
 250 int
 251 main (int argc __attribute__ ((unused)), char *argv[])
 252 {
 253   int master, slave;
 254   char *name;
 255   pid_t pid;
 256   struct sigaction sa;
 257   int pipe_fds[2];
 258   struct itimerval zero_itimerval, old_itimerval;
 259
 260   if (argc < 2)
 261     {
 262       fprintf (stderr,
 263                "usage: squish-pty COMMAND [ARG]...\n"
 264                "Squishes both stdin and stdout into a single pseudoterminal,\n"
 265                "which is passed as stdout to run the specified COMMAND.\n");
 266       return EXIT_FAILURE;
 267     }
 268
 269   /* Open master side of pty and get ready to open slave. */
 270   master = open ("/dev/ptmx", O_RDWR | O_NOCTTY);
 271   if (master < 0)
 272     fail_io ("open \"/dev/ptmx\"");
 273   if (grantpt (master) < 0)
 274     fail_io ("grantpt");
 275   if (unlockpt (master) < 0)
 276     fail_io ("unlockpt");
 277
 278   /* Open slave side of pty. */
 279   name = ptsname (master);
 280   if (name == NULL)
 281     fail_io ("ptsname");
 282   slave = open (name, O_RDWR);
 283   if (slave < 0)
 284     fail_io ("open \"%s\"", name);
 285
 286   /* System V implementations need STREAMS configuration for the
 287      slave. */
 288   if (isastream (slave))
 289     {
 290       if (ioctl (slave, I_PUSH, "ptem") < 0
 291           || ioctl (slave, I_PUSH, "ldterm") < 0)
 292         fail_io ("ioctl");
 293     }
 294
 295   /* Arrange to get notified when a child dies, by writing a byte
 296      to a pipe fd.  We really want to use pselect() and
 297      sigprocmask(), but Solaris 2.7 doesn't have it. */
 298   if (pipe (pipe_fds) < 0)
 299     fail_io ("pipe");
 300   dead_child_fd = pipe_fds[1];
 301
 302   memset (&sa, 0, sizeof sa);
 303   sa.sa_handler = sigchld_handler;
 304   sigemptyset (&sa.sa_mask);
 305   sa.sa_flags = SA_RESTART;
 306   if (sigaction (SIGCHLD, &sa, NULL) < 0)
 307     fail_io ("sigaction");
 308
 309   /* Save the virtual interval timer, which might have been set
 310      by the process that ran us.  It really should be applied to
 311      our child process. */
 312   memset (&zero_itimerval, 0, sizeof zero_itimerval);
 313   if (setitimer (ITIMER_VIRTUAL, &zero_itimerval, &old_itimerval) < 0)
 314     fail_io ("setitimer");
 315
 316   pid = fork ();
 317   if (pid < 0)
 318     fail_io ("fork");
 319   else if (pid != 0)
 320     {
 321       /* Running in parent process. */
 322       int status;
 323       close (slave);
 324       relay (master, pipe_fds[0]);
 325
 326       /* If the subprocess has died, die in the same fashion.
 327          In particular, dying from SIGVTALRM tells the pintos
 328          script that we ran out of CPU time. */
 329       if (waitpid (pid, &status, WNOHANG) > 0)
 330         {
 331           if (WIFEXITED (status))
 332             return WEXITSTATUS (status);
 333           else if (WIFSIGNALED (status))
 334             raise (WTERMSIG (status));
 335         }
 336       return 0;
 337     }
 338   else
 339     {
 340       /* Running in child process. */
 341       if (setitimer (ITIMER_VIRTUAL, &old_itimerval, NULL) < 0)
 342         fail_io ("setitimer");
 343       if (dup2 (slave, STDOUT_FILENO) < 0)
 344         fail_io ("dup2");
 345       if (close (pipe_fds[0]) < 0 || close (pipe_fds[1]) < 0
 346           || close (slave) < 0 || close (master) < 0)
 347         fail_io ("close");
 348       execvp (argv[1], argv + 1);
 349       fail_io ("exec");
 350     }
 351 }