/*
 * LINKTEST [-m minimum[:maximum]] [-s minsize[:maxsize]] [-f] [-rx] [-tx]
 */

#include <sys/types.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/file.h>
#include <sys/time.h>

#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <errno.h>
#include <time.h>

typedef struct Host {
    struct sockaddr_in h_Sin;
    char	*h_Name;
    int		h_Fd;
    int		h_PipeFds[2];
    int		h_IsDynamic;
    int		h_RxTimeStamp;
    int		h_SimulateLostOutput;
    int		h_RxSeq;
    int		h_TxSeq;
    int		h_RxCount;
    int		h_RxMissed;
    int		h_RemoteRxMissed;
} Host;

typedef struct Head {
    char	b_Type;
    char	b_Flags;
    short	b_Len;
    int		b_Seq;
    int		b_TimeStamp;
    int		b_RxMissed;
} Head;

#define B_SIMRECVLOST	0x01

void linksend(Host *h);
void linkrecv(Host *h, int fd, struct sockaddr_in *skin, Head *data, int l);

char	MyHostName[64];
Host	*HostAry;
int	NumHosts;
int	MaxHosts = 1024;
int	MinTime = 1000;
int	MaxTime = 50000;
int	SimulateLostInput;
int	SimulateLostOutput;
int	MinSize = sizeof(Head);
int	MaxSize = 4096;
int	TimeStamp;
int	FastMode = 0;

