pintos-os.org Git - pintos-anon/blob - src/utils/squish-unix.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 <stddef.h>
   8 #include <stdio.h>
   9 #include <stdlib.h>
  10 #include <string.h>
  11 #include <stropts.h>
  12 #include <sys/ioctl.h>
  13 #include <sys/stat.h>
  14 #include <sys/time.h>
  15 #include <sys/types.h>
  16 #include <sys/wait.h>
  17 #include <sys/socket.h>
  18 #include <sys/un.h>
  19 #include <termios.h>
  20 #include <unistd.h>
  21
  22 static void
  23 fail_io (const char *msg, ...)
  24      __attribute__ ((noreturn))
  25      __attribute__ ((format (printf, 1, 2)));
  26
  27 /* Prints MSG, formatting as with printf(),
  28    plus an error message based on errno,
  29    and exits. */
  30 static void
  31 fail_io (const char *msg, ...)
  32 {
  33   va_list args;
  34
  35   va_start (args, msg);
  36   vfprintf (stderr, msg, args);
  37   va_end (args);
  38
  39   if (errno != 0)
  40     fprintf (stderr, ": %s", strerror (errno));
  41   putc ('\n', stderr);
  42   exit (EXIT_FAILURE);
  43 }
  44
  45 /* If FD is a terminal, configures it for noncanonical input mode
  46    with VMIN and VTIME set as indicated.
  47    If FD is not a terminal, has no effect. */
  48 static void
  49 make_noncanon (int fd, int vmin, int vtime)
  50 {
  51   if (isatty (fd))
  52     {
  53       struct termios termios;
  54       if (tcgetattr (fd, &termios) < 0)
  55         fail_io ("tcgetattr");
  56       termios.c_lflag &= ~(ICANON | ECHO);
  57       termios.c_cc[VMIN] = vmin;
  58       termios.c_cc[VTIME] = vtime;
  59       if (tcsetattr (fd, TCSANOW, &termios) < 0)
  60         fail_io ("tcsetattr");
  61     }
  62 }
  63
  64 /* Make FD non-blocking if NONBLOCKING is true,
  65    or blocking if NONBLOCKING is false. */
  66 static void
  67 make_nonblocking (int fd, bool nonblocking)
  68 {
  69   int flags = fcntl (fd, F_GETFL);
  70   if (flags < 0)
  71     fail_io ("fcntl");
  72   if (nonblocking)
  73     flags |= O_NONBLOCK;
  74   else
  75     flags &= ~O_NONBLOCK;
  76   if (fcntl (fd, F_SETFL, flags) < 0)
  77     fail_io ("fcntl");
  78 }
  79
  80 /* Handle a read or write on *FD, which is the socket if
  81    FD_IS_SOCK is true, that returned end-of-file or error
  82    indication RETVAL.  The system call is named CALL, for use in
  83    error messages.  Returns true if processing may continue,
  84    false if we're all done. */
  85 static bool
  86 handle_error (ssize_t retval, int *fd, bool fd_is_sock, const char *call)
  87 {
  88   if (retval == 0)
  89     {
  90       if (fd_is_sock)
  91         return false;
  92       else
  93         {
  94           *fd = -1;
  95           return true;
  96         }
  97     }
  98   else
  99     fail_io (call);
 100 }
 101
 102 /* Copies data from stdin to SOCK and from SOCK to stdout until no
 103    more data can be read or written. */
 104 static void
 105 relay (int sock)
 106 {
 107   struct pipe
 108     {
 109       int in, out;
 110       char buf[BUFSIZ];
 111       size_t size, ofs;
 112       bool active;
 113     };
 114   struct pipe pipes[2];
 115
 116   /* In case stdin is a file, go back to the beginning.
 117      This allows replaying the input on reset. */
 118   lseek (STDIN_FILENO, 0, SEEK_SET);
 119
 120   /* Make SOCK, stdin, and stdout non-blocking. */
 121   make_nonblocking (sock, true);
 122   make_nonblocking (STDIN_FILENO, true);
 123   make_nonblocking (STDOUT_FILENO, true);
 124
 125   /* Configure noncanonical mode on stdin to avoid waiting for
 126      end-of-line. */
 127   make_noncanon (STDIN_FILENO, 1, 0);
 128
 129   memset (pipes, 0, sizeof pipes);
 130   pipes[0].in = STDIN_FILENO;
 131   pipes[0].out = sock;
 132   pipes[1].in = sock;
 133   pipes[1].out = STDOUT_FILENO;
 134
 135   while (pipes[0].in != -1 || pipes[1].in != -1
 136          || (pipes[1].size && pipes[1].out != -1))
 137     {
 138       fd_set read_fds, write_fds;
 139       sigset_t empty_set;
 140       int retval;
 141       int i;
 142
 143       FD_ZERO (&read_fds);
 144       FD_ZERO (&write_fds);
 145       for (i = 0; i < 2; i++)
 146         {
 147           struct pipe *p = &pipes[i];
 148
 149           /* Don't do anything with the stdin->sock pipe until we
 150              have some data for the sock->stdout pipe.  If we get
 151              too eager, vmplayer will throw away our input. */
 152           if (i == 0 && !pipes[1].active)
 153             continue;
 154
 155           if (p->in != -1 && p->size + p->ofs < sizeof p->buf)
 156             FD_SET (p->in, &read_fds);
 157           if (p->out != -1 && p->size > 0)
 158             FD_SET (p->out, &write_fds);
 159         }
 160       sigemptyset (&empty_set);
 161       retval = pselect (FD_SETSIZE, &read_fds, &write_fds, NULL, NULL,
 162                         &empty_set);
 163       if (retval < 0)
 164         {
 165           if (errno == EINTR)
 166             {
 167               /* Child died.  Do final relaying. */
 168               struct pipe *p = &pipes[1];
 169               if (p->out == -1)
 170                 exit (0);
 171               make_nonblocking (STDOUT_FILENO, false);
 172               for (;;)
 173                 {
 174                   ssize_t n;
 175
 176                   /* Write buffer. */
 177                   while (p->size > 0)
 178                     {
 179                       n = write (p->out, p->buf + p->ofs, p->size);
 180                       if (n < 0)
 181                         fail_io ("write");
 182                       else if (n == 0)
 183                         fail_io ("zero-length write");
 184                       p->ofs += n;
 185                       p->size -= n;
 186                     }
 187                   p->ofs = 0;
 188
 189                   p->size = n = read (p->in, p->buf, sizeof p->buf);
 190                   if (n <= 0)
 191                     exit (0);
 192                 }
 193             }
 194           fail_io ("select");
 195         }
 196
 197       for (i = 0; i < 2; i++)
 198         {
 199           struct pipe *p = &pipes[i];
 200           if (p->in != -1 && FD_ISSET (p->in, &read_fds))
 201             {
 202               ssize_t n = read (p->in, p->buf + p->ofs + p->size,
 203                                 sizeof p->buf - p->ofs - p->size);
 204               if (n > 0)
 205                 {
 206                   p->active = true;
 207                   p->size += n;
 208                   if (p->size == BUFSIZ && p->ofs != 0)
 209                     {
 210                       memmove (p->buf, p->buf + p->ofs, p->size);
 211                       p->ofs = 0;
 212                     }
 213                 }
 214               else if (!handle_error (n, &p->in, p->in == sock, "read"))
 215                 return;
 216             }
 217           if (p->out != -1 && FD_ISSET (p->out, &write_fds))
 218             {
 219               ssize_t n = write (p->out, p->buf + p->ofs, p->size);
 220               if (n > 0)
 221                 {
 222                   p->ofs += n;
 223                   p->size -= n;
 224                   if (p->size == 0)
 225                     p->ofs = 0;
 226                 }
 227               else if (!handle_error (n, &p->out, p->out == sock, "write"))
 228                 return;
 229             }
 230         }
 231     }
 232 }
 233
 234 static void
 235 sigchld_handler (int signo __attribute__ ((unused)))
 236 {
 237   /* Nothing to do. */
 238 }
 239
 240 int
 241 main (int argc __attribute__ ((unused)), char *argv[])
 242 {
 243   pid_t pid;
 244   struct itimerval zero_itimerval;
 245   struct sockaddr_un sun;
 246   sigset_t sigchld_set;
 247   int sock;
 248
 249   if (argc < 3)
 250     {
 251       fprintf (stderr,
 252                "usage: squish-unix SOCKET COMMAND [ARG]...\n"
 253                "Squishes both stdin and stdout into a single Unix domain\n"
 254                "socket named SOCKET, and runs COMMAND as a subprocess.\n");
 255       return EXIT_FAILURE;
 256     }
 257
 258   /* Create socket. */
 259   sock = socket (PF_LOCAL, SOCK_STREAM, 0);
 260   if (sock < 0)
 261     fail_io ("socket");
 262
 263   /* Configure socket. */
 264   sun.sun_family = AF_LOCAL;
 265   strncpy (sun.sun_path, argv[1], sizeof sun.sun_path);
 266   sun.sun_path[sizeof sun.sun_path - 1] = '\0';
 267   if (unlink (sun.sun_path) < 0 && errno != ENOENT)
 268     fail_io ("unlink");
 269   if (bind (sock, (struct sockaddr *) &sun,
 270             (offsetof (struct sockaddr_un, sun_path)
 271              + strlen (sun.sun_path) + 1)) < 0)
 272     fail_io ("bind");
 273
 274   /* Listen on socket. */
 275   if (listen (sock, 1) < 0)
 276     fail_io ("listen");
 277
 278   /* Block SIGCHLD and set up a handler for it. */
 279   sigemptyset (&sigchld_set);
 280   sigaddset (&sigchld_set, SIGCHLD);
 281   if (sigprocmask (SIG_BLOCK, &sigchld_set, NULL) < 0)
 282     fail_io ("sigprocmask");
 283   if (signal (SIGCHLD, sigchld_handler) == SIG_ERR)
 284     fail_io ("signal");
 285
 286   /* Save the virtual interval timer, which might have been set
 287      by the process that ran us.  It really should be applied to
 288      our child process. */
 289   memset (&zero_itimerval, 0, sizeof zero_itimerval);
 290   if (setitimer (ITIMER_VIRTUAL, &zero_itimerval, NULL) < 0)
 291     fail_io ("setitimer");
 292
 293   pid = fork ();
 294   if (pid < 0)
 295     fail_io ("fork");
 296   else if (pid != 0)
 297     {
 298       /* Running in parent process. */
 299       make_nonblocking (sock, true);
 300       for (;;)
 301         {
 302           fd_set read_fds;
 303           sigset_t empty_set;
 304           int retval;
 305           int conn;
 306
 307           /* Wait for connection. */
 308           FD_ZERO (&read_fds);
 309           FD_SET (sock, &read_fds);
 310           sigemptyset (&empty_set);
 311           retval = pselect (sock + 1, &read_fds, NULL, NULL, NULL, &empty_set);
 312           if (retval < 0)
 313             {
 314               if (errno == EINTR)
 315                 break;
 316               fail_io ("select");
 317             }
 318
 319           /* Accept connection. */
 320           conn = accept (sock, NULL, NULL);
 321           if (conn < 0)
 322             fail_io ("accept");
 323
 324           /* Relay connection. */
 325           relay (conn);
 326           close (conn);
 327         }
 328       return 0;
 329     }
 330   else
 331     {
 332       /* Running in child process. */
 333       if (close (sock) < 0)
 334         fail_io ("close");
 335       execvp (argv[2], argv + 2);
 336       fail_io ("exec");
 337     }
 338 }