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