/*
 * copytolog - designed for things that produce logfile output on
 *  stdout, but where you want to be able to rotate the logfile without
 *  killing and restarting the process.
 *
 * Usage: copytolog [options] logfile
 *
 * Read stdin.  Input is accumulated into lines; when a newline is
 *  read (or when a buffer fills), the logfile is opened (creating it
 *  if necessary), the accumulated input is written to it, and it's
 *  closed.  Thus, instead of
 *	daemon >> logfile &
 *  you would run
 *	daemon | copytolog logfile &
 *
 * To cut down on syscall overhead, copytolog will, if multiple lines
 *  are available on its input, write multiple lines at once to the
 *  logfile.  It will never write a line without its terminating
 *  newline unless either (a) the input fills up its internal buffer
 *  before it sees a newline, (b) it sees input without a newline and
 *  doesn't see a newline soon enough, or (c) it sees EOF on its input
 *  (which always forces it to flush any partial line).  The output
 *  logfile is never left open when waiting; the logfile is opened only
 *  when copytolog decides to write and is closed immediately after
 *  writing.
 *
 * copytolog blocks all signals it can.  It normally exits when it sees
 *  EOF on its standard input.  All other means of exiting (short of
 *  internal bugs) will produce an explanatory message on stderr and a
 *  nonzero exit code.  The other things that can cause it to exit are:
 *	- Startup failures, such as bad arguments or inability to
 *	   malloc() memory;
 *	- An error (as opposed to EOF) when reading from stdin;
 *	- An error is returned from select().
 *
 * Failure to open/create the output logfile will produce a message on
 *  stderr and copytolog will buffer the line; it will retry the open
 *  once a minute, until either the open succeeds or the error code
 *  returned from open() changes.  In the latter case it will report
 *  the new error and continue retrying.  While it is retrying, input
 *  is read and buffered up to the buffer size limit.  If the buffer
 *  size limit is reached and more input is read, copytolog continues
 *  reading input, throwing away old input to make room for the new,
 *  and if/when the logfile open succeeds, it generates a synthetic
 *  line saying that input was lost due to buffer overflow.  EOF on
 *  stdin while copytolog is retrying does not cause an immediate exit,
 *  but copytolog will exit immediately if it eventually does manage to
 *  write its buffered data in this case.
 *
 * Write errors on the logfile produce messages to stderr, a new
 *  message each time an error occurs.  When this happens, the data
 *  copytolog was attempting to write is lost.
 *
 * The options can be the following.  For defaults, see the source code
 *  immediately after this comment.
 *
 *	-timefmt fmt
 *		Causes each line of input to have a timestamp prepended
 *		to it.  fmt is a strftime format for the timestamp,
 *		with the additional format %f representing hundredths
 *		of a second.  The timestamp is obtained from
 *		gettimeofday(), called immediately after the read()
 *		that returned the first byte of the line (which may not
 *		be the same as the read() that returned the last byte
 *		of the line).  If copytolog flushes a partial line due
 *		to a timeout, it does not produce another timestamp
 *		when it finally does get the rest of that line.
 *
 *		Note that nothing is inserted between the last
 *		character generated by fmt and the first character of
 *		the line; fmt will normally end with whitespace, but it
 *		doesn't have to.
 *
 *	-timestamp
 *		Like "-timefmt %Y%m%d%H%M%S.%f ".
 *
 *	-mode mode
 *	-perm mode
 *		Specifies the mode bits with which to create the
 *		logfile, if it doesn't exist.  This value is always
 *		modified by the umask (see umask(2) and open(2)).  Note
 *		that the mode is always taken as an octal number; no
 *		leading zero is necessary.
 *
 *	-bufsize nbytes
 *		Specifies the size of the internal buffer referred to
 *		above.
 *
 *	-timeout milliseconds
 *		Specifies the time, in milliseconds, that copytolog
 *		should wait after reading a partial line in the hope
 *		that the rest of the line follows soon.  If the
 *		argument is -1, the timeout behavior described above is
 *		disabled; copytolog will wait forever for the rest of
 *		the line, never flushing a partial line unless forced
 *		to by its buffer filling up.  Other negative arguments
 *		are errors.
 *
 *	-logfile logfile
 *		Specifies the logfile pathname, same as giving it
 *		without the -logfile before it, but is unambiguous in
 *		case the pathname does or might begin with a -.
 *
 *	-flock
 *	-fcntl
 *		Specify that copytolog is to lock the logfile around
 *		writes.  -flock uses flock(2) with argument LOCK_EX;
 *		-fcntl uses fcntl(2)'s F_SETLKW operation with type
 *		F_WRLCK, applied to the entire file (l_start=0
 *		l_len=0).  copytolog always does blocking locks; if the
 *		attempt to get the lock blocks, copytolog will not read
 *		further input while waiting.  If both -flock and -fcntl
 *		are specified, whichever one appears first on the
 *		command line is done first, even if that one also
 *		appears last on the command line.
 */

