/* replace.c   by Michael Thorpe   2011-06-30 */
/*
 * This program opens a file and execs a subprogram.  The subprogram is
 * given pipes for stdin and stdout, which this program uses to send the
 * file to the subprogram and get data in return.  This program replaces
 * the data in the file with the data returned from the subprogram, being
 * careful not to overwrite the original data until it has been fed to the
 * subprogram.  If the subprogram exits with an error value or on a signal
 * and no data has yet been returned from it, the file will not be changed.
 * Otherwise, the file is truncated after the subprogram exits if more was
 * read than was written and the -t option was not given.
 *
 * Example usage:  replace largefile grep -v lines_i_dont_care_about
 */

#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#define BUF_SIZE 16384

struct buffer {
	struct buffer *next;
	char *buf;
	size_t start,stop;
};

struct bufhead {
	int fin,fout;
	struct buffer *first,*last;
};

const char usage[]="usage: replace [-fht] <file> <program> [<args>]\n";
const char help[]= "       -f  truncate file even if program exits non-zero\n"
                   "       -h  show this help message\n"
                   "       -t  don't truncate file on subprogram exit\n";

/* This is the pid and exit status of the child. */
pid_t child=0;
int child_status=0;

/* This comes into play if we get SIGCHLD before recording the child pid */
int need_to_waitpid=0;

/* Positions for reading and writing in the file, and the current position. */
size_t readpos=0,writepos=0,curpos=0;

/* These are the buffers between to/from the subprogram. */
struct bufhead inbuf,outbuf;

/* Have we reached end-of-file? */
int file_eof=0;

/* Have we read all that the subprogram has written? */
int subprog_eof=0;

/* This controls whether we truncate the file after a non-zero subprogram exit. */
int force_truncate=0;

/* This controls whether we truncate the file after the subprogram exits. */
int no_truncate=0;

/*
 * This function returns the index of the first non-option argument, or 0 on
 * success-but-don't-run or -1 on error.  It also makes sure there are at
 * least 2 non-option arguments.
 */
int parse_args(int argc,char **argv) {
	int i,j;

	for(i=1;i<argc && argv[i][0]=='-';i++)
		for(j=1;argv[i][j];j++)
			switch(argv[i][j]) {
			case '-':
				return(i+1);
			case 'f': /* force truncate */
				force_truncate=1;
				break;
			case 'h':
				fputs(usage,stdout);
				fputs(help,stdout);
				return(0);
			case 't': /* no truncate */
				no_truncate=1;
				break;
			default:
				fputs(usage,stderr);
				return(-1);
			}
	if(argc-i<2) {
		fputs(usage,stderr);
		return(-1);
	}
	if(force_truncate && no_truncate) {
		fputs("force-truncate and no-truncate together isn't sensible\n",stderr);
		return(-1);
	}
	return(i);
}

/*
 * This function forks, and execs the subprogram in the child.  It also
 * handles setting up the pipes.  It returns 0 on success (in the parent
 * only) or -1 on failure.
 */
int do_fork_exec(char **argv) {
	int in[2];
	int out[2];
	int tmp;

	if(pipe(in)) {
		perror("pipe");
		return(-1);
	}
	tmp=fcntl(in[1],F_GETFL);
	if(-1==fcntl(in[1],F_SETFL,tmp|O_NONBLOCK))
		return(-1);
	if(pipe(out)) {
		perror("pipe");
		return(-1);
	}
	child=fork();
	if(child<0) {
		perror("fork");
		return(-1);
	}
	if(child==0) {
		if(in[0] != 0) {
			dup2(in[0],0);
			close(in[0]);
		}
		close(in[1]);
		close(out[0]);
		if(out[1] != 1) {
			dup2(out[1],1);
			close(out[1]);
		}
		execvp(argv[0],argv);
		perror("exec");
		exit(-1);
	}
	close(in[0]);
	close(out[1]);
	inbuf.fout=in[1];
	outbuf.fin=out[0];
	return(0);
}

/*
 * I... am... the grim... rea-per!
 *
 * But seriously, folks, just nuke child and set child_status if the child
 * just exited.
 */
void reaper(int sig) {
	pid_t pid;
	int status;

	if(sig && !child) {
		need_to_waitpid=1;
		return;
	}
	while(1) {
		pid=waitpid(-1,&status,sig?WNOHANG:0);
		if(pid<=0)
			break;
		if(pid==child && WIFEXITED(status)) {
			child=0;
			child_status=WEXITSTATUS(status);
		}
	}
#ifdef REDO_SIGNALS
	signal(SIGCHLD,reaper);
#endif
}

