/* ISC license. */

#include <skalibs/sysdeps.h>

#if defined(SKALIBS_HASPRSETCHILDSUBREAPER) || defined(SKALIBS_HASKEVENT)

#include <unistd.h>
#include <signal.h>
#include <sys/wait.h>
#include <errno.h>
#ifdef SKALIBS_HASPRSETCHILDSUBREAPER
# include <sys/prctl.h>
#else
# include <sys/types.h>
# include <sys/event.h>
# include <sys/time.h>
# include <pthread.h>
# include <fcntl.h>
# include <limits.h>
#endif

#include <skalibs/posixplz.h>
#include <skalibs/types.h>
#include <skalibs/prog.h>
#include <skalibs/gol.h>
#include <skalibs/bytestr.h>
#include <skalibs/strerr.h>
#include <skalibs/tai.h>
#include <skalibs/selfpipe.h>
#include <skalibs/iopause.h>
#include <skalibs/cspawn.h>
#include <skalibs/djbunix.h>


 /* portable parts */

#define USAGE "s6-background-watch [ -t timeout ] [ -d notif ] pidfile prog..."
#define dieusage() strerr_dieusage(100, USAGE)

static int handle_signals_early (pid_t pid, char const *pidfile, char const *prog)
{
  int sig ;
  for (;;) switch (sig = selfpipe_read())
  {
    case -1 : strerr_diefu1sys(111, "read selfpipe") ;
    case 0 : return 0 ;
    case SIGCHLD :
    {
      int wstat ;
      pid_t r = wait_pid_nohang(pid, &wstat) ;
      if (r == -1) strerr_diefu1sys(111, "wait") ;
      if (!r) break ;
      if (WIFEXITED(wstat))
      {
        if (WEXITSTATUS(wstat))
        {
          char fmt[UINT_FMT] ;
          fmt[uint_fmt(fmt, WEXITSTATUS(wstat))] = 0 ;
          unlink_void(pidfile) ;
          strerr_dief3x(wait_estatus(wstat), prog, " (parent) exited ", fmt) ;
        }
        else return 1 ;
      }
      else
      {
        char fmt[UINT_FMT] ;
        fmt[uint_fmt(fmt, WTERMSIG(wstat))] = 0 ;
        unlink_void(pidfile) ;
        strerr_dief3x(wait_estatus(wstat), prog, " (parent) crashed with signal ", fmt) ;
      }
      break ;
    }
    default :
      kill(pid, sig) ;
      break ;
  }
}

static inline int read_file (char const *file, char *buf, size_t n)
{
  ssize_t r = openreadnclose_nb(file, buf, n) ;
  if (r == -1)
  {
    if (errno != ENOENT) strerr_diefu2sys(111, "open ", file) ;
    return 0 ;
  }
  buf[byte_chr(buf, r, '\n')] = 0 ;
  return 1 ;
}

static inline pid_t get_pid_from_pidfile (char const *pidfile, char const *prog)
{ /* for now, assume the pidfile is there once the parent has exited */
  pid_t pid ;
  char buf[128] ;
  if (!read_file(pidfile, buf, 128))
    strerr_diefu4sys(104, prog, " (parent) exited but pidfile ", pidfile, " does not exist") ;
  if (!pid0_scan(buf, &pid))
    strerr_dief3x(104, "pidfile ", pidfile, " does not contain a valid pid") ;
  return pid ;
}

static inline int handle_signals (pid_t pid, int issr, int *code)
{
  int sig ;
  for (;;) switch (sig = selfpipe_read())
  {
    case -1 : strerr_diefu1sys(111, "read selfpipe") ;
    case 0 : return 0 ;
    case SIGCHLD :
    {
      int wstat ;
      pid_t r = wait_pid_nohang(pid, &wstat) ;
      if (r == -1)
        if (errno != ECHILD || issr) strerr_diefu1sys(111, "wait") ;
        else break ;
      else if (!r) break ;
      *code = wait_estatus(wstat) ;
      return 1 ;
    }
    default :
      kill(pid, sig) ;
      break ;
  }
}


/* BSD part. kevent is a unique dildo-shaped snowflake, so we need to run it
   in its own thread for it to play nicely with our regular iopause loop. */

#ifdef SKALIBS_HASKEVENT

struct thinfo_s
{
  int kq ;
  int fdw ;
  pid_t pid ;
} ;

static void *keventwait (void *arg)
{
  struct thinfo_s *t = arg ;
  struct kevent ke ;
  for (;;)
  {
    int r = kevent(t->kq, 0, 0, &ke, 1, 0) ;
    if (r == -1 && errno != EINTR) strerr_diefu2sys(111, "read ", "kevent") ;
    else if (r && (pid_t)ke.ident == t->pid && ke.filter == EVFILT_PROC && ke.fflags & NOTE_EXIT) break ;
  }
  close(t->fdw) ;
  return (void *)(intptr_t)ke.data ;
}

