/*
 * blt - like dd, but overlaps reads and writes for speed.
 *
 * -help
 *	Print a help message, including the various options.
 *
 * -v
 *	Print block and byte counts to stderr on exit.
 *
 * -fill
 *	Treat the input as a byte-stream rather than a block-stream,
 *	doing multiple reads to fill up blocks
 *
 * -nbuf N
 *	Use N buffers.  N must be from 1 to 255.
 *
 * -bufsize N
 *	Use N-byte buffers.  N must fit in an int.
 *
 * -max N
 *	Copy no more than N bytes.  N can be large.
 *
 * -recover
 *	When getting a read error (as opposed to EOF), the reader is to
 *	break the read down into smaller pieces, trying to recover as
 *	much data as it can.  See also -grain.  Data that can't be read
 *	gets replaced with 0x00 bytes.
 *
 *	-recover works only when the object being read from supports
 *	seeking.
 *
 * -grain N
 *	When -recover is in use, the block size assumed for the
 *	underlying device is N bytes.  This is applied relative to the
 *	beginning of each buffer, so the grain value normally must
 *	divide the -bufsize value; blt will warn if this is not so.
 *	(If -recover is not in use -grain is accepted but the value is
 *	ignored.)  The default is 512, and the value must fit in an int.
 *
 * For -bufsize, -max, and -grain, the argument can have a
 *  (case-independent) suffix b, k, m, g, or t, to indicate
 *  multiplication by 512, 1024, 1048576, 1073741824, or 1099511627776.
 *
 * Data is always copied from stdin to stdout.
 *
 * SIGINFO produces byte and block counts (the same format as -v),
 *  written to /dev/tty.  If /dev/tty cannot be opened, the output is
 *  still generated, only it's sent to /dev/null.  If /dev/null cannot
 *  be opened either, the write() is still done, but to fd -1 (this may
 *  be visible in syscall traces, though if /dev/null is missing enough
 *  breaks that this is of questionable value).
 *
 * blt achieves its overlapping by forking.  The processes share some
 *  memory, which is created by mmap()ping with MAP_SHARED before
 *  forking.  The data is read by the reader into this memory and
 *  written by the writer out of it; ownership of these blocks is
 *  passed back and forth by writing small tokens to a socket
 *  connection between the two processes.
 *
 * In the reader->writer direction, the tokens are single-character key
 *  characters followed by zero or more additional characters:
 *	s	SIGINFO was noticed.
 *	e	All done; exit.  (This is a normal exit.  Certain
 *		abnormal conditions will produce an abrupt EOF on the
 *		connection; EOF also means exit.  The difference is
 *		that using e allows any remaining writes to drain.)
 *	fX	Full buffer.  X is a single byte holding the buffer
 *		number, from 0 through nbuf-1.
 *	pXN	Partial buffer.  X is as for f, and is followed by N,
 *		which is sizeof(int) bytes in native byte order,
 *		holding the number of bytes in the buffer.
 *
 * In the writer->reader direction, the tokens are single bytes which
 *  hold buffer numbers.  There is no message-type character because
 *  there is only one type of message.
 *
 * blt does assume the socket connection has enough buffering to hold
 *  nbuf bytes.  If this is not true, the two processes will probably
 *  deadlock, each one blocked trying to write more to the connection,
 *  data which cannot be written until the other reads something.
 *
 * This file is in the public domain.
 */

#include <stdio.h>
#include <errno.h>
#include <ctype.h>
#include <fcntl.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <stdarg.h>
#include <sys/uio.h>
#include <strings.h>
#include <sys/socket.h>

extern const char *__progname;

/*
 * Internal operation variables.
 *
 * got_siginfo is set on receipt of SIGINFO; output is generated by the
 *  main line, when it sees got_siginfo set.
 *
 * mem is the block of memory shared by the reader and writer.  It is
 *  conceptually divided into blocks of bufsize bytes each, but is
 *  mapped as one big chunk.
 *
 * rfd and wfd are the reader and writer ends of the socket connection
 *  between the processes.
 *
 * block_full, blocks_partial, and bytes are counters in which we keep
 *  the statistics used by -v and SIGINFO.
 *
 * ttyfd is a file descriptor on /dev/tty (or /dev/null if /dev/tty
 *  can't be opened).  This isn't actually opened until the first time
 *  we have occasion to want it.
 */