int
main(int ac, char **av)
{
    int i;
    int fd = -1;
    int pgmask = getpagesize() - 1;

    if (ac == 1) {
	printf(
	    "linktest [-v] [-m min[:max]] [-s [min[:max]]] [-tx] [-rx] [-f[#]] host...\n"
	    "    -v    print version\n"
	    "    -m    set minimum [and maximum] transmit rate, in milliseconds\n"
	    "    -s    set minimum [and maximum] packet size, in bytes\n"
	    "    -tx   introduce transmit data loss\n"
	    "    -rx   introduce receiver data loss\n"
	    "    -f    (-m ignored) transmit fast with up to 1 packet in transit\n"
	    "    -fN   (-m ignored) transmit fast with up to N packets in transit\n"
	);
	exit(0);
    }

    gethostname(MyHostName, sizeof(MyHostName) - 1);
    TimeStamp = (int)time(NULL);

    HostAry = mmap(
		NULL, 
		(sizeof(Host) * MaxHosts + pgmask) & ~pgmask,
		PROT_READ|PROT_WRITE,
		MAP_ANON|MAP_SHARED,
		-1,
		0
	    );

    for (i = 1; i < ac; ++i) {
	char *ptr = av[i];

	if (*ptr != '-')
	    break;
	ptr += 2;
	switch(ptr[-1]) {
	case 'v':
	    printf("%s version 1.2\n", av[0]);
	    exit(1);
	case 'r':
	    SimulateLostInput = 50;
	    break;
	case 't':
	    SimulateLostOutput = 50;
	    break;
	case 'f':
	    FastMode = (*ptr) ? strtol(ptr, NULL, 0) : 1;
	    break;
	case 'm':
	    ptr = (*ptr) ? ptr : av[++i];
	    MinTime = (int)(strtod(ptr, &ptr) * 1000.0 + 1.0);

	    if (*ptr == ':')
		MaxTime = (int)(strtod(ptr + 1, NULL) * 1000.0 + 1.0);
	    else
		MaxTime = 0;

	    if (MaxTime < MinTime)
		MaxTime = MinTime * 2;
	    break;
	case 's':
	    ptr = (*ptr) ? ptr : av[++i];
	    MinSize = strtol(ptr, &ptr, 0);
	    if (MinSize < sizeof(Head))
		MinSize = sizeof(Head);
	    if (MinSize > 4096)
		MinSize = 4096;
	    if (*ptr == ':')
		MaxSize = strtol(ptr + 1, NULL, 0);
	    else
		MaxSize = MinSize;
	    if (MaxSize < MinSize)
		MaxSize = MinSize;
	    if (MaxSize > 4096)
		MaxSize = 4096;
	    break;
	default:
	    fprintf(stderr, "Illegal option: %s\n", ptr - 2);
	    exit(1);
	}
    }

    if ((fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
	perror("socket");
	exit(1);
    }

    /*
     * Create listen socket
     */

    {
	struct sockaddr_in skin;

	bzero(&skin, sizeof(skin));
	skin.sin_family = AF_INET;
	skin.sin_addr.s_addr = INADDR_ANY;
	skin.sin_port = htons(1580);

	if (bind(fd, (void *)&skin, sizeof(skin)) < 0) {
	    perror("bind");
	    exit(1);
	}
    }

    /*
     * Create sending processes
     */

    while (i < ac) {
	struct hostent *host;
	struct sockaddr_in *skin;
	Host *h;

	h = &HostAry[NumHosts++];
	bzero(h, sizeof(Host));
	h->h_Name = av[i];
	h->h_Fd = fd;
	h->h_SimulateLostOutput = SimulateLostOutput;
	skin = &h->h_Sin;

	bzero(skin, sizeof(*skin));
	skin->sin_family = AF_INET;
	skin->sin_addr.s_addr = INADDR_ANY;
	skin->sin_port = htons(1580);

        if ((host = gethostbyname(h->h_Name)) != NULL) {
            skin->sin_family = host->h_addrtype;
            bcopy(host->h_addr, &skin->sin_addr, host->h_length);
	} else {
	    skin->sin_addr.s_addr = inet_addr(h->h_Name);
	}
	linksend(h);
	++i;
    }

    /*
     * Deal with incoming packets and display statistics
     */

    for (;;) {
	union {
	    char buf[32768];
	    Head head;
	} ubuf;
	struct sockaddr_in skin;
	int sLen = sizeof(skin);
	int l;

	if ((l = recvfrom(fd, &ubuf, sizeof(ubuf), 0, (void *)&skin, &sLen)) > 0) {
	    struct Host *h = NULL;

	    for (i = 0; i < NumHosts; ++i) {
		if (skin.sin_addr.s_addr == HostAry[i].h_Sin.sin_addr.s_addr) {
		    h = &HostAry[i];
		    break;
		}
	    }

	    /*
	     * Add a new host to the test, both receiving and sending
	     */
	    if (h == NULL && NumHosts < MaxHosts) {
		struct hostent *he;

		h = &HostAry[NumHosts];
		bzero(h, sizeof(Host));
		h->h_Sin = skin;
		if ((he = gethostbyaddr((void *)&h->h_Sin.sin_addr, sizeof(h->h_Sin.sin_addr), AF_INET)) != NULL) {
		    h->h_Name = strdup(he->h_name);
		} else {
		    h->h_Name = strdup(inet_ntoa(h->h_Sin.sin_addr));
		}
		h->h_IsDynamic = 1;
		h->h_Fd = fd;
		++NumHosts;
		linksend(h);
	    }
	    linkrecv(h, fd, &skin, &ubuf.head, l);
	} else {
	    perror("recvfrom");
	}
    }
    return(0);
}

#define MSG_PING	0x55
#define MSG_PONG	0x33

void
linksend(Host *h)
{
    union {
	char buf[32768];
	Head head;
    } ubuf;
    Head *data = &ubuf.head;
    fd_set RFds;

    if (pipe(h->h_PipeFds) < 0) {
	perror("pipe");
	exit(1);
    }
    fcntl(h->h_PipeFds[0], F_SETFL, O_NONBLOCK);
    fcntl(h->h_PipeFds[1], F_SETFL, O_NONBLOCK);

    if (fork() != 0) {
	close(h->h_PipeFds[0]);
	return;
    }
    close(h->h_PipeFds[1]);
    FD_ZERO(&RFds);

    bzero(&ubuf, sizeof(ubuf));

    for (;;) {
	int len;
	char pbuf[32];

	if (MinSize == MaxSize)
	    len = MinSize;
	else
	    len = (random() % (MaxSize - MinSize)) + MinSize;

	bzero(data, sizeof(Head));
	data->b_Type = MSG_PING;
	data->b_Len = htons(len);
	data->b_Seq = htonl(h->h_TxSeq);
	data->b_TimeStamp = htonl(TimeStamp);
	data->b_RxMissed = htonl(h->h_RxMissed);
	if (SimulateLostInput)
	    data->b_Flags |= B_SIMRECVLOST;
	if (h->h_TxSeq < 16 || (h->h_TxSeq & 15) == 0) {
	    int i;
	    unsigned int v = 0;

	    for (i = sizeof(Head); i < len; ++i) {
		if (v == 0)
		    v = random();
		((char *)(data + 1))[i] ^= v;
		v >>= 8;
	    }
	}
	if (h->h_SimulateLostOutput == 0 || 
	    random() % h->h_SimulateLostOutput != 1
	) {
	    int r;

	    r = sendto(h->h_Fd, data, len, 0, (void *)&h->h_Sin, sizeof(h->h_Sin));
	    if (r < 0)
		perror("sendto");
	}
	++h->h_TxSeq;
	if (FastMode) {
	    while (h->h_TxSeq - h->h_RxSeq > FastMode) {
		struct timeval tv = { 0, 100000 };

		FD_SET(h->h_PipeFds[0], &RFds);
		if (select(h->h_PipeFds[0] + 1, &RFds, NULL, NULL, &tv) == 0)
		    break;
		read(h->h_PipeFds[0], pbuf, sizeof(pbuf));
	    }
	} else {
	    usleep(random() % (MaxTime - MinTime) + MinTime);
	}
	read(h->h_PipeFds[0], pbuf, sizeof(pbuf));	/* drain the pipe */
    }
}

void
linkrecv(Host *h, int fd, struct sockaddr_in *skin, Head *data, int l)
{
    if (l < sizeof(Head))
	return;

    data->b_Len = ntohs(data->b_Len);
    data->b_Seq = ntohl(data->b_Seq);
    data->b_RxMissed = ntohl(data->b_RxMissed);
    data->b_TimeStamp = ntohl(data->b_TimeStamp);

    {
	char c = 0;
	if (write(h->h_PipeFds[1], &c, 1) < 0)
	    perror("pipe write");
    }

    if (data->b_Len != l)
	return;

    switch(data->b_Type) {
    case MSG_PING:
	if (h->h_RxTimeStamp != data->b_TimeStamp) {
	    h->h_RxTimeStamp = data->b_TimeStamp;
	    h->h_RxSeq = data->b_Seq;
	    h->h_RxMissed = 0;
	    h->h_RxCount = 0;
	    if (h->h_IsDynamic)
		h->h_SimulateLostOutput = 0;
	    else
		h->h_SimulateLostOutput = SimulateLostOutput;
	    printf("Receiving from %s (%s)\n", h->h_Name, inet_ntoa(h->h_Sin.sin_addr));
	}

	if (data->b_Flags & B_SIMRECVLOST) {
	    if (h->h_SimulateLostOutput == 0)
		h->h_SimulateLostOutput = 50;
	}
	h->h_RxCount += data->b_Seq - h->h_RxSeq + 1;
	if (data->b_Seq != h->h_RxSeq) {
	    ++h->h_RxMissed;
	    printf("%s(%s)->%s\tlost %d/%d\n",
		h->h_Name,
		inet_ntoa(h->h_Sin.sin_addr),
		MyHostName,
		h->h_RxMissed,
		h->h_RxCount
	    );
	}
	if (h->h_RemoteRxMissed != data->b_RxMissed) {
	    h->h_RemoteRxMissed = data->b_RxMissed;
	    printf("%s->%s(%s)\tlost %d/%d\n",
		MyHostName,
		h->h_Name,
		inet_ntoa(h->h_Sin.sin_addr),
		h->h_RemoteRxMissed,
		h->h_TxSeq
	    );
	}
	h->h_RxSeq = data->b_Seq + 1;
	break;
    }
}