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