static volatile int got_siginfo;
static char *mem;
static int rfd;
static int wfd;
static unsigned long int blocks_full;
static unsigned long int blocks_partial;
static unsigned long long int bytes;
static int ttyfd = -1;

/*
 * Variables holding command-line stuff.
 */
static int verbose = 0;				/* was -v given? */
static int nbuf = 10;				/* -nbuf value */
static int bufsize = 65536;			/* -bufsize value */
static unsigned long long int max = ~0ULL;	/* -max value */
static int fill = 0;				/* was -fill given? */
static int recover = 0;				/* was -recover given? */
static int grain = 512;				/* -grain value */

/*
 * We need a way to map memory.  Unfortunately this is OS-dependent.
 *  Collected here are the variants for the OSes we've got code for.
 *  The interface is:
 *	void *sysdep_map(unsigned int nb)
 *  Allocates nb bytes of memory.  On success, returns the memory.  On
 *  failure, prints a complaint and exits with exit code 1.  The memory
 *  must be such that it's shared after a fork().
 */

#define NEED_SYSDEP_MAP

/* NetBSD version: use mmap with MAP_ANON|MAP_SHARED. */
#if defined(NEED_SYSDEP_MAP) && defined(__NetBSD__)
#include <sys/mman.h>
#ifndef MAP_FAILED
#define MAP_FAILED ((void *)(-1))
#endif
static void *sysdep_map(unsigned int nb)
{
 void *rv;

 rv = mmap(0,nb,PROT_READ|PROT_WRITE,MAP_ANON|MAP_SHARED,-1,0);
 if (rv == MAP_FAILED)
  { fprintf(stderr,"%s: mmap: %s\n",__progname,strerror(errno));
    exit(1);
  }
 return(rv);
}
#undef NEED_SYSDEP_MAP
#endif

/* Linux version: use mmap with MAP_ANON|MAP_SHARED. */
#if defined(NEED_SYSDEP_MAP) && defined(__linux__)
#include <sys/mman.h>
#ifndef MAP_FAILED
#define MAP_FAILED ((void *)(-1))
#endif
static void *sysdep_map(unsigned int nb)
{
 void *rv;

 rv = mmap(0,nb,PROT_READ|PROT_WRITE,MAP_ANON|MAP_SHARED,-1,0);
 if (rv == MAP_FAILED)
  { fprintf(stderr,"%s: mmap: %s\n",__progname,strerror(errno));
    exit(1);
  }
 return(rv);
}
#undef NEED_SYSDEP_MAP
#endif

/* SunOS version: use mmap on /dev/zero, with MAP_SHARED. */
#if defined(NEED_SYSDEP_MAP) && defined(sun)
#include <sys/mman.h>
static void *sysdep_map(unsigned int nb)
{
 int fd;
 void *rv;

 fd = open("/dev/zero",O_RDWR,0);
 if (fd < 0)
  { fprintf(stderr,"%s: /dev/zero: %s\n",__progname,strerror(errno));
    exit(1);
  }
 rv = mmap(0,nb,PROT_READ|PROT_WRITE,MAP_SHARED,fd,0);
 if (rv == (void *)-1)
  { fprintf(stderr,"%s: mmap: %s\n",__progname,strerror(errno));
    exit(1);
  }
 return(rv);
}
#undef NEED_SYSDEP_MAP
#endif

/* If this is still defined, we don't have a sysdep_map(). */
#ifdef NEED_SYSDEP_MAP
#error You need to write a sysdep_map() for this system.
#endif

/*
 * Open /dev/tty if it's not already open.  If /dev/tty isn't openable,
 *  open /dev/null instead.  (If that's not openable, ttyfd is left set
 *  to -1, which makes the write() fail and we'll get called again if
 *  there is a next time.)
 */
