/*
 * Copyright 2006-2007 Free Software Foundation, Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 * Reap any leftover children possibly holding file descriptors.
 * Children are identified by the stale file descriptor or PGID / SID.
 * Both can be missed but only the stale file descriptors are important for us.
 * PGID / SID may be set by the children on their own.
 * If we fine a candidate we kill it will all its process tree (grandchildren).
 * The child process is run with `2>&1' redirection (due to forkpty(3)).
 * 2007-07-10  Jan Kratochvil  <jan.kratochvil@redhat.com>
 */

#define _GNU_SOURCE 1

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <dirent.h>
#include <unistd.h>
#include <errno.h>
#include <ctype.h>
#include <string.h>
#include <limits.h>
#include <fcntl.h>
#include <assert.h>
#include <pty.h>
#include <poll.h>

static const char *progname;

static volatile int signal_child_hit = 0;

static void signal_child (int signo)
{
  signal_child_hit = 1;
}

static char childptyname[LINE_MAX];
static pid_t child;

static int spawn (char **argv)
{
  pid_t child_got;
  int status, amaster, i;
  struct sigaction act;
  sigset_t sigset;

  /* Never miss SIGCHLD.  */
  i = sigemptyset (&sigset);
  assert (i == 0);
  i = sigaddset (&sigset, SIGCHLD);
  assert (i == 0);
  i = sigprocmask (SIG_BLOCK, &sigset, NULL);
  assert (i == 0);
  i = sigemptyset (&sigset);
  assert (i == 0);

  /* We do not use signal(2) to be sure we have SA_RESTART unset.  */
  memset (&act, 0, sizeof (act));
  act.sa_handler = signal_child;
  i = sigemptyset (&act.sa_mask);
  assert (i == 0);
  act.sa_flags = 0;
  i = sigaction (SIGCHLD, &act, NULL);
  assert (i == 0);

  child = forkpty (&amaster, childptyname, NULL, NULL);
  switch (child)
    {
      case -1:
	perror ("forkpty(3)");
	exit (EXIT_FAILURE);
      case 0:
	/* Do not setpgrp(2) in the parent process as the process-group
	   is shared for the whole sh(1) pipeline we could be a part
	   of.  The process-group is set according to PID of the first
	   command in the pipeline.
	   We would rip even vi(1) in the case of:
		./orphanripper sh -c 'sleep 1&' | vi -
	   */
	/* Do not setpgrp(2) as our pty would not be ours and we would
	   get `SIGSTOP' later, particularly after spawning gdb(1).
	   setsid(3) was already executed by forkpty(3) and it would fail if
	   executed again.  */
	if (getpid() != getpgrp ())
	  {
	    perror ("getpgrp(2)");
	    exit (EXIT_FAILURE);
	  }
	execvp (argv[1], argv + 1);
	perror ("execvp(2)");
	exit (EXIT_FAILURE);
      default:
	break;
    }
  i = fcntl (amaster, F_SETFL, O_RDWR | O_NONBLOCK);
  if (i != 0)
    {
      perror ("fcntl (amaster, F_SETFL, O_NONBLOCK)");
      exit (EXIT_FAILURE);
    }
  for (;;)
    {
      struct pollfd pollfd;
      char buf[LINE_MAX];
      ssize_t buf_got;

      pollfd.fd = amaster;
      pollfd.events = POLLIN;
      i = ppoll (&pollfd, 1, NULL, &sigset);
      assert (i == -1 || i == 1);
      if (i == -1)
        {
	  assert (errno == EINTR && signal_child_hit != 0);
	  break;
	}
      /* i == 1 */
      assert (pollfd.revents != 0);
      if (pollfd.revents & POLLHUP)
        break;
      if (pollfd.revents != POLLIN)
        {
	  fprintf (stderr, "%s: ppoll(2): revents 0x%x\n", progname,
		   (unsigned) pollfd.revents);
	  exit (EXIT_FAILURE);
	}
      buf_got = read (amaster, buf, sizeof buf);
      if (buf_got == 0)
	break;
      if (buf_got == -1)
	{
	  perror ("read (amaster)");
	  exit (EXIT_FAILURE);
	}
      if (write (STDOUT_FILENO, buf, buf_got) != buf_got)
	{
	  perror ("write(2)");
	  exit (EXIT_FAILURE);
	}
    }
  /* WNOHANG still could fail.  */
  child_got = waitpid (child, &status, 0);
  if (child != child_got)
    {
      fprintf (stderr, "waitpid (%d) = %d: %m\n", (int) child, (int) child_got);
      exit (EXIT_FAILURE);
    }
  if (!WIFEXITED (status))
    {
      fprintf (stderr, "waitpid (%d): !WIFEXITED (%d)\n", (int) child, status);
      exit (EXIT_FAILURE);
    }
  /* Do not close the master FD as the child would have `/dev/pts/23 (deleted)'
     entries which are not expected (and expecting ` (deleted)' would be
     a race.  */
#if 0
  i = close (amaster);
  if (i != 0)
    {
      perror ("close (forkpty ()'s amaster)");
      exit (EXIT_FAILURE);
    }
#endif

  return WEXITSTATUS (status);
}

