spawn-process-helper.c

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