static void open_tty(void)
{
 if (ttyfd >= 0) return;
 ttyfd = open("/dev/tty",O_WRONLY,0);
 if (ttyfd < 0) ttyfd = open("/dev/null",O_WRONLY,0);
}

/*
 * Map the memory area to be shared between the processes.
 */
static void setup_map(void)
{
 mem = sysdep_map(nbuf*bufsize);
}

/*
 * Create the socketpair that will connect the processes.
 */
static void setup_pipe(void)
{
 int fd[2];

 if (socketpair(AF_LOCAL,SOCK_STREAM,0,&fd[0]) < 0)
  { fprintf(stderr,"%s: socketpair: %s\n",__progname,strerror(errno));
    exit(1);
  }
 rfd = fd[0];
 wfd = fd[1];
}

/*
 * Read one byte from the inter-process connection.  This routine
 *  `cannot fail'; errors produce complaints and exit(1) rather than
 *  failure returns.
 */
static void read1(int fd, unsigned char *cp, const char *otherend)
{
 int n;

 n = read(fd,cp,1);
 if (n < 0)
  { fprintf(stderr,"%s: read error from %s: %s\n",__progname,otherend,strerror(errno));
    exit(1);
  }
 if (n == 0)
  { fprintf(stderr,"%s: EOF from %s\n",__progname,otherend);
    exit(1);
  }
 if (n != 1) abort();
}

/*
 * Signal handler for SIGINFO.  Just set got_siginfo and let the main
 *  line notice it.
 */
static void siginfo(int sig __attribute__((__unused__)))
{
 got_siginfo = 1;
}

/*
 * Print out block and byte counts, to the specified fd.  This is used
 *  both for -v and for SIGINFO output.
 */
static void print_status(int to)
{
 char *msg;
 int l;

 l = asprintf(&msg,"blocks: full %ld, partial %ld; bytes: %qd\n",blocks_full,blocks_partial,bytes);
 if (msg)
  { write(to,msg,l);
    free(msg);
  }
}

/*
 * The reader process.  The main loop is:
 *	Check for SIGINFO
 *	Get a buffer to fill from the writer
 *	Read stuff into it (read repeatedly if -fill)
 *		Handle errors (see -recover)
 *	Write an appropriate token into the pipe to the writer
 */
