/* Copyright (C) 2010 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 3 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, see <http://www.gnu.org/licenses/>.  */

#include <error.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <libintl.h>
#include <stdio.h>
#include <string.h>
#include <signal.h>
#include <stdarg.h>
#include <fcntl.h>
#include <poll.h>
#include <assert.h>

#ifndef STAP_SCRIPT_FILENAME
#error "STAP_SCRIPT_FILENAME is required"
#endif

#define STRINGIFY1(lit) #lit
#define STRINGIFY(lit) STRINGIFY1 (lit)
#define _(x) gettext (x)

static int opt_d, opt_f, opt_q, opt_p;
/* Verified we should run the -ff postprocessor.  */
static int opt_ff;
static char *opt_o;
static size_t opt_o_len;

/* Script filename.  */
static char *opt_Z = STRINGIFY (STAP_SCRIPT_FILENAME);

static void *
xmalloc (size_t size)
{
  void *retval = malloc (size);

  if (retval == NULL)
    error (EXIT_FAILURE, errno, _("memory allocation of %zu bytes"), size);

  return retval;
}

/* warning: EOF is also an error.  */
static size_t
xread (int fd, void *buf, size_t count)
{
  ssize_t retval = read (fd, buf, count);

  if (retval <= 0)
    error (EXIT_FAILURE, errno, _("Error reading data"));

  return retval;
}

/* warning: No events are also an error.  */
static void
xpoll (struct pollfd *fds, nfds_t nfds, int timeout)
{
  if (poll (fds, nfds, timeout) != 1)
    error (EXIT_FAILURE, errno, _("poll returned error"));
}

static pid_t
xfork (void)
{
  pid_t retval = fork ();

  if (retval == -1)
    error (EXIT_FAILURE, errno, _("Error calling fork"));

  return retval;
}

static void
xsnprintf (char *str, size_t size, const char *format, ...)
{
  int i;
  va_list ap;

  va_start (ap, format);
  i = vsnprintf (str, size, format, ap);
  va_end (ap);
  if (i <= 0 || strlen (str) >= size - 1)
    error (EXIT_FAILURE, errno, _("Error calling sprintf"));
}

static void
xfflush (FILE *stream)
{
  if (fflush (stream) != 0)
    error (EXIT_FAILURE, errno, _("Error calling fflush"));
}

static void
xdup2 (int oldfd, int newfd)
{
  if (dup2 (oldfd, newfd) != newfd)
    error (EXIT_FAILURE, errno, _("Error calling dup2 (%d, %d)"), oldfd,
	   newfd);
}

static void
xclose (int fd)
{
  if (close (fd) != 0)
    error (EXIT_FAILURE, errno, _("Error calling close (%d)"), fd);
}

static void
xpipe (int pipefd[2])
{
  if (pipe (pipefd) != 0)
    error (EXIT_FAILURE, errno, _("Error creating internal pipe"));
}

static void
xraise (int sig)
{
  if (raise (sig) != 0)
    error (EXIT_FAILURE, errno, _("Error calling raise (%d)"), sig);
}

static void
xfcntl_setfl (int fd, int cmdarg)
{
  if (fcntl (fd, F_SETFL, cmdarg) != 0)
    error (EXIT_FAILURE, errno,
	   _("Error calling fcntl (%d, F_SETFL, 0%o)"), fd, cmdarg);
}

static void
xwrite (int fd, const void *buf, size_t count)
{
  if (!count)
    return;

  if ((ssize_t) count < 0 || write (fd, buf, count) != (ssize_t) count)
    error (EXIT_FAILURE, errno,
	   _("Error calling write (%d, ..., %zu)"), fd, count);
}

static int
xopen (const char *pathname, int flags, mode_t mode)
{
  int retval = open (pathname, flags, mode);

  if (retval == -1)
    error (EXIT_FAILURE, errno,
	   _("Error %s file \"%s\" (flags 0x%x, mode 0%o)"),
	   flags & O_APPEND ? _("appending")
			    : (flags & O_CREAT ? _("creating")
					       : _("opening")),
	   pathname, flags, mode);

  return retval;
}

static void
usage (int status)
{
  printf ("\
usage: staptrace [-dfq] [-o file] [-p pid] [command [arg ...]]\n\
                 [-Z <stap script>]\n\
\n\
Default -Z is: %s\n\
",
	  opt_Z);
  exit (status);
}

static void
dump_args (char **args)
{
  char **argp;

  for (argp = args; *argp; argp++)
    {
      if (argp > args)
	fputc (' ', stderr);
      fputs (*argp, stderr);
    }
  fputc ('\n', stderr);
}

