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    Returns true if processing may continue, false if we're all
  81    done. */
  82 static bool
  83 handle_error (ssize_t retval, int *fd, bool fd_is_pty, const char *call)
  84 {
  85   if (fd_is_pty)
  86     {
  87       if (retval < 0)
  88         {
  89           if (errno == EIO)
  90             {
  91               /* Slave side of pty has been closed. */
  92               return false;
  93             }
  94           else
  95             fail_io (call);
  96         }
  97       else
  98         return true;
  99     }
 100   else
 101     {
 102       if (retval == 0)
 103         {
 104           close (*fd);
 105           *fd = -1;
 106           return true;
 107         }
 108       else
 109         fail_io (call);
 110     }
 111 }
 112
 113 /* Copies data from stdin to PTY and from PTY to stdout until no
 114    more data can be read or written. */
 115 static void
 116 relay (int pty, int dead_child_fd)
 117 {
 118   struct pipe
 119     {
 120       int in, out;
 121       char buf[BUFSIZ];
 122       size_t size, ofs;
 123       bool active;
 124     };
 125   struct pipe pipes[2];
 126
 127   /* Make PTY, stdin, and stdout non-blocking. */
 128   make_nonblocking (pty, true);
 129   make_nonblocking (STDIN_FILENO, true);
 130   make_nonblocking (STDOUT_FILENO, true);
 131
 132   /* Configure noncanonical mode on PTY and stdin to avoid
 133      waiting for end-of-line.  We want to minimize context
 134      switching on PTY (for efficiency) and minimize latency on
 135      stdin to avoid a laggy user experience. */
 136   make_noncanon (pty, 16, 1);
 137   make_noncanon (STDIN_FILENO, 1, 0);
 138
 139   memset (pipes, 0, sizeof pipes);
 140   pipes[0].in = STDIN_FILENO;
 141   pipes[0].out = pty;
 142   pipes[1].in = pty;
 143   pipes[1].out = STDOUT_FILENO;
 144
 145   while (pipes[0].in != -1 || pipes[1].in != -1)
 146     {
 147       fd_set read_fds, write_fds;
 148       int retval;
 149       int i;
 150
 151       FD_ZERO (&read_fds);
 152       FD_ZERO (&write_fds);
 153       for (i = 0; i < 2; i++)
 154         {
 155           struct pipe *p = &pipes[i];
 156
 157           /* Don't do anything with the stdin->pty pipe until we
 158              have some data for the pty->stdout pipe.  If we get
 159              too eager, Bochs will throw away our input. */
 160           if (i == 0 && !pipes[1].active)
 161             continue;
 162
 163           if (p->in != -1 && p->size + p->ofs < sizeof p->buf)
 164             FD_SET (p->in, &read_fds);
 165           if (p->out != -1 && p->size > 0)
 166             FD_SET (p->out, &write_fds);
 167         }
 168       FD_SET (dead_child_fd, &read_fds);
 169
 170       do
 171         {
 172           retval = select (FD_SETSIZE, &read_fds, &write_fds, NULL, NULL);
 173         }
 174       while (retval < 0 && errno == EINTR);
 175       if (retval < 0)
 176         fail_io ("select");
 177
 178       if (FD_ISSET (dead_child_fd, &read_fds))
 179         {
 180           /* Child died.  Do final relaying. */
 181           struct pipe *p = &pipes[1];
 182           if (p->out == -1)
 183             return;
 184           make_nonblocking (STDOUT_FILENO, false);
 185           for (;;)
 186             {
 187               ssize_t n;
 188
 189               /* Write buffer. */
 190               while (p->size > 0)
 191                 {
 192                   n = write (p->out, p->buf + p->ofs, p->size);
 193                   if (n < 0)
 194                     fail_io ("write");
 195                   else if (n == 0)
 196                     fail_io ("zero-length write");
 197                   p->ofs += n;
 198                   p->size -= n;
 199                 }
 200               p->ofs = 0;
 201
 202               p->size = n = read (p->in, p->buf, sizeof p->buf);
 203               if (n <= 0)
 204                 return;
 205             }
 206         }
 207
 208       for (i = 0; i < 2; i++)
 209         {
 210           struct pipe *p = &pipes[i];
 211           if (p->in != -1 && FD_ISSET (p->in, &read_fds))
 212             {
 213               ssize_t n = read (p->in, p->buf + p->ofs + p->size,
 214                                 sizeof p->buf - p->ofs - p->size);
 215               if (n > 0)
 216                 {
 217                   p->active = true;
 218                   p->size += n;
 219                   if (p->size == BUFSIZ && p->ofs != 0)
 220                     {
 221                       memmove (p->buf, p->buf + p->ofs, p->size);
 222                       p->ofs = 0;
 223                     }
 224                 }
 225               else if (!handle_error (n, &p->in, p->in == pty, "read"))
 226                 return;
 227             }
 228           if (p->out != -1 && FD_ISSET (p->out, &write_fds))
 229             {
 230               ssize_t n = write (p->out, p->buf + p->ofs, p->size);
 231               if (n > 0)
 232                 {
 233                   p->ofs += n;
 234                   p->size -= n;
 235                   if (p->size == 0)
 236                     p->ofs = 0;
 237                 }
 238               else if (!handle_error (n, &p->out, p->out == pty, "write"))
 239                 return;
 240             }
 241         }
 242     }
 243 }
 244
 245 static int dead_child_fd;
 246
 247 static void
 248 sigchld_handler (int signo __attribute__ ((unused)))
 249 {
 250   if (write (dead_child_fd, "", 1) < 0)
 251     _exit (1);
 252 }
 253
 254 int
 255 main (int argc __attribute__ ((unused)), char *argv[])
 256 {
 257   int master, slave;
 258   char *name;
 259   pid_t pid;
 260   struct sigaction sa;
 261   int pipe_fds[2];
 262   struct itimerval zero_itimerval, old_itimerval;
 263
 264   if (argc < 2)
 265     {
 266       fprintf (stderr,
 267                "usage: squish-pty COMMAND [ARG]...\n"
 268                "Squishes both stdin and stdout into a single pseudoterminal,\n"
 269                "which is passed as stdout to run the specified COMMAND.\n");
 270       return EXIT_FAILURE;
 271     }
 272
 273   /* Open master side of pty and get ready to open slave. */
 274   master = open ("/dev/ptmx", O_RDWR | O_NOCTTY);
 275   if (master < 0)
 276     fail_io ("open \"/dev/ptmx\"");
 277   if (grantpt (master) < 0)
 278     fail_io ("grantpt");
 279   if (unlockpt (master) < 0)
 280     fail_io ("unlockpt");
 281
 282   /* Open slave side of pty. */
 283   name = ptsname (master);
 284   if (name == NULL)
 285     fail_io ("ptsname");
 286   slave = open (name, O_RDWR);
 287   if (slave < 0)
 288     fail_io ("open \"%s\"", name);
 289
 290   /* System V implementations need STREAMS configuration for the
 291      slave. */
 292   if (isastream (slave))
 293     {
 294       if (ioctl (slave, I_PUSH, "ptem") < 0
 295           || ioctl (slave, I_PUSH, "ldterm") < 0)
 296         fail_io ("ioctl");
 297     }
 298
 299   /* Arrange to get notified when a child dies, by writing a byte
 300      to a pipe fd.  We really want to use pselect() and
 301      sigprocmask(), but Solaris 2.7 doesn't have it. */
 302   if (pipe (pipe_fds) < 0)
 303     fail_io ("pipe");
 304   dead_child_fd = pipe_fds[1];
 305
 306   memset (&sa, 0, sizeof sa);
 307   sa.sa_handler = sigchld_handler;
 308   sigemptyset (&sa.sa_mask);
 309   sa.sa_flags = SA_RESTART;
 310   if (sigaction (SIGCHLD, &sa, NULL) < 0)
 311     fail_io ("sigaction");
 312
 313   /* Save the virtual interval timer, which might have been set
 314      by the process that ran us.  It really should be applied to
 315      our child process. */
 316   memset (&zero_itimerval, 0, sizeof zero_itimerval);
 317   if (setitimer (ITIMER_VIRTUAL, &zero_itimerval, &old_itimerval) < 0)
 318     fail_io ("setitimer");
 319
 320   pid = fork ();
 321   if (pid < 0)
 322     fail_io ("fork");
 323   else if (pid != 0)
 324     {
 325       /* Running in parent process. */
 326       int status;
 327       close (slave);
 328       relay (master, pipe_fds[0]);
 329
 330       /* If the subprocess has died, die in the same fashion.
 331          In particular, dying from SIGVTALRM tells the pintos
 332          script that we ran out of CPU time. */
 333       if (waitpid (pid, &status, WNOHANG) > 0)
 334         {
 335           if (WIFEXITED (status))
 336             return WEXITSTATUS (status);
 337           else if (WIFSIGNALED (status))
 338             raise (WTERMSIG (status));
 339         }
 340       return 0;
 341     }
 342   else
 343     {
 344       /* Running in child process. */
 345       if (setitimer (ITIMER_VIRTUAL, &old_itimerval, NULL) < 0)
 346         fail_io ("setitimer");
 347       if (dup2 (slave, STDOUT_FILENO) < 0)
 348         fail_io ("dup2");
 349       if (close (pipe_fds[0]) < 0 || close (pipe_fds[1]) < 0
 350           || close (slave) < 0 || close (master) < 0)
 351         fail_io ("close");
 352       execvp (argv[1], argv + 1);
 353       fail_io ("exec");
 354     }
 355 }