/* $Id: edinplace.c,v 1.11 2006/02/12 22:31:12 dm Exp $ */

/*
 *
 * Copyright (C) 2004 David Mazieres (dm@uun.org)
 *
 * 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, 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
 *
 */

#include "avutil.h"
#include <signal.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/resource.h>
#include "getopt_long.h"

#ifdef NEED_PREAD_DECL
ssize_t pread(int fd, void *buf, size_t size, off_t offset);
#endif /* NEED_PREAD_DECL */
#ifdef NEED_PWRITE_DECL
ssize_t pwrite(int fd, const void *buf, size_t size, off_t offset);
#endif /* NEED_PWRITE_DECL */

char *progname;

struct filter_state {
  int file;
  int filter;
  char file_eof;
  char filter_eof;
  char filter_weof;
  off_t read_offset;
  off_t write_offset;
  size_t read_pos;
  size_t read_lim;
  size_t write_buf_size;
  size_t write_pos;
  size_t write_lim;
  char read_buf[8192];
  char *write_buf;
};
typedef struct filter_state filter_state;
pid_t filter_pid;
int error_exit = 1;

static void
close_on_exec (int s)
{
  if (fcntl (s, F_SETFD, 1) < 0) {
    perror ("F_SETFD");
    exit (error_exit);
  }
}
static int
make_async (int s)
{
  int n;
  if ((n = fcntl (s, F_GETFL)) < 0
      || fcntl (s, F_SETFL, n | O_NONBLOCK) < 0)
    return -1;
  return 0;
}

int
spawnit (int argc, char **argv)
{
  char **av = xmalloc ((argc + 1) * sizeof (*av));
  int socks[2];
  int errpipe[2];

  memcpy (av, argv, argc * sizeof (*av));
  av[argc] = NULL;

  if (socketpair (AF_UNIX, SOCK_STREAM, 0, socks) < 0) {
    perror ("socketpair");
    exit (error_exit);
  }
  close_on_exec (socks[0]);
  close_on_exec (socks[1]);
  make_async (socks[0]);
  if (pipe (errpipe) < 0) {
    perror ("socketpair");
    exit (error_exit);
  }
  close_on_exec (errpipe[0]);
  close_on_exec (errpipe[1]);

  filter_pid = fork ();
  if (filter_pid < 0) {
    perror ("fork");
    exit (error_exit);
  }
  if (!filter_pid) {
    dup2 (socks[1], 0);
    dup2 (socks[1], 1);
    execvp (av[0], av);
    perror (av[0]);
    write (errpipe[1], &errno, sizeof (errno));
    _exit (1);
  }

  close (socks[1]);
  close (errpipe[1]);
  if (read (errpipe[0], &errno, sizeof (errno)) > 0)
    exit (error_exit);
  close (errpipe[0]);
  return socks[0];
}

static void
maybe_write_to_file (filter_state *sp)
{
  ssize_t n;

  n = sp->write_lim - sp->write_pos;
  if (!sp->file_eof) {
    off_t d = sp->read_offset - sp->write_offset;
    if (n > d)
      n = d;
  }

  if (n > 0) {
    n = pwrite (sp->file, sp->write_buf + sp->write_pos, n, sp->write_offset);
    if (n < 0) {
      perror ("file");
      exit (error_exit);
      sp->write_offset += n;
      sp->write_pos += n;
    }
    sp->write_pos += n;
    sp->write_offset += n;
  }
}

void
tofilter (filter_state *sp)
{
  int n;

  if (!sp->file_eof && sp->read_pos == sp->read_lim) {
    sp->read_pos = sp->read_lim = 0;
    n = pread (sp->file, sp->read_buf, sizeof (sp->read_buf), sp->read_offset);
    if (n > 0) {
      sp->read_lim = n;
      sp->read_offset += n;
    }
    else {
      sp->file_eof = 1;
      sp->read_offset = -1;
    }
    maybe_write_to_file (sp);
  }

  if (sp->read_lim > sp->read_pos) {
    n = write (sp->filter, sp->read_buf + sp->read_pos,
	       sp->read_lim - sp->read_pos);
    if (n < 0 && errno != EAGAIN) {
      if (errno != EPIPE)
	perror ("filter");
      sp->filter_weof = 1;
      return;
    }
    else if (n >= 0)
      sp->read_pos += n;
  }

  if (sp->file_eof && sp->read_pos == sp->read_lim) {
    shutdown (sp->filter, 1);
    sp->filter_weof = 1;
  }
}

void
fromfilter (filter_state *sp)
{
  ssize_t n;

  if (sp->filter_eof)
    return;

  if (sp->write_pos == sp->write_lim)
    sp->write_pos = sp->write_lim = 0;
  else if (sp->write_pos > 0) {
    memmove (sp->write_buf, sp->write_buf + sp->write_pos,
	     sp->write_lim - sp->write_pos);
    sp->write_lim -= sp->write_pos;
    sp->write_pos = 0;
  }

  if (sp->write_lim == sp->write_buf_size) {
    sp->write_buf_size *= 2;
    sp->write_buf = realloc (sp->write_buf, sp->write_buf_size);
    if (!sp->write_buf) {
      fprintf (stderr, "out of memory\n");
      abort ();
    }
  }

  n = read (sp->filter, sp->write_buf + sp->write_lim,
	    sp->write_buf_size - sp->write_lim);
  if (n > 0) {
    sp->write_lim += n;
    maybe_write_to_file (sp);
  }
  else if (!n || errno != EAGAIN)
    sp->filter_eof = 1;
}