static void reader(void)
{
 unsigned char b;
 int n;
 off_t beginning;

 void set_beginning(void)
  { if (recover)
     { beginning = lseek(0,0,SEEK_CUR);
       if (beginning == -1)
	{ fprintf(stderr,"%s: input seek error: %s\n",
		     __progname,strerror(errno));
	  fprintf(stderr,"%s: disabling -recover\n",__progname);
	  recover = 0;
	}
     }
  }

 auto void eofactions(int) __attribute__((__noreturn__));
 void eofactions(int es)
  { write(rfd,"e",1);
    while (read(rfd,&b,1) == 1) ;
    if (verbose) print_status(2);
    exit(es);
  }

 void fullblock(unsigned char bno)
  { static char key = 'f';
    struct iovec iov[2];
    iov[0].iov_base = (void *) &key;
    iov[0].iov_len = 1;
    iov[1].iov_base = (void *) &bno;
    iov[1].iov_len = 1;
    writev(rfd,&iov[0],2);
    blocks_full ++;
  }

 void partialblock(unsigned char bno, int n)
  { static char key = 'p';
    struct iovec iov[3];
    iov[0].iov_base = (void *) &key;
    iov[0].iov_len = 1;
    iov[1].iov_base = (void *) &bno;
    iov[1].iov_len = 1;
    iov[2].iov_base = (void *) &n;
    iov[2].iov_len = sizeof(n);
    writev(rfd,&iov[0],3);
    blocks_partial ++;
  }

 got_siginfo = 0;
 blocks_full = 0;
 blocks_partial = 0;
 bytes = 0;
#ifdef SIGINFO
 signal(SIGINFO,siginfo);
#else
 (void)&siginfo;
#endif
 close(wfd);
 while (1)
  { if (got_siginfo)
     { got_siginfo = 0;
       open_tty();
       print_status(ttyfd);
       write(rfd,"s",1);
     }
    n = read(rfd,&b,1);
    if (n < 0)
     { fprintf(stderr,"%s: read error from writer: %s\n",__progname,strerror(errno));
       if (verbose) print_status(2);
       exit(1);
     }
    if (n == 0)
     { fprintf(stderr,"%s: EOF from writer\n",__progname);
       if (verbose) print_status(2);
       exit(1);
     }
    if (n != 1) abort();
    if (b >= nbuf)
     { fprintf(stderr,"%s: protocol error from writer (buf=%d)\n",__progname,(int)b);
       exit(1);
     }
    set_beginning();
    if (bytes >= max)
     { n = 0;
     }
    else if (fill)
     { int l;
       char *bp;
       int r;
       l = (bufsize > max-bytes) ? max-bytes : bufsize;
       bp = mem + (bufsize * b);
       n = 0;
       while (l > 0)
	{ r = read(0,bp,l);
	  if (r < 0)
	   { n = -1;
	     break;
	   }
	  if (r == 0) break;
	  l -= r;
	  bp += r;
	  n += r;
	}
     }
    else
     { n = read(0,mem+(bufsize*b),(bufsize>max-bytes)?max-bytes:bufsize);
     }
    if ((n < 0) && recover)
     { int i;
       int r;
       char *base;
       int size;
       void seekto(off_t loc)
	{ if (lseek(0,beginning+loc,SEEK_SET) == -1)
	   { fprintf(stderr,"%s: recovery seek: %s\n",__progname,strerror(errno));
	     exit(1);
	   }
	}
       void readsome(int baseoff, int len)
	{ char *bp;
	  int left;
	  left = len;
	  bp = base + baseoff;
	  while (left > 0)
	   { r = read(0,bp,left);
	     if (r < 0)
	      { fprintf(stderr,"%s: read %u @ %llu: %s\n",
		   __progname ,grain,
		   (unsigned long long int)(beginning+(i*grain)),
		   strerror(errno) );
		bzero(base+baseoff,len);
		return;
	      }
	     else if (r == 0)
	      { if (i > 0) partialblock(b,i*grain);
		eofactions(0);
	      }
	     else
	      { bp += r;
		left -= r;
	      }
	   }
	}
       size = (bufsize > max-bytes) ? max-bytes : bufsize;
       n = size / grain;
       base = mem + (bufsize * b);
       for (i=0;i<n;i++)
	{ seekto(i*grain);
	  readsome(i*grain,grain);
	}
       if (n*grain < size)
	{ seekto(n*grain);
	  readsome(n*grain,size-(n*grain));
	}
       seekto(size);
       n = size;
     }
    if (n <= 0)
     { if (n < 0) fprintf(stderr,"%s: read error: %s\n",__progname,strerror(errno));
       eofactions(n<0);
     }
    if (n == bufsize)
     { fullblock(b);
     }
    else
     { partialblock(b,n);
     }
    bytes += n;
  }
}

/*
 * The writer process.  First we preload the connection with nbuf
 *  bytes, which amounts to passing all the buffers to the raeder
 *  process.  Then we drop into the writer main loop:
 *	Read a message type byte from the reader
 *	Handle the message (which usually means writing a buffer's
 *	  data)
 */
