/*
 * Copyright 2001, 2002, 2003, 2013 Adam Sampson <ats@offog.org>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <stdlib.h>
#include <time.h>
#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include "iolib.h"
#include "freedt.h"
#include "config.h"
const char *progname = "ratelimit";
const char *proghelp =
	"Usage: ratelimit [OPTIONS] command ...\n"
	"Run a command limiting stdin/stdout bandwidth and/or dumping "
	"IO to stderr.\n\n"
	"-i N      Limit input bandwidth to N bytes per second\n"
	"-o N      Limit output bandwidth to N bytes per second\n"
	"-r        Print all traffic to stderr\n";

/* The size of the data buffers either way. */
#define BUFFERSIZE 64

pid_t pid;

void recordio(const char *data, size_t length, const char *prefix) {
	buffer buf = BUFFER;

	if (length == 0) {
		bformat(&buf, "@i@c[EOF]\n", pid, prefix);
	} else {
		do {
			char c;
	
			bformat(&buf, "@i@c", pid, prefix);
			while (length > 0) {
				c = *data++;
				length--;
				if (c == '\n') break;
				bappendc(&buf, c);
			}
			if (c == '\n')
				bappends(&buf, " \n");
			else
				bappends(&buf, "+\n");
		} while (length > 0);
	}
	writeba(fd_err, &buf);
	bfree(&buf);
}

int main(int argc, char **argv) {
	int inlimit = 0, outlimit = 0, record = 0;
	size_t incount, outcount;
	int inopen = 1, rc;
	int inpipe[2], outpipe[2];
	char inbuf[BUFFERSIZE], outbuf[BUFFERSIZE];
	ssize_t inbsize = 0, outbsize = 0;
	time_t instart = -1, outstart = -1;

	pid = getpid();

	while (1) {
		int c = getopt(argc, argv, "+V?i:o:r");
		if (c == -1)
			break;
		switch (c) {
		case 'i':
			inlimit = atoi(optarg);
			break;
		case 'o':
			outlimit = atoi(optarg);
			break;
		case 'r':
			record = 1;
			break;
		case 'V':
			version();
		default:
			help();
		}
	}

	if ((argc - optind) < 1)
		help();

	if (pipe(inpipe) < 0 || pipe(outpipe) < 0)
		die("unable to create pipe");
	set_fd_cloexec(inpipe[0]);
	set_fd_cloexec(inpipe[1]);
	set_fd_cloexec(outpipe[0]);
	set_fd_cloexec(outpipe[1]);

	rc = fork();
	if (rc < 0)
		die("fork failed");
	if (rc > 0) {
		dup2(inpipe[0], 0);
		dup2(outpipe[1], 1);

		execvp(argv[optind], &argv[optind]);
		die2(argv[optind], "unable to exec");
	}
	close(inpipe[0]);
	close(outpipe[1]);

	incount = outcount = 0;
	while (1) {
		int rc, maxfd = -1;
		fd_set infds, outfds;
		struct timeval tv;
		time_t now = time(NULL);

		tv.tv_sec = 1;
		tv.tv_usec = 0;

#define ADD(fd, set) \
	do { \
		FD_SET(fd, set); \
		if (fd > maxfd) \
			maxfd = fd; \
	} while (0)

		FD_ZERO(&infds);
		if (inopen && inbsize == 0)
			ADD(fd_in, &infds);
		if (outbsize == 0)
			ADD(outpipe[0], &infds);

/* This is somewhat messy, but it's less messy than inlining it twice.
   A write is allowed if some data is available to be written,
   and some time has elapsed so that the rate can be measured,
   and the average rate so far doesn't exceed the limit. */
#define CAN_WRITE(bsize, limit, count, start) \
	((bsize > 0) && \
		((limit == 0) \
		 || ((now != start) && (count / (now - start) <= limit))))

		FD_ZERO(&outfds);
		if (inopen && CAN_WRITE(inbsize, inlimit, incount, instart))
			ADD(inpipe[1], &outfds);
		if (CAN_WRITE(outbsize, outlimit, outcount, outstart))
			ADD(fd_out, &outfds);

		rc = select(maxfd + 1, &infds, &outfds, NULL, &tv);
		if (rc < 0)
			die("select failed");

		if (rc == 0)
			continue;

		if (FD_ISSET(fd_in, &infds)) {
			inbsize = read(fd_in, inbuf, BUFFERSIZE);
			if (inbsize < 0)
				die("read from stdin failed");
			if (record)
				recordio(inbuf, inbsize, " < ");
			if (inbsize == 0) {
				close(inpipe[1]);
				inopen = 0;
			}
			incount += inbsize;
			if (instart == -1)
				instart = time(NULL);
		}

		if (FD_ISSET(outpipe[0], &infds)) {
			outbsize = read(outpipe[0], outbuf, BUFFERSIZE);
			if (outbsize < 0)
				die("read from child failed");
			if (record)
				recordio(outbuf, outbsize, " > ");
			if (outbsize == 0) {
				break;
			}
			outcount += outbsize;
			if (outstart == -1)
				outstart = time(NULL);
		}

		if (FD_ISSET(inpipe[1], &outfds)) {
			ssize_t n = write(inpipe[1], inbuf, inbsize);
			if (n < 0)
				die("write to child failed");
			inbsize -= n;
			memmove(inbuf, inbuf + n, inbsize);
		}

		if (FD_ISSET(fd_out, &outfds)) {
			ssize_t n = write(fd_out, outbuf, outbsize);
			if (n < 0)
				die("write to stdout failed");
			outbsize -= n;
			memmove(outbuf, outbuf + n, outbsize);
		}
	}

	return 0;
}