/* Detected commandline may look weird due to a race:
   Original command:
	./orphanripper sh -c 'sleep 1&' &
   Correct output:
	[1] 29610
	./orphanripper: Killed -9 orphan PID 29612 (PGID 29611): sleep 1
   Raced output (sh(1) child still did not update its argv[]):
	[1] 29613
	./orphanripper: Killed -9 orphan PID 29615 (PGID 29614): sh -c sleep 1&
   We could delay a bit before ripping the children.  */
static const char *read_cmdline (pid_t pid)
{
  char cmdline_fname[32];
  static char cmdline[LINE_MAX];
  int fd;
  ssize_t got;
  char *s;

  if (snprintf (cmdline_fname, sizeof cmdline_fname, "/proc/%d/cmdline",
      (int) pid) < 0)
    return NULL;
  fd = open (cmdline_fname, O_RDONLY);
  if (fd == -1)
    {
      /* It may have already exited - ENOENT.  */
#if 0
      fprintf (stderr, "%s: open (\"%s\"): %m\n", progname, cmdline_fname);
#endif
      return NULL;
    }
  got = read (fd, cmdline, sizeof (cmdline) - 1);
  if (got == -1)
    fprintf (stderr, "%s: read (\"%s\"): %m\n", progname,
       cmdline_fname);
  if (close (fd) != 0)
    fprintf (stderr, "%s: close (\"%s\"): %m\n", progname,
       cmdline_fname);
  if (got < 0)
    return NULL;
  /* Convert '\0' argument delimiters to spaces.  */
  for (s = cmdline; s < cmdline + got; s++)
    if (!*s)
      *s = ' ';
  /* Trim the trailing spaces (typically single '\0'->' ').  */
  while (s > cmdline && isspace (s[-1]))
    s--;
  *s = 0;
  return cmdline;
}

static int fd_fs_scan (pid_t pid, int (*func) (pid_t pid, const char *link))
{
  DIR *dir;
  struct dirent *dirent;
  char dirname[64];
  int rc = 0;

  if (snprintf (dirname, sizeof dirname, "/proc/%d/fd", (int) pid) < 0)
    {
      perror ("snprintf(3)");
      exit (EXIT_FAILURE);
    }
  dir = opendir (dirname);
  if (dir == NULL)
    {
      if (errno == EACCES || errno == ENOENT)
	return 0;
      fprintf (stderr, "%s: opendir (\"%s\"): %m\n", progname, dirname);
      exit (EXIT_FAILURE);
    }
  while ((errno = 0, dirent = readdir (dir)))
    {
      char linkname[LINE_MAX], buf[LINE_MAX];
      int linkname_len;
      ssize_t buf_len;

      /* FIXME: POSIX portability.  */
      if ((dirent->d_type != DT_DIR && dirent->d_type != DT_LNK)
	  || (dirent->d_type == DT_DIR && strcmp (dirent->d_name, ".") != 0
	      && strcmp (dirent->d_name, "..") != 0)
	  || (dirent->d_type == DT_LNK && strspn (dirent->d_name, "0123456789")
	      != strlen (dirent->d_name)))
	{
	  /* There is a race, D_TYPE may be uninitialized if stat(2) fails.  */
	  if (dirent->d_type != DT_UNKNOWN)
	    fprintf (stderr, "Unexpected entry \"%s\" (d_type %u)"
			     " on readdir (\"%s\"): %m\n",
		     dirent->d_name, (unsigned) dirent->d_type, dirname);
	  continue;
	}
      if (dirent->d_type == DT_DIR)
	continue;
      linkname_len = snprintf (linkname, sizeof linkname, "%s/%s", dirname,
			       dirent->d_name);
      if (linkname_len <= 0 || linkname_len >= sizeof linkname)
	{
	  fprintf (stderr, "Link content too long: `%s' / `%s'\n",
		   dirent->d_name, dirent->d_name);
	  continue;
	}
      buf_len = readlink (linkname, buf, sizeof buf - 1);
      if (buf_len <= 0 || buf_len >= sizeof buf - 1)
	{
	  if (errno != ENOENT && errno != EACCES)
	    fprintf (stderr, "Error reading link \"%s\": %m\n",
		     linkname);
	  continue;
	}
      buf[buf_len] = 0;
      rc = (*func) (pid, buf);
      if (rc != 0)
	{
	  errno = 0;
	  break;
	}
    }
  if (errno != 0)
    {
      fprintf (stderr, "%s: readdir (\"%s\"): %m\n", progname, dirname);
      exit (EXIT_FAILURE);
    }
  if (closedir (dir) != 0)
    {
      fprintf (stderr, "%s: closedir (\"%s\"): %m\n", progname, dirname);
      exit (EXIT_FAILURE);
    }
  return rc;
}