/* Defaults for the options */
static const char *timestamp_fmt = 0;
static int creation_mode = 0666;
static int bufsize = 8192;
static int timeout = 1000;
#define LOCK_NONE   0
#define LOCK_FIRST  1
#define LOCK_SECOND 2
static int flock_lock = LOCK_NONE;
static int fcntl_lock = LOCK_NONE;

#include <time.h>
#include <poll.h>
#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <strings.h>
#include <sys/uio.h>
#include <sys/sysctl.h>

extern const char *__progname;

static const char *logfile = 0;

/* input buffer and fill value */
static char *ibuf;
static int ibfill;
/* output buffer and head and tail pointers (this is a ring buffer) */
static char *obuf;
static int obhead;
static int obtail;
static int obfill;
/* boolean: next byte read first of new line? */
static int atnl;
/* boolean: have we read EOF? */
static int goteof;
/* boolean: partial-line timeout running? */
static int timing_out;
/* boolean: last open() attempt failed? */
static int open_failed;
/* boolean: dropped data due to buffer overrun? */
static int lost_data;
/* errno from last failed open attempt */
static int open_errno;
/* nil if -timefmt doesn't use %f, buffer for %f processing if it does */
static char *stamptmp;
/* buffer to generate timestamps in */
static char *stamp;
/* amount of space allocated to stamp */
static int stampalloc;
/* current length of stamp */
static int stamplen;
/* kern.iov_max, capped to a reasonable maximum */
static int iov_max;
/* ring buffer of iovecs */
static struct iovec *iovv;
/* head and tail indices and fill counter for iovv */
static int iovh;
static int iovt;
static int iovf;
/* byte count for iovv */
static int iovnb;
/* boolean: tried to write already this bufferful? */
static int tried_write;

static __inline__ int inc_iov(int x)
{
 x ++;
 if (x >= iov_max) x = 0;
 return(x);
}

static void usage(void)
{
 fprintf(stderr,"Usage: %s [options] logfile\n",__progname);
 fprintf(stderr,"Options: -timefmt format\n");
 fprintf(stderr,"         -timestamp\n");
 fprintf(stderr,"         -mode octal-mode   (also -perm)\n");
 fprintf(stderr,"         -bufsize nbytes\n");
 fprintf(stderr,"         -timeout milliseconds\n");
 fprintf(stderr,"         -logfile logfile\n");
 fprintf(stderr,"         -flock\n");
 fprintf(stderr,"         -fcntl\n");
}

static int setint(const char *arg, int *vp, int base, const char *tag)
{
 char *cp;
 long int v;

 v = strtol(arg,&cp,base);
 if ((cp == arg) || *cp)
  { fprintf(stderr,"%s: invalid %s value `%s'\n",__progname,tag,arg);
    return(1);
  }
 *vp = v;
 return(0);
}

static void set_locking(int *thisvar, int *othervar)
{
 if (*thisvar != LOCK_NONE) return;
 if (*othervar == LOCK_NONE)
  { *thisvar = LOCK_FIRST;
  }
 else
  { *thisvar = LOCK_SECOND;
  }
}