void
dofilter (int file, int filter, off_t start_offset)
{
  int n;
  filter_state state;
  fd_set rfds, wfds;

  bzero (&state, sizeof (state));
  state.read_offset = state.write_offset = start_offset;
  state.file = file;
  state.filter = filter;
  state.write_buf_size = 8192;
  state.write_buf = xmalloc (state.write_buf_size);

  FD_ZERO (&rfds);
  FD_ZERO (&wfds);

  for (;;) {
    if (state.filter_eof)
      FD_CLR (state.filter, &rfds);
    else
      FD_SET (state.filter, &rfds);
    if (state.filter_weof)
      FD_CLR (state.filter, &wfds);
    else
      FD_SET (state.filter, &wfds);

    n = select (1 + (state.filter > state.file ? state.filter : state.file),
		&rfds, &wfds, NULL, NULL);
    if (n < 0) {
      perror ("select");
      exit (error_exit);
    }
    if (FD_ISSET (state.filter, &wfds))
      tofilter (&state);
    else if (FD_ISSET (state.filter, &rfds))
      fromfilter (&state);

    if (state.filter_eof && state.filter_weof) {
      int status;
      ftruncate (state.file, state.write_offset);
      lseek (state.file, 0, SEEK_SET);
      close (state.file);
      close (state.filter);
      if (waitpid (filter_pid, &status, 0) < 0) {
	perror ("waitpid");
	exit (error_exit);
      }
      if (WIFEXITED (status))
	exit (WEXITSTATUS (status));
      if (WIFSIGNALED (status)) {
	struct rlimit rl;
	bzero (&rl, sizeof (rl));
	setrlimit (RLIMIT_CORE, &rl);
	raise (WTERMSIG (status));
      }
      exit (error_exit);
    }
  }
}

/* Assumes fd is already at offset zero; leaves offset at arbitrary
 * location. */
off_t
skipfrom (int fd)
{
  char buf[1024], *e;
  int n;
  off_t pos = 0;

  while (pos < 5) {
    n = read (fd, buf + pos, sizeof (buf) - pos);
    if (n == 0)
      return 0;
    if (n < 0)
      return -1;
    pos += n;
  }
  if (strncmp (buf, "From ", 5))
    return 0;

  n = pos;
  pos = 0;
  for (;;) {
    if ((e = memchr (buf, '\n', n)))
      return pos + (e - buf) + 1;
    pos += n;
    n = read (fd, buf, sizeof (buf));
    if (n < 0)
      return -1;
    if (n == 0)
      return pos;
  }
}

static void usage (void) __attribute__ ((noreturn));
static void
usage (void)
{
  fprintf (stderr,
	   "usage: %s [--skipfrom] [-x err_exit] [[-f file]"
	   " filter [arg ...]]\n",
	   progname);
  exit (error_exit);
}

int
main (int argc, char **argv)
{
  int c;
  struct option o[] = {
    { "skipfrom", no_argument, NULL, 'S' },
    { "file", required_argument, NULL, 'f' },
    { "error", required_argument, NULL, 'x' },
    { "version", no_argument, NULL, 'v' },
    { NULL, 0, NULL, 0 }
  };
  char *opt_file = NULL;
  int opt_skipfrom = 0;
  int fd = 0;
  int fil = -1;
  off_t start_offset = 0;

  progname = strrchr (argv[0], '/');
  if (progname)
    progname++;
  else
    progname = argv[0];

  while ((c = getopt_long (argc, argv, "+f:x:", o, NULL)) != -1)
    switch (c) {
    case 'S':
      opt_skipfrom = 1;
      break;
    case 'f':
      if (opt_file)
	usage ();
      opt_file = optarg;
      break;
    case 'x':
      error_exit = atoi (optarg);
      if (error_exit <= 0 || error_exit > 255)
	usage ();
      break;
    case 'v':
      version (progname, 1);
      break;
    default:
      usage ();
      break;
    }

  argv += optind;
  argc -= optind;

  if (opt_file) {
    if (!argc)
      usage ();
    fd = open (opt_file, O_RDWR, 0);
    if (fd < 0) {
      perror (opt_file);
      exit (error_exit);
    }
  }

  if (lseek (fd, 0, SEEK_SET) == -1) {
    fprintf (stderr, "cannot seek in input file (%s)\n", strerror (errno));
    exit (error_exit);
  }
  if (opt_skipfrom) {
    start_offset = skipfrom (fd);
    if (start_offset == -1) {
      fprintf (stderr, "cannot read input file (%s)\n", strerror (errno));
      exit (error_exit);
    }
  }

  if (!argc) {
    if (lseek (fd, start_offset, SEEK_SET) == -1) {
      fprintf (stderr, "cannot seek in input file (%s)\n", strerror (errno));
      exit (error_exit);
    }
    exit (0);
  }

  fil = spawnit (argc, argv);

  signal (SIGPIPE, SIG_IGN);

  dofilter (fd, fil, start_offset);

  return 0;
}