/*
 * Read a chunk into buffer; !!isfile if readpos should be updated.
 */
int slurp(struct bufhead *bh,int isfile) {
	struct buffer *sb;

/* FIXME: This ought to use any fractional buffer in bh. */
	sb=(struct buffer *)malloc(sizeof(struct buffer));
	if(!sb)
		return(-1);
	sb->buf=(char *)malloc(BUF_SIZE);
	if(!sb->buf)
		return(-1);
	if(isfile && curpos != readpos) {
		if(readpos != lseek(bh->fin,readpos,SEEK_SET))
			return(-1);
		curpos=readpos;
	}
	sb->stop=read(bh->fin,sb->buf,BUF_SIZE);
	if(sb->stop==-1) {
		perror("read");
		return(-1);
	}
	if(sb->stop==0) {
		free(sb->buf);
		free(sb);
		if(isfile)
			file_eof=1;
		else
			subprog_eof=1;
	} else {
		if(isfile) {
			readpos+=sb->stop;
			curpos=readpos;
		}
		sb->next=0;
		sb->start=0;
		if(bh->last)
			bh->last->next=sb;
		bh->last=sb;
		if(!bh->first)
			bh->first=sb;
	}
	return(0);
}

/*
 * Try to write out the buffer. !!isfile if read/writepos should be obeyed.
 */
int barf(struct bufhead *bh,int isfile) {
	size_t l;
	struct buffer *sb;

	if(!bh->first)
		return(0);
	l=bh->first->stop-bh->first->start;
	if(isfile && !file_eof && l>(readpos-writepos))
		l=readpos-writepos;
	if(!l)
		return(0);
	if(isfile && curpos != writepos) {
		if(writepos != lseek(bh->fout,writepos,SEEK_SET))
			return(-1);
		curpos=writepos;
	}
	l=write(bh->fout,bh->first->buf+bh->first->start,l);
	if(l==-1) {
		if(errno==EINTR)
			return(0);
		return(-1);
	}
	bh->first->start+=l;
	if(isfile) {
		writepos+=l;
		curpos=writepos;
	}
	if(bh->first->start==bh->first->stop) {
		free(bh->first->buf);
		sb=bh->first->next;
		free(bh->first);
		bh->first=sb;
		if(!sb)
			bh->last=0;
	}
	return(0);
}

/*
 * This space for rent.
 */
int do_it() {
	fd_set r,w;
	int i,maxfd;

	while(child || !subprog_eof) {
		FD_ZERO(&r);
		FD_SET(outbuf.fin,&r);
		maxfd=outbuf.fin;
		if(!file_eof) {
			FD_ZERO(&w);
			FD_SET(inbuf.fout,&w);
			if(maxfd<inbuf.fout)
				maxfd=inbuf.fout;
		}
		i=select(1+maxfd,&r,(file_eof?0:&w),0,0);
		if(i<0 && errno==EINTR)
			continue;
		if(i<=0)
			return(-1);
		if(FD_ISSET(outbuf.fin,&r)) {
			if(slurp(&outbuf,0))
				return(-1);
		}
		if(!file_eof && FD_ISSET(inbuf.fout,&w)) {
			if(!inbuf.first && slurp(&inbuf,1))
				return(-1);
			if(file_eof && close(inbuf.fout))
				return(-1);
			if(barf(&inbuf,0))
				return(-1);
		}
		if(barf(&outbuf,1))
			return(-1);
	}
	if(file_eof)
		while(outbuf.first)
			if(barf(&outbuf,1))
				return(-1);
	return(0);
}

/*
 * This function... well, I think you know what it does.
 */
int main(int argc,char **argv) {
	int i,file;

	i=parse_args(argc,argv);
	if(i<=0)
		return(i);
	file=open(argv[i],O_RDWR);
	if(file==-1) {
		perror("open");
		return(-1);
	}
	if(-1==fcntl(file,F_SETFD,1)) {
		perror("fcntl");
		return(-1);
	}
	signal(SIGCHLD,reaper);
	signal(SIGPIPE,SIG_IGN); /* Just in case we have unlucky timing. */
	if(do_fork_exec(argv+i+1))
		return(-1);
	if(need_to_waitpid)
		reaper(SIGCHLD);
	inbuf.fin=file;
	inbuf.first=0;
	inbuf.last=0;
	outbuf.fout=file;
	outbuf.first=0;
	outbuf.last=0;
	if(do_it())
		return(-1);
	reaper(0);
	if(!no_truncate && (force_truncate || !child_status) && readpos>writepos)
		ftruncate(file,writepos);
	if(close(file))
		return(-1);
	return(child_status);
}