static void handleargs(int ac, char **av)
{
 int skip;
 int errs;

 skip = 0;
 errs = 0;
 for (ac--,av++;ac;ac--,av++)
  { if (skip > 0)
     { skip --;
       continue;
     }
    if (**av != '-')
     { if (! logfile)
	{ logfile = *av;
	}
       else
	{ fprintf(stderr,"%s: extra argument `%s'\n",__progname,*av);
	  errs ++;
	}
       continue;
     }
    if (0)
     {
needarg:;
       fprintf(stderr,"%s: %s needs a following argument\n",__progname,*av);
       errs ++;
       continue;
     }
#define WANTARG() do { if (++skip >= ac) goto needarg; } while (0)
    if (!strcmp(*av,"-logfile"))
     { WANTARG();
       if (logfile)
	{ fprintf(stderr,"%s: %s %s: logfile already specified\n",__progname,*av,av[skip]);
	  errs ++;
	}
       else
	{ logfile = av[skip];
	}
       continue;
     }
    if (!strcmp(*av,"-timefmt"))
     { WANTARG();
       timestamp_fmt = av[skip];
       continue;
     }
    if (!strcmp(*av,"-timestamp"))
     { timestamp_fmt = "%Y%m%d%H%M%S.%f ";
       continue;
     }
    if (!strcmp(*av,"-mode") || !strcmp(*av,"-perm"))
     { WANTARG();
       errs += setint(av[skip],&creation_mode,8,"mode");
       continue;
     }
    if (!strcmp(*av,"-bufsize"))
     { WANTARG();
       errs += setint(av[skip],&bufsize,0,"buffer size");
       continue;
     }
    if (!strcmp(*av,"-timeout"))
     { WANTARG();
       errs += setint(av[skip],&timeout,0,"timeout");
       continue;
     }
    if (!strcmp(*av,"-flock"))
     { set_locking(&flock_lock,&fcntl_lock);
       continue;
     }
    if (!strcmp(*av,"-fcntl"))
     { set_locking(&fcntl_lock,&flock_lock);
       continue;
     }
#undef WANTARG
    fprintf(stderr,"%s: unrecognized option `%s'\n",__progname,*av);
    errs ++;
  }
 if (errs)
  { usage();
    exit(1);
  }
}

static void percent_f(const char *from, char *to, int n)
{
 char c;

 while (1)
  { c = *from++;
    switch (c)
     { default:
	  *to++ = c;
	  continue;
	  break;
       case '\0':
	  *to = '\0';
	  return;
	  break;
       case '%':
	  c = *from++;
	  switch (c)
	   { case '\0':
		*to++ = '%';
		*to = '\0';
		return;
		break;
	     case 'f':
		*to++ = '0' + (n / 10);
		*to++ = '0' + (n % 10);
		break;
	     default:
		*to++ = '%';
		*to++ = c;
		break;
	   }
	  break;
     }
  }
}

static void stamp_setup(void)
{
 if (! timestamp_fmt) return;
 stamptmp = malloc(strlen(timestamp_fmt)+1);
 percent_f(timestamp_fmt,stamptmp,0);
 if (!strcmp(timestamp_fmt,stamptmp))
  { free(stamptmp);
    stamptmp = 0;
  }
 stamp = malloc(stampalloc=16);
}

static void gen_init(void)
{
 if (iov_max < 0)
  { int mib[2];
    size_t len;
    mib[0] = CTL_KERN;
    mib[1] = KERN_IOV_MAX;
    len = sizeof(iov_max);
    if (sysctl(&mib[0],2,&iov_max,&len,0,0) < 0)
     { fprintf(stderr,"%s: can't get kern.iov_max: %s\n",__progname,strerror(errno));
       exit(1);
     }
    if (iov_max > 64) iov_max = 64;
    if (iov_max < 2)
     { fprintf(stderr,"%s: kern.iov_max unreasonable (%d)\n",__progname,iov_max);
       exit(1);
     }
    iovv = malloc(iov_max*sizeof(struct iovec));
  }
 iovt = 0;
 iovh = 0;
 iovf = 0;
 iovnb = 0;
 tried_write = 0;
}

static void drop_written(int nb)
{
 struct iovec *i;

 while (nb > 0)
  { if (iovf < 1) abort();
    i = &iovv[iovt];
    if (nb >= i->iov_len)
     { nb -= i->iov_len;
       iovnb -= i->iov_len;
       iovt = inc_iov(iovt);
       iovf --;
     }
    else
     { i->iov_base = ((char *)i->iov_base) + nb;
       i->iov_len -= nb;
       iovnb -= nb;
       nb = 0;
     }
  }
}