static void pid_fs_scan (void (*func) (pid_t pid, void *data), void *data)
{
  DIR *dir;
  struct dirent *dirent;

  dir = opendir ("/proc");
  if (dir == NULL)
    {
      perror ("opendir (\"/proc\")");
      exit (EXIT_FAILURE);
    }
  while ((errno = 0, dirent = readdir (dir)))
    {
      /* FIXME: POSIX portability.  */
      if (dirent->d_type != DT_DIR
	  || strspn (dirent->d_name, "0123456789") != strlen (dirent->d_name))
	  continue;
      (*func) (atoi (dirent->d_name), data);
    }
  if (errno != 0)
    {
      perror ("readdir (\"/proc\")");
      exit (EXIT_FAILURE);
    }
  if (closedir (dir) != 0)
    {
      perror ("closedir (\"/proc\")");
      exit (EXIT_FAILURE);
    }
}

static int rip_check (pid_t pid, const char *link)
{
  assert (pid != getpid ());

  return strcmp (link, childptyname) == 0;
}

struct pid
  {
    struct pid *next;
    pid_t pid;
  };
static struct pid *pid_list;

static int pid_found (pid_t pid)
{
  struct pid *entry;

  for (entry = pid_list; entry != NULL; entry = entry->next)
    if (entry->pid == pid)
      return 1;
  return 0;
}

static void pid_record (pid_t pid)
{
  struct pid *entry;

  if (pid_found (pid))
    return;

  entry = malloc (sizeof (*entry));
  if (entry == NULL)
    {
      fprintf (stderr, "%s: malloc: %m\n", progname);
      exit (EXIT_FAILURE);
    }
  entry->pid = pid;
  entry->next = pid_list;
  pid_list = entry;
}

static void pid_forall (void (*func) (pid_t pid))
{
  struct pid *entry;

  for (entry = pid_list; entry != NULL; entry = entry->next)
    (*func) (entry->pid);
}

/* Returns 0 on failure.  */
static pid_t pid_get_parent (pid_t pid)
{
  char fname[64];
  FILE *f;
  char line[LINE_MAX];
  pid_t retval = 0;

  if (snprintf (fname, sizeof fname, "/proc/%d/status", (int) pid) < 0)
    {
      perror ("snprintf(3)");
      exit (EXIT_FAILURE);
    }
  f = fopen (fname, "r");
  if (f == NULL)
    {
      return 0;
    }
  while (errno = 0, fgets (line, sizeof line, f) == line)
    {
      if (strncmp (line, "PPid:\t", sizeof "PPid:\t" - 1) != 0)
	continue;
      retval = atoi (line + sizeof "PPid:\t" - 1);
      errno = 0;
      break;
    }
  if (errno != 0)
    {
      fprintf (stderr, "%s: fgets (\"%s\"): %m\n", progname, fname);
      exit (EXIT_FAILURE);
    }
  if (fclose (f) != 0)
    {
      fprintf (stderr, "%s: fclose (\"%s\"): %m\n", progname, fname);
      exit (EXIT_FAILURE);
    }
  return retval;
}

static void killtree (pid_t pid);

static void killtree_pid_fs_scan (pid_t pid, void *data)
{
  pid_t parent_pid = *(pid_t *) data;

  /* Optimization.  */
  if (pid_found (pid))
    return;

  if (pid_get_parent (pid) != parent_pid)
    return;

  killtree (pid);
}

static void killtree (pid_t pid)
{
  pid_record (pid);
  pid_fs_scan (killtree_pid_fs_scan, &pid);
}

static void rip_pid_fs_scan (pid_t pid, void *data)
{
  pid_t pgid;

  /* Shouldn't happen.  */
  if (pid == getpid ())
    return;

  /* Check both PGID and the stale file descriptors.  */
  pgid = getpgid (pid);
  if (pgid == child
      || fd_fs_scan (pid, rip_check) != 0)
    killtree (pid);
}

static void killproc (pid_t pid)
{
  const char *cmdline;

  cmdline = read_cmdline (pid);
  /* Avoid printing the message for already gone processes.  */
  if (kill (pid, 0) != 0 && errno == ESRCH)
    return;
  if (cmdline == NULL)
    cmdline = "<error>";
  fprintf (stderr, "%s: Killed -9 orphan PID %d: %s\n", progname, (int) pid, cmdline);
  if (kill (pid, SIGKILL))
    {
      if (errno != ESRCH)
	fprintf (stderr, "%s: kill (%d, SIGKILL): %m\n", progname, (int) pid);
      return;
    }
  /* Do not waitpid(2) as it cannot be our direct descendant and it gets
     cleaned up by init(8).  */
#if 0
  pid_t pid_got;
  pid_got = waitpid (pid, NULL, 0);
  if (pid != pid_got)
    {
      fprintf (stderr, "%s: waitpid (%d) != %d: %m\n", progname,
	 (int) pid, (int) pid_got);
      return;
    }
#endif
}

static void rip ()
{
  pid_fs_scan (rip_pid_fs_scan, NULL);
  pid_forall (killproc);
}

int main (int argc, char **argv)
{
  int rc;

  if (argc < 2 || strcmp (argv[1], "-h") == 0
      || strcmp (argv[1], "--help") == 0)
    {
      fputs ("Syntax: orphanripper <execvp(3) commandline>\n", stdout);
      exit (EXIT_FAILURE);
    }
  progname = argv[0];
  rc = spawn (argv);
  rip ();
  return rc;
}