pintos: Avoid apparent name collision in definition of SIGALRM subroutine.
[pintos-anon] / 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 }