static void try_flush(void)
{
 int fd;
 int w;

 fd = open(logfile,O_WRONLY|O_APPEND|O_CREAT,creation_mode);
 if (fd < 0)
  { if (!open_failed || (errno != open_errno))
     { open_errno = errno;
       open_failed = 1;
       fprintf(stderr,"%s: open %s: %s\n",__progname,logfile,strerror(open_errno));
     }
    tried_write = 1;
    return;
  }
 open_failed = 0;
 if (!obfill && !iovf) abort();
 do <"ret">
  { if (obfill)
     { struct iovec iov[2];
       int niov;
       if (obhead <= obtail)
	{ if (obhead+bufsize-obtail != obfill) abort();
	  iov[0] = (struct iovec) { .iov_base = obuf+obtail, .iov_len = bufsize-obtail };
	  if (obhead > 0)
	   { iov[1] = (struct iovec) { .iov_base = obuf, .iov_len = obhead };
	     niov = 2;
	   }
	  else
	   { niov = 1;
	   }
	}
       else
	{ if (obhead-obtail != obfill) abort();
	  iov[0] = (struct iovec) { .iov_base = obuf+obtail, .iov_len = obfill };
	  niov = 1;
	}
       w = writev(fd,&iov[0],niov);
       if (w < 0)
	{ fprintf(stderr,"%s: write to %s: %s\n",__progname,logfile,strerror(errno));
	  break <"ret">;
	}
       if (w == obfill)
	{ obhead = 0;
	  obtail = 0;
	  obfill = 0;
	}
       else
	{ obtail += w;
	  if (obtail >= bufsize) obtail -= bufsize;
	  obfill -= w;
	  break <"ret">;
	}
     }
    if (iovf)
     { if (iovh <= iovt)
	{ w = writev(fd,iovv+iovt,iov_max-iovt);
	  if (w < 0)
	   { fprintf(stderr,"%s: write to %s: %s\n",__progname,logfile,strerror(errno));
	     break <"ret">;
	   }
	  drop_written(w);
	  if ((iovf < 1) || (iovh <= iovt)) break <"ret">;
	}
       if (iovh <= iovt) abort();
       w = writev(fd,iovv+iovt,iovh-iovt);
       if (w < 0)
	{ fprintf(stderr,"%s: write to %s: %s\n",__progname,logfile,strerror(errno));
	  break <"ret">;
	}
       drop_written(w);
     }
  } while (0);
 close(fd);
 if (iovf) tried_write = 1;
}

/* Push iovv[iovt] to the obuf ring buffer.  If this involves dropping
   old data, set lost_data. */
static void gen_push_ring(void)
{
 struct iovec *iov;

 iov = iovv + iovt;
 iovnb -= iov->iov_len;
 if (iov->iov_len >= bufsize)
  { lost_data = (obfill || (iov->iov_len > bufsize));
    obhead = 0;
    obtail = 0;
    obfill = bufsize;
    bcopy(((char *)iov->iov_base)+(iov->iov_len-bufsize),obuf,bufsize);
  }
 else
  { if (obfill+iov->iov_len > bufsize)
     { int n;
       lost_data = 1;
       n = obfill + iov->iov_len - bufsize;
       obtail += n;
       if (obtail >= bufsize) obtail -= bufsize;
       obfill -= n;
     }
    if (obhead+iov->iov_len > bufsize)
     { bcopy(iov->iov_base,obuf+obhead,bufsize-obhead);
       bcopy(((char *)iov->iov_base)+(bufsize-obhead),obuf,iov->iov_len-(bufsize-obhead));
       obhead += iov->iov_len - bufsize;
     }
    else
     { bcopy(iov->iov_base,obuf+obhead,iov->iov_len);
       obhead += iov->iov_len;
     }
    obfill += iov->iov_len;
  }
 iovt = inc_iov(iovt);
 iovf --;
}

static void gen_data(char *data, int len)
{
 if (len < 1) return;
 if (iovf >= iov_max)
  { if (! tried_write) try_flush();
    if (iovf >= iov_max) gen_push_ring();
    if (iovf >= iov_max) abort();
  }
 iovv[iovh] = (struct iovec) { .iov_base = data, .iov_len = len };
 iovh = inc_iov(iovh);
 iovf ++;
 iovnb += len;
}

static void gen_flush(void)
{
 if (iovf)
  { if (! tried_write) try_flush();
    while (iovf) gen_push_ring();
  }
}

static void gen_writepush(char *data, int len)
{
 gen_data(data,len);
 if ((obfill+iovnb > bufsize) && !tried_write) try_flush();
 while (iovf) gen_push_ring();
}

