Logo Search packages:      
Sourcecode: mswatch version File versions  Download package

socketwatch.cc

#define _XOPEN_SOURCE 500
#include <assert.h>
#include <glib.h>
#include <libgen.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <signal.h>
#include <time.h>
#include <unistd.h>
#include <fcntl.h>

#include <string>

#include "compat.h"
#include "lib/util.h"
#include "socketwatch.h"

using std::string;


#define STRINGIFY(x) #x
#define TOSTRING(x) STRINGIFY(x)

#define BIN_NAME "socketwatch"
// Default heartbeat timeout (seconds)
#define DEFAULT_TIMEOUT 120
// Location of the LD_PRELOAD library
#define SOCKETWATCH_SO LIBDIR "/libsocketwatch.so"


static pid_t child_pid = -1;
static int sigchld_pipe[2] = {-1, -1};


static void at_exit(void)
{
      // Avoid potential (?) child_pid race
      pid_t child = child_pid;

      if (child >= 0)
      {
            int r;

            child_pid = -1;

            r = kill(child, SIGKILL);
            die_if(r < 0, "kill(pid = %d)", child);

            // assume child is killed. we probably shouldn't block on wait().
      }
}

static void signal_exit(int)
{
      // exit() will call at_exit()
      exit(EXIT_SUCCESS);
}

static void signal_child(int)
{
      // Send a note to the watch loop instead of handling here, because
      // checking child_alive_fd here races with closing the write end.
      char c = 0;
      int r = TEMP_FAILURE_RETRY(write(sigchld_pipe[1], &c, 1));
      die_if(r != 1, "write(sigchld_pipe)");
}

static int watch(int timeout, const time_t *heartbeat)
{
      struct pollfd pollfd;

      pollfd.fd = sigchld_pipe[0];
      pollfd.events = POLLIN;

      while (1)
      {
            int timeout_ms = timeout * 1000;
            int poll_ret;
            int r;

            poll_ret = TEMP_FAILURE_RETRY(poll(&pollfd, 1, timeout_ms));
            die_if(poll_ret < 0, "poll");

            if (!poll_ret)
            {
                  if ((time(NULL) - *heartbeat) > timeout)
                  {
                        pid_t child = child_pid;

                        // Reset before the wait() to avoid stale use by at_exit()
                        child_pid = -1;

                        fprintf(stderr, BIN_NAME ": Not receiving heartbeats,"
                                        " killing child.\n");

                        r = kill(child, SIGKILL);
                        die_if(r < 0, "kill(pid = %d)", child);

                        // assume child is killed.
                        // we probably shouldn't block on wait().

                        exit(EXIT_FAILURE);
                  }
            }
            else if (pollfd.revents)
            {
                  pid_t child = child_pid;
                  pid_t child_ret;
                  int status;
                  int v;

                  // Reset before the wait() to avoid stale use by at_exit()
                  child_pid = -1;

                  child_ret = wait(&status);
                  assert(child == child_ret);

                  if (WIFEXITED(status))
                        v = WEXITSTATUS(status);
                  else
                        v = 1;

                  exit(v);
            }
      }     
}

static void configure_socketwatch(int argc, char **argv, char ***child_command,
                                  int *timeout)
{
      gboolean version = 0;
      GOptionContext *ctx;
      GError *gerror = NULL;
      gboolean r;

      GOptionEntry options[] =
      {
            { "timeout", 't', 0, G_OPTION_ARG_INT, timeout,
              "No network activity timeout (in seconds, default is "
              TOSTRING(DEFAULT_TIMEOUT) ")", NULL },
            { "version", 'V', 0, G_OPTION_ARG_NONE, &version,
              "Display version", NULL },
            { G_OPTION_REMAINING, 0, 0, G_OPTION_ARG_STRING_ARRAY, child_command,
              "Command to run", NULL },
            { NULL, 0, 0, G_OPTION_ARG_NONE, NULL, NULL, NULL }
      };

      ctx = g_option_context_new("<PROGRAM> [ARGUMENT...]\n\t- Run PROGRAM."
                                 " Kill it if it stops using the network.");
      g_option_context_add_main_entries(ctx, options, NULL);
      r = g_option_context_parse(ctx, &argc, &argv, &gerror);
      if (r == FALSE)
      {
            GQuark error = G_OPTION_ERROR;
            assert(error == G_OPTION_ERROR_BAD_VALUE);
            fprintf(stderr, "%s\n", gerror->message);
            g_error_free(gerror);
            exit(EXIT_FAILURE);
      }
      assert(!gerror);
      g_option_context_free(ctx);
      ctx = NULL;

      if (*timeout < 1)
      {
            fprintf(stderr, "Timeout (%d second) must be at least 1 second.\n", *timeout);
            exit(EXIT_FAILURE);
      }

      if (version)
      {
            printf(BIN_NAME "\n");
            printf(PACKAGE " " VERSION "\n");
            exit(EXIT_FAILURE);
      }

      if (!child_command || !*child_command)
      {
            fprintf(stderr, "No child command arguments. Use -? or --help for help.\n");
            exit(EXIT_FAILURE);
      }
}

