shwait.c

/* gcc -O3 -ansi -Wall -Wextra -pedantic -static shwait.c -lrt -lutil -o shwait */
#define _XOPEN_SOURCE 700
#include <errno.h>
#include <fcntl.h>
#include <pty.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <termios.h>
#include <time.h>
#include <unistd.h>
#include <utmp.h>

/* Structure of shared object */
struct msg {
  struct timespec ts_shell; /* execing shell */
  struct timespec ts_prompt; /* something (prompt?) received over pty */
  struct timespec ts_grandchild; /* shell spawned a job and it ran */
};

static struct timespec TS_INF = { -1, -1 };

/* Shared object name */
#define SHM_PATH "/shwait.shm."
#define SHM_PATH_SLEN 64
#define SHM_ENV "SHWAIT_SHM"

/* Prepare randomized name for shared object */
static char *randpath(void *seed, size_t len) {
  static char hex[16] = "0123456789abcdef";
  static char path[SHM_PATH_SLEN] = SHM_PATH;
  unsigned k, i = sizeof(SHM_PATH)-1, h = getpid();
#define MIX(h) { (h) ^= (h) << 13; (h) ^= (h) >> 17; (h) ^= (h) << 5; }
  for (k = 0; h >> k; k+=4)
    path[i++] = hex[(h>>k) & 0xf];
  path[i++] = '.';
  h ^= 2463534242;
  for (k = 0; k < len; k++) {
    MIX(h); h ^= ((unsigned char*)seed)[k];
  }
  for (; i < sizeof(path)-1; i++) {
    MIX(h); path[i] = hex[h & 0xf];
  }
#undef MIX
  path[i] = '\0';
  return path;
}

/* Prepare shared memory object and expose its name in environment */
static struct msg *msg_create(char *path) {
  struct msg *shm;
  int md;
  if (setenv(SHM_ENV, path, 1)) {
    return NULL;
  }
  md = shm_open(path, O_RDWR|O_CREAT|O_EXCL, S_IRUSR|S_IWUSR);
  if (md < 0)
    return NULL;
  (void) ftruncate(md, sizeof *shm);
  shm = (struct msg*) mmap(NULL, sizeof *shm,
      PROT_READ|PROT_WRITE, MAP_SHARED, md, 0);
  (void) close(md);
  if (shm == MAP_FAILED) {
    (void) shm_unlink(SHM_PATH);
    return NULL;
  }
  return shm;
}

/* Remove shared memory object */
static int msg_remove(struct msg *shm) {
  (void) munmap(shm, sizeof *shm);
  return shm_unlink(getenv(SHM_ENV));
}

/* Push ts_grandchild to shared object */
static int msg_send(struct timespec ts_main) {
  struct msg *shm;
  char *path = getenv(SHM_ENV);
  int md;
  if (!path)
    return -1;
  if ((md = shm_open(path, O_RDWR, 0)) < 0)
    return errno;
  shm = (struct msg*) mmap(NULL, sizeof *shm,
      PROT_READ|PROT_WRITE, MAP_SHARED, md, 0);
  (void) close(md);
  if (shm == MAP_FAILED)
    return errno;
  shm->ts_grandchild = ts_main;
  (void) munmap(shm, sizeof *shm);
  return 0;
}

/* Inhibit echo (so the first read will hopefully get the prompt) */
static int noecho(int tfd) {
  struct termios tp;
  if (tcgetattr(tfd, &tp))
    return -1;
  tp.c_lflag &= ~ECHO;
  return tcsetattr(tfd, TCSAFLUSH, &tp);
}

/* Sink data from input to the terminal's buffer */
static int sinkterm(int tfd) {
  int l, flags = fcntl(tfd, F_GETFL, 0);
  char buf[64];
  if (fcntl(tfd, F_SETFL, flags | O_NONBLOCK) == -1)
    return errno;
  errno = 0;
  for (;;) {
    l = read(0, buf, sizeof(buf));
    if (l < 1) {
      if (l < 0)
        perror("read");
      break;
    }
    l = write(tfd, buf, l);
    if (l < 1) {
      if (errno == EWOULDBLOCK)
        fprintf(stderr, "Script too large to fit in kernel's buffer\n");
      else
        perror("write");
      break;
    }
  }
  (void) fcntl(tfd, F_SETFL, flags &~O_NONBLOCK);
  return errno;
}

/* Subtract time specs */
static struct timespec ts_sub(struct timespec a, struct timespec b) {
  if (a.tv_nsec == -1)
    return a; /* infty */
  if (b.tv_nsec > a.tv_nsec) {
    a.tv_nsec += 1000000000L;
    a.tv_sec--;
  }
  a.tv_nsec -= b.tv_nsec;
  a.tv_sec  -= b.tv_sec;
  return a;
}

/* Print results */
static void ts_print(struct timespec ts) {
  if (ts.tv_nsec == -1)
    printf("inf\n");
  else
    printf("%lu.%09ld\n", ts.tv_sec, ts.tv_nsec);
}


int main(int argc, char **argv) {
  struct timespec ts_main;
  (void) clock_gettime(CLOCK_REALTIME, &ts_main);
  if (argc < 2) {
    if (msg_send(ts_main) == -1)
      fprintf(stderr, "Usage: %s shell < script\n", *argv);
    exit(errno);
  }
  else {
    struct msg *shm = msg_create(randpath(&ts_main, sizeof(ts_main)));
    int ptm, pts;
    if (openpty(&ptm, &pts, NULL, NULL, NULL))
      perror("openpty");
    else {
      (void) noecho(ptm);
      if (!sinkterm(ptm)) {
        pid_t child;
        char buf;
        int len;
        shm->ts_shell = TS_INF;
        shm->ts_grandchild = TS_INF;
        if (!(child = fork())) {
          close(ptm);
          login_tty(pts);
          argv[argc] = NULL; /* sic */
          (void) clock_gettime(CLOCK_REALTIME, &shm->ts_shell);
          execv(argv[1], argv+1);
          perror("execv");
          exit(errno);
        }
        len = read(ptm, &buf, sizeof(buf));
        (void) clock_gettime(CLOCK_REALTIME, &shm->ts_prompt);
        if (len < 0)
          perror("read");
        (void) waitpid(child, &len, 0);
        if (len)
          fprintf(stderr, "Exit status: %d\n", WEXITSTATUS(len));

        ts_print(ts_sub(shm->ts_prompt, shm->ts_shell));
        ts_print(ts_sub(shm->ts_grandchild, shm->ts_shell));
      }
      close(pts);
      close(ptm);
    }
    (void) msg_remove(shm);
  }
  exit(errno);
}