static void writer(void)
{
 char k;
 unsigned char b;
 int n;
 int i;

 blocks_full = 0;
 blocks_partial = 0;
 bytes = 0;
 close(rfd);
 for (b=0;b<nbuf;b++)
  { write(wfd,&b,1);
  }
 while (1)
  { read1(wfd,&k,"reader");
    switch (k)
     { case 'e':
	  exit(0);
	  break;
       case 'f':
	  read1(wfd,&b,"reader");
	  if (b >= nbuf)
	   { fprintf(stderr,"%s: protocol error from reader (buf=%d)\n",__progname,(int)b);
	     exit(1);
	   }
	  n = bufsize;
	  blocks_full ++;
	  break;
       case 'p':
	  read1(wfd,&b,"reader");
	  if (b >= nbuf)
	   { fprintf(stderr,"%s: protocol error from reader (buf=%d)\n",__progname,(int)b);
	     exit(1);
	   }
	  for (i=0;i<sizeof(n);i++) read1(wfd,i+(char *)&n,"reader");
	  if ((n < 0) || (n > bufsize))
	   { fprintf(stderr,"%s: protocol error from reader (len=%d)\n",__progname,n);
	     exit(1);
	   }
	  blocks_partial ++;
	  break;
       case 's':
	  /* no point - numbers match what reader printed */
	  continue;
	  break;
       default:
	  fprintf(stderr,"%s: protocol error from reader (key=0x%02x)\n",__progname,(int)(unsigned char)k);
	  exit(1);
	  break;
     }
    bytes += n;
    i = write(1,mem+(bufsize*b),n);
    if (i < 0)
     { fprintf(stderr,"%s: write error: %s\n",__progname,strerror(errno));
       exit(1);
     }
    if (i != n)
     { char *bp;
       bp = mem + (bufsize * b);
       fprintf(stderr,"%s: short write: wanted %d, did %d (retrying remainder)\n",__progname,n,i);
       while (n > 0)
	{ bp += i;
	  n -= i;
	  i = write(1,bp,n);
	  if (i < 0)
	   { fprintf(stderr,"%s: write error: %s\n",__progname,strerror(errno));
	     exit(1);
	   }
	}
     }
    write(wfd,&b,1);
  }
}

/*
 * Turn a number in string form (from the command line) into an int.
 *  This works only for comparatively small numbers; it uses strtol()
 *  but then blindly converts to int, with no checks to see if anything
 *  of importance was lost in the conversion.
 */
static int number(const char *s)
{
 char *e;
 int v;

 v = strtol(s,&e,0);
 if (*e || (e == s))
  { fprintf(stderr,"%s: %s: invalid number\n",__progname,s);
    exit(1);
  }
 return(v);
}

/*
 * Turn a number in string form (from the command line) into an
 *  unsigned long long int.  This also supports the b/k/m/g suffixes,
 *  to scale the values.
 */
static unsigned long long int sufnumber(const char *s)
{
 char *e;
 unsigned long long int v;

 v = strtouq(s,&e,0);
 if (e != s)
  { switch (*e++)
     { case 'b': case 'B':
	  v <<= 9;
	  break;
       case 'k': case 'K':
	  v <<= 10;
	  break;
       case 'm': case 'M':
	  v <<= 20;
	  break;
       case 'g': case 'G':
	  v <<= 30;
	  break;
       case 't': case 'T':
	  v <<= 40;
	  break;
       default:
	  e --;
	  break;
     }
  }
 if (*e || (e == s))
  { fprintf(stderr,"%s: %s: invalid number\n",__progname,s);
    exit(1);
  }
 return(v);
}

/*
 * Print the internal help.
 */
static void printhelp(void)
{
 fprintf(stderr,"Usage: %s [options]\n",__progname);
 fprintf(stderr,"Data is read from stdin and written to stdout.\n");
 fprintf(stderr,"Options:\n");
 fprintf(stderr,"\t-help       Print this help message\n");
 fprintf(stderr,"\t-nbuf N     Use N buffers (default 10)\n");
 fprintf(stderr,"\t-bufsize N  Make buffers N bytes long (default 65536)\n");
 fprintf(stderr,"\t-max N      Don't copy more than N bytes (default no limit)\n");
 fprintf(stderr,"\t-v          Print block/byte counts on exit\n");
 fprintf(stderr,"\t-fill       Do multiple reads to fill buffers\n");
 fprintf(stderr,"-bufsize and -max: argument can have a suffix letter\n");
 fprintf(stderr,"\tSuffix  Argument multiplied by\n");
 fprintf(stderr,"\t  b     512 (disk blocks)\n");
 fprintf(stderr,"\t  k     1024 (kilobytes)\n");
 fprintf(stderr,"\t  m     1048576 (megabytes)\n");
 fprintf(stderr,"\t  g     1073741824 (gigabytes)\n");
}