int main(int argc, char **argv)
{
      char **child_command = NULL;
      int timeout = DEFAULT_TIMEOUT;
      size_t page_size = sysconf(_SC_PAGE_SIZE);
      char shared_name[] = "/dev/shm/socketwatch.XXXXXX";
      int shared_fd;
      void *shared;
      int r;

      configure_socketwatch(argc, argv, &child_command, &timeout);

      //
      // Set up shared heartbeat page

      shared_fd = mkstemp(shared_name);
      die_if(shared_fd < 0, "mkstemp(%s)", shared_name);

      // ftruncate() before unlink() because I think some file systems
      // do not permit file extension after the inode is no longer
      // referenced by a directory entry.
      r = ftruncate(shared_fd, page_size);
      if (r < 0)
      {
            perror("ftruncate");

            r = unlink(shared_name);
            if (r < 0)
                  fprintf(stderr, "unlink(%s)\n", shared_name);

            exit(EXIT_FAILURE);
      }

      r = unlink(shared_name);
      die_if(r < 0, "unlink(%s)", shared_name);

      shared = mmap(NULL, page_size, PROT_READ, MAP_SHARED, shared_fd, 0);
      die_if(shared == MAP_FAILED, "mmap");


      r = pipe(sigchld_pipe);
      die_if(r < 0, "pipe");
      r = fcntl(sigchld_pipe[0], F_SETFL, O_NONBLOCK);
      die_if(r < 0, "fcntl(NONBLOCK)");

      r = atexit(at_exit);
      die_if(r < 0, "atexit");

      // SIGPIPE ignored in case STDOUT_FILENO is a pipe without a reader
      if ((set_signal_handler(SIGHUP, signal_exit) == -1)
          || (set_signal_handler(SIGINT, signal_exit) == -1)
          || (set_signal_handler(SIGTERM, signal_exit) == -1)
          || (set_signal_handler(SIGCHLD, signal_child) == -1)
          || (set_signal_handler(SIGPIPE, SIG_IGN) == -1))
            die_if(1, "set_signal_handler\n");

      // NOTE: it is possible that the child is created, an exit signal
      // is delivered to the parent, and then child_pid is updated.
      // This is racey: the parent may not set child_pid and thus will
      // not kill the child.
      // TODO: perhaps this could be fixed with a pair of pipes? or...
      if (!(child_pid = fork()))
      {
            char *ldpreload_cstr = getenv("LD_PRELOAD");
            string ldpreload = ldpreload_cstr ? ldpreload_cstr : "";
            string shared_fd_str = itos(shared_fd);
            struct rlimit rlimit_zero = {0, 0};

            // sigchld_pipe is only for the parent.
            // However, it must be created prior to the fork.
            r = close(sigchld_pipe[0]);
            die_if(r < 0, "close(sigchld_pipe[0]");
            r = close(sigchld_pipe[1]);
            die_if(r < 0, "close(sigchld_pipe[1]");

            if (ldpreload.size())
                  ldpreload.append(":");
            ldpreload.append(SOCKETWATCH_SO);
            r = setenv("LD_PRELOAD", ldpreload.c_str(), 1);
            die_if(r < 0, "setenv(LD_PRELOAD)");

            r = setenv(SOCKETWATCH_SHARED_FD, shared_fd_str.c_str(), 1);
            die_if(r < 0, "setenv(%s)", SOCKETWATCH_SHARED_FD);

            // Prevent this process from forking, because the parent cannot
            // reliably discover (and thus terminate) our children.
            r = setrlimit(RLIMIT_NPROC, &rlimit_zero);
            die_if(r < 0, "setrlimit(RLIMIT_NPROC, 0)");

            execvp(child_command[0], &child_command[0]);
            error("%s: Unable to execute \"%s\"", BIN_NAME, child_command[0]);
            exit(EXIT_FAILURE);
      }
      die_if(child_pid < 0, "fork");

      return watch(timeout, (const time_t*) shared);
}

Generated by  Doxygen 1.6.0   Back to index