static void gen_stamp(void)
{
 struct timeval nowtv;
 time_t nowtt;
 struct tm *nowtm;
 const char *f;
 int rv;

 if (! timestamp_fmt)
  { stamp = 0;
    stamplen = 0;
    return;
  }
 gettimeofday(&nowtv,0);
 nowtt = nowtv.tv_sec;
 nowtm = localtime(&nowtt);
 if (stamptmp)
  { percent_f(timestamp_fmt,stamptmp,nowtv.tv_usec/10000);
    f = stamptmp;
  }
 else
  { f = timestamp_fmt;
  }
 while (1)
  { rv = strftime(stamp,stampalloc,f,nowtm);
    if (! ((rv == 0) && *f && (stampalloc < 2048))) break;
    free(stamp);
    stamp = malloc(stampalloc*=2);
  }
 stamplen = rv;
}

static void try_read(void)
{
 int x;
 char *nl;
 int havestamp;

 static void do_stamp(void)
  { if (! havestamp)
     { gen_stamp();
       havestamp = 1;
     }
    gen_data(stamp,stamplen);
  }

 ibfill = read(0,ibuf,bufsize);
 if (ibfill < 0)
  { switch (errno)
     { case EINTR:
       case EWOULDBLOCK:
	  return;
	  break;
     }
    fprintf(stderr,"%s: input read: %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (ibfill == 0)
  { goteof = 1;
    return;
  }
 havestamp = 0;
 gen_init();
 x = 0;
 while (1)
  { if (atnl)
     { do_stamp();
       atnl = 0;
     }
    nl = memchr(ibuf+x,'\n',ibfill-x);
    if (nl)
     { gen_data(ibuf+x,(nl+1)-(ibuf+x));
       atnl = 1;
       x = (nl+1) - ibuf;
       if (x >= ibfill)
	{ gen_flush();
	  timing_out = 0;
	  break;
	}
     }
    else
     { gen_writepush(ibuf+x,ibfill-x);
       timing_out = 1;
       break;
     }
  }
}

static void tmo_try_open(void)
{
 if (obfill)
  { gen_init();
    try_flush();
  }
}

static void tmo_timed_out(void)
{
 tmo_try_open();
 timing_out = 0;
}

static void block_signals(void)
{
/*
 sigset_t s;

 sigfillset(&s);
 sigprocmask(SIG_BLOCK,&s,0);
*/
}

int main(int, char **);
int main(int ac, char **av)
{
 handleargs(ac,av);
 stamp_setup();
 if (! logfile)
  { fprintf(stderr,"%s: must specify a logfile\n",__progname);
    usage();
    exit(1);
  }
 ibuf = malloc(bufsize);
 obuf = malloc(bufsize);
 if (!ibuf || !obuf)
  { fprintf(stderr,"%s: can't malloc buffers (bufsize=%d)\n",__progname,bufsize);
    exit(1);
  }
 if ((timeout < 0) && (timeout != -1))
  { fprintf(stderr,"%s: negative timeout value\n",__progname);
    exit(1);
  }
 ibfill = 0;
 obhead = 0;
 obtail = 0;
 obfill = 0;
 atnl = 1;
 goteof = 0;
 timing_out = 0;
 open_failed = 0;
 lost_data = 0;
 iov_max = -1;
 block_signals();
 while (1)
  { struct pollfd pfd;
    int npfd;
    int prv;
    int tmo;
    void (*tmofn)(void);
    static void set_tmo(int ms, void (*fn)(void))
     { if ((tmo == INFTIM) || (ms < tmo))
	{ tmo = ms;
	  tmofn = fn;
	}
     }
    if (goteof)
     { if (obfill == 0) break;
       if (!timing_out && !open_failed) abort();
       npfd = 0;
     }
    else
     { pfd.fd = 0;
       pfd.events = POLLIN | POLLRDNORM;
       npfd = 1;
     }
    tmo = INFTIM;
    if (open_failed) set_tmo(60000,&tmo_try_open);
    if (timing_out)
     { if (timeout < 0)
	{ timing_out = 0;
	}
       else
	{ set_tmo(timeout,&tmo_timed_out);
	}
     }
    prv = poll(&pfd,npfd,tmo);
    if (prv < 0)
     { switch (errno)
	{ case EINTR:
	     continue;
	     break;
	}
       fprintf(stderr,"%s: poll; %s\n",__progname,strerror(errno));
       exit(1);
     }
    if (prv == 0)
     { (*tmofn)();
       continue;
     }
    if ((npfd > 0) && (pfd.revents & (POLLIN|POLLRDNORM))) try_read();
  }
 exit(0);
}