static void
ff_filter (int fd_read)
{
  char buf[getpagesize ()];
  size_t buf_len = 0;
  /* Does the buf start on a line boundary?  */
  int buf_nl = 1;
  /* Initial fd if no `\n[' is seen at the start.  */
  int fd_write = STDOUT_FILENO;

  xfcntl_setfl (fd_read, O_NONBLOCK);

  for (;;)
    {
      struct pollfd pollfd;

      pollfd.fd = fd_read;
      pollfd.events = POLLIN;

      xpoll (&pollfd, 1, -1);

      if (pollfd.revents & ~(POLLIN | POLLHUP) || pollfd.revents == 0)
	error (EXIT_FAILURE, errno, _("poll returned error revents 0x%x"),
	       pollfd.revents);

      if (pollfd.revents & POLLIN)
	{
	  ssize_t want, got;
	  char *s;
	  char *start;
	  int start_nl;

	  want = sizeof (buf) - buf_len - 1;
	  assert (want > 0);
	  got = xread (fd_read, &buf[buf_len], want);
	  if (memchr (&buf[buf_len], 0, got) != NULL)
	    error (EXIT_FAILURE, 0,
		   _("Unexpected \\0 in the internal -ff stream"));
	  buf_len += got;
	  buf[buf_len] = 0;
	  if (opt_d >= 2)
	    fprintf (stderr, "buf_len %zu + %zd = %zu, read: {%s}\n",
		     buf_len - got, got, buf_len, &buf[buf_len - got]);
	  start = buf;
	  start_nl = buf_nl;

	  /* Keep the last character for `c' or `n' line continuations.  */
	  while (start + 1 < &buf[buf_len])
	    {
	      unsigned long ul;
	      char *end;
	      static char fname[FILENAME_MAX];

	      if (!start_nl)
		{
		  ssize_t len;

		  s = strchr (start, '\n');
		  if (!s)
		    {
		      len = strlen (start);
		      assert (len > 0);
		      if (start[len - 1] == 'c' || start[len - 1] == 'n')
			len--;
		      xwrite (fd_write, start, len);
		      start += len;
		      continue;
		    }
		  len = s - start - 1;
		  assert (len >= 0);
		  if (start[len] == 'n')
		    start[len++] = '\n';
		  xwrite (fd_write, start, len);
		  start = s + 1;
		  start_nl = 1;
		  continue;
		}

	      if (*start != '[')
		error (EXIT_FAILURE, 0,
		       _("'[' expected in the internal -ff stream"));
	      if (start[1] == 0)
		{
		  /* We need more data.  */
		  break;
		}
	      ul = strtoul (&start[1], &end, 10);
	      assert (end != NULL);
	      if (*end == 0)
		{
		  /* We need more data.  */
		  break;
		}
	      if (end > start + 64 || *end != ']' || end == &start[1])
		error (EXIT_FAILURE, 0,
		       _("Invalid TID in the internal -ff stream: %s"), start);
	      start = end + 1;
	      start_nl = 0;

	      xsnprintf (fname, sizeof (fname), "%s.%lu", opt_o, ul);
	      if (fd_write != STDOUT_FILENO)
		xclose (fd_write);
	      fd_write = xopen (fname, O_WRONLY | O_APPEND | O_CREAT, 0644);

	      continue;
	    }

	  want = &buf[buf_len] - start;
	  memmove (buf, start, want);
	  buf_len = want;
	  buf_nl = start_nl;
	  continue;
	}
      if (pollfd.revents & POLLHUP)
	break;
      assert (0);
    }
  if (fd_write != STDOUT_FILENO)
    xclose (fd_write);
}

int
main (int argc, char **argv)
{
  char **args, **argp;

  args = xmalloc (sizeof (*args) * argc * 2 + 32);
  argp = args;
  *argp++ = "stap";

  for (;;)
    {
      int i;

      i = getopt (argc, argv, "+dfqo:p:hZ:");
      if (i == EOF)
	break;
      switch (i)
	{
	case 'd':
	  opt_d++;
	  break;
	case 'f':
	  opt_f++;
	  break;
	case 'q':
	  *argp++ = "-G";
	  *argp++ = "opt_q=1";
	  opt_q = 1;
	  break;
	case 'o':
	  opt_o = optarg;
	  opt_o_len = strlen (opt_o);
	  break;
	case 'p':
	  *argp++ = "-x";
	  *argp++ = optarg;
	  opt_p = 1;
	  break;
	case 'Z':
	  opt_Z = optarg;
	  break;
	case 'h':
	  usage (EXIT_SUCCESS);
	default:
	  usage (EXIT_FAILURE);
	}
    }

  if ((optind == argc) == !opt_p)
    usage (EXIT_FAILURE);

  if (optind < argc)
    {
      pid_t child;
      static char kill_s[256];
      static char option_s[256];

      child = xfork ();
      if (!child)
	{
	  xraise (SIGSTOP);

	  execvp (argv[optind], &argv[optind]);
	  error (EXIT_FAILURE, errno, _("Error executing child command"));
	}

      /* FIXME: Call `kill -CONT' from the systamtap script.  How?  */
      xsnprintf (kill_s, sizeof (kill_s),
		 "kill -CONT %lu; "
		 "while [ -e /proc/%lu/root ]; do sleep 0.1; done",
		 (unsigned long) child, (unsigned long) child);
      *argp++ = "-c";
      *argp++ = kill_s;
      xsnprintf (option_s, sizeof (option_s), "opt_child=%lu",
		 (unsigned long) child);
      *argp++ = "-G";
      *argp++ = option_s;
    }

  xfflush (stdout);
  xdup2 (STDERR_FILENO, STDOUT_FILENO);

  switch (opt_f)
    {
    case 1:
      *argp++ = "-G";
      *argp++ = "opt_f=1";
      /* FALLTHRU */
    case 0:
      if (opt_o)
	{
	  *argp++ = "-o";
	  *argp++ = opt_o;
	}
      break;
    default:
      *argp++ = "-G";
      *argp++ = "opt_f=1";
      if (opt_o)
	{
	  *argp++ = "-G";
	  *argp++ = "opt_ff=1";
	  opt_ff = 1;
	}
      else
	{
	  /* -ff is the same like -f when no -o is specified.  */
	}
      break;
    }

  *argp++ = opt_Z;
  *argp = NULL;

  if (opt_d)
    dump_args (args);

  if (opt_ff)
    {
      int pipefd[2];

      xpipe (pipefd);

      if (xfork ())
	{
	  xclose (pipefd[1]);
	  ff_filter (pipefd[0]);
	  exit (EXIT_SUCCESS);
	}

      xclose (pipefd[0]);
      xdup2 (pipefd[1], STDOUT_FILENO);
      xclose (pipefd[1]);
    }

  execvp ("stap", args);

  if (!opt_d)
    dump_args (args);
  error (EXIT_FAILURE, errno, _("Error executing stap"));
}