/*
 * Parse the command line.  Fairly straightforward.
 */
static void handleargs(int ac, char **av)
{
 int errs;
 int skip;

 skip = 0;
 errs = 0;
 for (ac--,av++;ac;ac--,av++)
  { if (skip > 0)
     { skip --;
       continue;
     }
    if (**av != '-')
     { fprintf(stderr,"%s: unrecognized 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,"-help"))
     { printhelp();
       exit(0);
     }
    if (!strcmp(*av,"-v"))
     { verbose = 1;
       continue;
     }
    if (!strcmp(*av,"-fill"))
     { fill = 1;
       continue;
     }
    if (!strcmp(*av,"-nbuf"))
     { WANTARG();
       nbuf = number(av[skip]);
       continue;
     }
    if (!strcmp(*av,"-bufsize"))
     { WANTARG();
       bufsize = sufnumber(av[skip]);
       continue;
     }
    if (!strcmp(*av,"-max"))
     { WANTARG();
       max = sufnumber(av[skip]);
       continue;
     }
    if (!strcmp(*av,"-recover"))
     { recover = 1;
       continue;
     }
    if (!strcmp(*av,"-grain"))
     { WANTARG();
       grain = sufnumber(av[skip]);
       continue;
     }
#undef WANTARG
    fprintf(stderr,"%s: unrecognized option `%s' (use -help for a list)\n",__progname,*av);
    errs ++;
  }
 if (errs) exit(1);
}

/*
 * Check that argument values are sane.  Specifically:
 *
 * nbuf must be 1..255.  The 255 limit is not casually raiseable;
 *  buffer indexes are passed around in unsigned chars (through the
 *  connection); it could be raised to 256 at the price of complicating
 *  a little code that assumes not just nbuf-1 but also nbuf fits in an
 *  unsigned char, but that hardly seems worth it.
 *
 * bufsize must be from 1 to 100M.
 *
 * nbuf*bufsize must not overflow a signed int.
 *
 * If -recover was given, grain must be >=1 and must divide bufsize.
 */
static void check_args(void)
{
 if ((nbuf < 1) || (nbuf > 255))
  { fprintf(stderr,"%s: number of buffers must be 1..255\n",__progname);
    exit(1);
  }
 if ((bufsize < 1) || (bufsize > 104857600))
  { fprintf(stderr,"%s: buffer size must be 1..100M\n",__progname);
    exit(1);
  }
 if ((nbuf*bufsize)/nbuf != bufsize)
  { fprintf(stderr,"%s: buffer-size * buffer-count overflows\n",__progname);
    exit(1);
  }
 if (recover)
  { if (grain < 1)
     { fprintf(stderr,"%s: bad grian size %d\n",__progname,grain);
       exit(1);
     }
    else if (bufsize % grain)
     { fprintf(stderr,"%s: warning: buffer-size %% grain-size is nonzero\n",__progname);
     }
  }
}

/*
 * Fork into two processes.  Return true in the parent, false in the
 *  child.  On error, complain and die.
 */
static int dofork(void)
{
 switch (fork())
  { case -1:
       fprintf(stderr,"%s: fork: %s\n",__progname,strerror(errno));
       exit(1);
       break;
    case 0:
       return(0);
       break;
    default:
       return(1);
       break;
  }
}

/*
 * main() is very simple by this point.  The only thing of note here is
 *  that we ignore SIGPIPE; we'd rather get EPIPE errors from I/O
 *  attempts than be killed by a SIGPIPE.
 */
int main(int, char **);
int main(int ac, char **av)
{
 handleargs(ac,av);
 check_args();
 setup_map();
 setup_pipe();
 signal(SIGPIPE,SIG_IGN);
 if (dofork()) reader(); else writer();
 exit(0);
}