static void keventstart (pid_t pid, int *fd, pthread_t *th)
{
  struct thinfo_s t ;
  int p[2] ;
  pthread_attr_t attr ;
  struct kevent ke ;
  int e ;

  t.kq = kqueue() ;
  if (t.kq == -1) strerr_diefu1sys(111, "kqueue") ;
  EV_SET(&ke, pid, EVFILT_PROC, EV_ADD | EV_ONESHOT, NOTE_EXIT, 0, 0) ;
  if (kevent(t.kq, &ke, 1, 0, 0, 0) == -1) strerr_diefu1sys(111, "register kevent") ;
  if (pipe(p) == -1) strerr_diefu2sys(111, "pipe") ;
  t.fdw = p[1] ;
  t.pid = pid ;
  e = pthread_attr_init(&attr) ;
  if (e) { errno = e ; strerr_diefu1sys(111, "pthread_attr_init") ; }
  pthread_attr_setstacksize(&attr, PTHREAD_STACK_MIN) ;
  e = pthread_create(th, &attr, &keventwait, &t) ;
  if (e) { errno = e ; strerr_diefu1sys(111, "pthread_create") ; }
  *fd = p[0] ;
}

static int getexitcode (pthread_t th)
{
  void *ptr ;
  int e = pthread_join(th, &ptr) ;
  if (e) { errno = e ; strerr_diefu1sys(111, "pthread_join") ; }
  return (int)(intptr_t)ptr ;
}

#endif


 /* main */

enum rgola_e
{
  GOLA_TIMEOUT,
  GOLA_NOTIF,
  GOLA_N
} ;

static gol_arg const rgola[GOLA_N] =
{
  { .so = 't', .lo = "timeout-ready", .i = GOLA_TIMEOUT },
  { .so = 'd', .lo = "notification-fd", .i = GOLA_NOTIF }
} ;

int main (int argc, char const *const *argv)
{
  iopause_fd x[2] = { { .events = IOPAUSE_READ }, { .fd = -2, .events = IOPAUSE_READ, .revents = 0 } } ;
  tain tto = TAIN_INFINITE_RELATIVE ;
  tain deadline ;
  unsigned int notif = 0 ;
  pid_t pid ;
  
  char const *wgola[GOLA_N] = { 0 } ;
  unsigned int golc ;
  PROG = "s6-background-watch" ;
  golc = gol_main(argc, argv, 0, 0, rgola, GOLA_N, 0, wgola) ;
  argc -= golc ; argv += golc ;
  if (argc < 2) dieusage() ;
  if (wgola[GOLA_TIMEOUT])
  {
    unsigned int t ;
    if (!uint0_scan(wgola[GOLA_TIMEOUT], &t)) dieusage() ;
    if (!tain_from_millisecs(&tto, t)) dieusage() ;
  }
  if (wgola[GOLA_NOTIF])
  {
    if (!uint0_scan(wgola[GOLA_NOTIF], &notif)) dieusage() ;
    if (notif < 3) strerr_dief2x(100, "notification-fd", " cannot be 0, 1 or 2") ;
    if (fcntl(notif, F_GETFD) == -1) strerr_diefu2sys(111, "check ", "notification-fd") ;
    if (coe(notif) == -1) strerr_diefu1sys(111, "coe notification-fd") ;
  }

  {
    sigset_t full ;
    sigfillset(&full) ;
    if (!selfpipe_trapset(&full))
      strerr_diefu1sys(111, "trap all signals") ;
    x[0].fd = selfpipe_fd() ;
  }

#ifdef SKALIBS_HASPRSETCHILDSUBREAPER
  if (prctl(PR_SET_CHILD_SUBREAPER, 1) == -1)
    strerr_diefu1sys(111, "prctl to become a subreaper") ;
#endif

  pid = cspawn(argv[1], argv + 1, (char const *const *)environ, CSPAWN_FLAGS_SELFPIPE_FINISH, 0, 0) ;
  if (!pid) strerr_diefu2sys(errno == ENOEXEC ? 126 : 127, "spawn ", argv[1]) ;

  tain_now_set_stopwatch_g() ;
  tain_add_g(&deadline, &tto) ;
  for (;;)
  {
    int r = iopause_g(x, 1, &deadline) ;
    if (r == -1) strerr_diefu1sys(111, "iopause") ;
    if (!r)
    {
      strerr_warnw2x(argv[1], " (parent) did not exit before timeout-ready, killing it") ;
      kill(pid, SIGKILL) ;
    }
    else if (handle_signals_early(pid, argv[0], argv[1])) break ;
  }

  pid = get_pid_from_pidfile(argv[0], argv[1]) ;
  if (kill(pid, 0) == -1) strerr_diefu1sys(111, "check daemon with a null signal") ;

#ifdef SKALIBS_HASKEVENT
  pthread_t th = PTHREAD_NULL ;
  keventstart(pid, &x[1].fd, &th) ;
#endif
  if (notif)
  {
    write(notif, "\n", 1) ;
    close(notif) ;
  }

  for (;;)
  {
    int r = iopause_g(x, 1 + (x[1].fd >= 0), 0) ;
    if (r == -1) strerr_diefu1sys(111, "iopause") ;
    if (x[0].revents & IOPAUSE_READ)
    {
      int code = 0 ;
      if (handle_signals(pid, x[1].fd == -2, &code)) _exit(code) ;
    }
#ifdef SKALIBS_KEVENT
    else if (x[1].revents & IOPAUSE_READ) _exit(getexitcode(th)) ;
#endif
  }
}

#else

#include <skalibs/prog.h>
#include <skalibs/strerr.h>

int main (int argc, char const *const *argv)
{
  PROG = "s6-background-watch" ;
  strerr_dief1x(112, "system does not meet the requirements to watch auto-backgrounding daemons") ;
}

#endif