conkeror-spawn-helper.c

   1 /**
   2  * (C) Copyright 2008 Jeremy Maitin-Shepard
   3  *
   4  * Use, modification, and distribution are subject to the terms specified in the
   5  * COPYING file.
   6  **/
   7
   8 #include <sys/types.h>
   9 #include <sys/socket.h>
  10 #include <unistd.h>
  11 #include <stdio.h>
  12 #include <errno.h>
  13 #include <stdlib.h>
  14 #include <signal.h>
  15 #include <sys/wait.h>
  16 #include <sys/stat.h>
  17 #include <string.h>
  18 #include <fcntl.h>
  19 #include <netinet/in.h>
  20 #include <dirent.h>
  21 #include <sys/resource.h>
  22
  23 void fail(const char *msg) {
  24   fprintf(stderr, "%s\n", msg);
  25   exit(1);
  26 }
  27
  28 void failerr(const char *msg) {
  29   perror(msg);
  30   exit(1);
  31 }
  32
  33 #define TRY(var, foo) var = foo; while (var == -1) { if(errno != EINTR) failerr(#foo); }
  34
  35 void *Malloc(size_t count) { void *r = malloc(count); if (!r) fail("malloc"); return r; }
  36
  37 /**
  38  * read_all: read from the specified file descriptor, returning a
  39  * malloc-allocated buffer containing the data that was read; the
  40  * number of bytes read is stored in *bytes_read.  If max_bytes is
  41  * non-negative, it specifies the maximum number of bytes to read.
  42  * Otherwise, read_all reads from the file descriptor until the end of
  43  * file is reached.
  44  */
  45 char *read_all(int fd, int max_bytes, int *bytes_read) {
  46   int capacity = 256;
  47   if (max_bytes > 0)
  48     capacity = max_bytes;
  49   char *buffer = Malloc(capacity);
  50   int count = 0;
  51   if (max_bytes < 0 || max_bytes > 0) {
  52     while (1) {
  53       int remain;
  54       if (count == capacity) {
  55         capacity *= 2;
  56         buffer = realloc(buffer, capacity);
  57         if (!buffer)
  58           fail("realloc failed");
  59       }
  60       remain = capacity - count;
  61       if (max_bytes > 0 && remain > max_bytes)
  62         remain = max_bytes;
  63       TRY(remain, read(fd, buffer + count, remain));
  64       count += remain;
  65       if (remain == 0 || count == max_bytes)
  66         break;
  67     }
  68   }
  69   *bytes_read = count;
  70   return buffer;
  71 }
  72
  73 /**
  74  * next_term: return the next NUL terminated string from buffer, and
  75  * adjust buffer and len accordingly.
  76  */
  77 char *next_term(char **buffer, int *len) {
  78   char *p = *buffer;
  79   int x = 0;
  80   int max_len = *len;
  81   while (x < max_len && p[x])
  82     ++x;
  83   if (x == max_len)
  84     fail("error parsing");
  85   *buffer += x + 1;
  86   *len -= (x + 1);
  87   return p;
  88 }
  89
  90 struct fd_info {
  91   int desired_fd;
  92   int orig_fd;
  93   char *path;
  94   int open_mode;
  95   int perms;
  96 };
  97
  98 void write_all(int fd, const char *buf, int len) {
  99   int result;
 100   do {
 101     TRY(result, write(fd, buf, len));
 102     buf += result;
 103     len -= result;
 104   } while (len > 0);
 105 }
 106
 107 /**
 108  * my_connect: Create a connection to the local Conkeror process on
 109  * the specified TCP port.  After connecting, the properly formatted
 110  * header specifying the client_key and the "role" (file descriptor or
 111  * -1 to indicate the control socket) are sent as well.  The file
 112  * descriptor for the socket is returned.
 113  */
 114 int my_connect(int port, char *client_key, int role) {
 115   int sockfd;
 116   int result;
 117   struct sockaddr_in sa;
 118
 119   TRY(sockfd, socket(PF_INET, SOCK_STREAM, 0));
 120   sa.sin_family = AF_INET;
 121   sa.sin_port = htons(port);
 122   sa.sin_addr.s_addr = inet_addr("127.0.0.1");
 123   memset(sa.sin_zero, 0, sizeof(sa.sin_zero));
 124
 125   TRY(result, connect(sockfd, (struct sockaddr *)&sa, sizeof(sa)));
 126
 127   /* Send the client key */
 128   write_all(sockfd, client_key, strlen(client_key));
 129
 130   /* Send the role */
 131   if (role < 0) {
 132     write_all(sockfd, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0", 15);
 133   }
 134   else {
 135     char buf[16];
 136     snprintf(buf, 16, "%15d", role);
 137     write_all(sockfd, buf, 15);
 138   }
 139
 140   return sockfd;
 141 }
 142
 143 int child_pid = 0;
 144 int control_fd;
 145
 146 /**
 147  * sigchld_handler: reap any waitable children.  Once the child
 148  * process exits, send the exit status back over the control socket,
 149  * then exit. */
 150 void sigchld_handler(int sig) {
 151   int status;
 152   int pid;
 153   int err;
 154
 155   while (1) {
 156     pid = waitpid(-1, &status, WNOHANG);
 157     if (pid == 0)
 158       return;
 159     if (pid == -1) {
 160       if (errno == ECHILD)
 161         break;
 162       failerr("waitpid");
 163     }
 164
 165     /* Our child process exited */
 166     if (pid == child_pid && (WIFEXITED(status) || WIFSIGNALED(status))) {
 167       char buf[30];
 168       snprintf(buf, 30, "%d", status);
 169       write_all(control_fd, buf, strlen(buf) + 1);
 170       exit(0);
 171     }
 172   }
 173 }
 174
 175 void check_duplicate_fds(struct fd_info *fds, int fd_count) {
 176   int i, j;
 177   for (i = 0; i < fd_count; ++i) {
 178     for (j = i + 1; j < fd_count; ++j) {
 179       if (fds[i].desired_fd == fds[j].desired_fd)
 180         fail("duplicate redirection requested");
 181     }
 182   }
 183 }
 184
 185 /**
 186  * setup_fds: Make the requested redirections.  For each entry in the
 187  * fds array, rename orig_fd to desired_fd.
 188  */
 189 void setup_fds(struct fd_info *fds, int fd_count) {
 190   int i, j, result;
 191   for (i = 0; i < fd_count; ++i) {
 192     int fd = fds[i].desired_fd;
 193     /* Check if this file descriptor is still in use by any subsequent
 194        redirection. */
 195     for (j = i + 1; j < fd_count; ++j) {
 196       if (fd == fds[j].orig_fd) {
 197         /* It is in use.  Pick a new file descriptor for fds[j]. */
 198         int fd_new;
 199         TRY(fd_new, dup(fds[j].orig_fd));
 200         close(fds[j].orig_fd);
 201         fds[j].orig_fd = fd_new;
 202         break;
 203       }
 204     }
 205     TRY(result, dup2(fds[i].orig_fd, fd));
 206     close(fds[i].orig_fd);
 207   }
 208 }
 209
 210 int main(int argc, char **argv) {
 211
 212   int port;
 213   char *client_key, *server_key, *executable, *workdir;
 214   char **my_argv;
 215   struct fd_info *fds;
 216   int fd_count;
 217   int i;
 218   sigset_t my_mask, my_old_mask;
 219
 220   if (argc != 3 || (port = atoi(argv[2])) == 0)
 221     fail("Invalid arguments");
 222
 223   sigemptyset(&my_mask);
 224   sigaddset(&my_mask, SIGCHLD);
 225
 226   /* Block SIGPIPE to avoid a signal being generated while writing to a socket */
 227   signal(SIGPIPE, SIG_IGN);
 228
 229   /* Close everything except STDERR.  Mozilla leaves us with a bunch
 230      of junk file descriptors. */
 231   {
 232     DIR *dir = opendir("/proc/self/fd");
 233     if (!dir) {
 234       /* No proc filesystem available, just loop through file descriptors */
 235       struct rlimit file_lim;
 236       int max_fileno = 1024;
 237       if (getrlimit(RLIMIT_NOFILE, &file_lim) == 0)
 238         max_fileno = file_lim.rlim_cur;
 239       for (i = 0; i < max_fileno; ++i) {
 240         if (i == STDERR_FILENO)
 241           continue;
 242         close(i);
 243       }
 244     } else {
 245       struct dirent *dir_ent;
 246       int dir_fd = dirfd(dir);
 247       while ((dir_ent = readdir(dir)) != NULL) {
 248         int file_desc = atoi(dir_ent->d_name);
 249         if (file_desc == STDERR_FILENO || file_desc == dir_fd)
 250           continue;
 251         close(file_desc);
 252       }
 253       closedir(dir);
 254     }
 255   }
 256
 257   /* Parse key file */
 258   {
 259     char *buf;
 260     int len;
 261     int my_argc;
 262     /* Read the entire file into buf. */
 263     {
 264       int file;
 265       TRY(file, open(argv[1], O_RDONLY));
 266       buf = read_all(file, -1, &len);
 267       close(file);
 268
 269       /* Remove the temporary file */
 270       remove(argv[1]);
 271     }
 272     client_key = next_term(&buf, &len);
 273     server_key = next_term(&buf, &len);
 274     executable = next_term(&buf, &len);
 275     workdir = next_term(&buf, &len);
 276     my_argc = atoi(next_term(&buf, &len));
 277     my_argv = Malloc(sizeof(char *) * (my_argc + 1));
 278     for (i = 0; i < my_argc; ++i)
 279       my_argv[i] = next_term(&buf, &len);
 280     my_argv[my_argc] = NULL;
 281     fd_count = atoi(next_term(&buf, &len));
 282     if (fd_count < 0) fail("invalid fd count");
 283     fds = Malloc(sizeof(struct fd_info) * fd_count);
 284     for (i = 0; i < fd_count; ++i) {
 285       fds[i].desired_fd = atoi(next_term(&buf, &len));
 286       fds[i].path = next_term(&buf, &len);
 287       if (fds[i].path[0]) {
 288         fds[i].open_mode = atoi(next_term(&buf, &len));
 289         fds[i].perms = atoi(next_term(&buf, &len));
 290       }
 291     }
 292     if (len != 0)
 293       fail("invalid input file");
 294   }
 295
 296   /* Validate the file descriptor redirection request. */
 297   check_duplicate_fds(fds, fd_count);
 298
 299   /* Create the control socket connection. */
 300   control_fd = my_connect(port, client_key, -1);
 301
 302   /* Create a socket connection or open a local file for each
 303      requested file descriptor redirection. */
 304   for (i = 0; i < fd_count; ++i) {
 305     if (fds[i].path[0]) {
 306       TRY(fds[i].orig_fd, open(fds[i].path, fds[i].open_mode, fds[i].perms));
 307     } else {
 308       fds[i].orig_fd = my_connect(port, client_key, fds[i].desired_fd);
 309     }
 310   }
 311
 312   /* Check server key */
 313   {
 314     int len = strlen(server_key);
 315     int read_len;
 316     char *buf = read_all(control_fd, len, &read_len);
 317     if (len != read_len || memcmp(buf, server_key, len) != 0)
 318       fail("server key mismatch");
 319     free(buf);
 320   }
 321
 322   /* Block SIGCHLD */
 323   sigprocmask(SIG_BLOCK, &my_mask, &my_old_mask);
 324
 325   /* Create the child process */
 326   child_pid = fork();
 327   if (child_pid == 0) {
 328     int result;
 329     /* Unblock SIGCHLD */
 330     sigprocmask(SIG_SETMASK, &my_old_mask, NULL);
 331
 332     /* Reset the SIGPIPE signal handler. */
 333     signal(SIGPIPE, SIG_DFL);
 334
 335     /* Close the control socket, as it isn't needed from the child. */
 336     close(control_fd);
 337
 338     /* Change to the specified working directory. */
 339     if (workdir[0] != 0) {
 340       if (chdir(workdir) == -1)
 341         failerr(workdir);
 342     }
 343
 344     /* Rearrange file descriptors according to the user specification */
 345     setup_fds(fds, fd_count);
 346
 347     /* Exec */
 348     TRY(result, execv(executable, my_argv));
 349
 350   } else if (child_pid == -1) {
 351     failerr("fork");
 352   } else {
 353     /* We are in the parent process */
 354     char msg;
 355     int count;
 356
 357     /* Install SIGCHLD handler */
 358     {
 359       struct sigaction act;
 360       act.sa_handler = sigchld_handler;
 361       sigemptyset(&act.sa_mask);
 362       act.sa_flags = SA_NOCLDSTOP;
 363       sigaction(SIGCHLD, &act, NULL);
 364     }
 365     /* Unblock SIGCHLD */
 366     sigprocmask(SIG_SETMASK, &my_old_mask, NULL);
 367
 368     /* Close all of the redirection file descriptors, as we don't need
 369        them from the parent. */
 370     for (i = 0; i < fd_count; ++i)
 371       close(fds[i].orig_fd);
 372
 373     /* Wait for a message from the server telling us to exit early. */
 374     TRY(count, read(control_fd, &msg, 1));
 375
 376     if (count == 0) {
 377       /* End of file received: exit without killing child */
 378       return;
 379     }
 380
 381     /* Assume msg == 0 until we support more messages */
 382     TRY(count, kill(child_pid, SIGTERM));
 383     return;
 384   }
 385 }