/*
   Primitive TCP speedometer: blast data through the echo port and
   count the bytes read and written in a specified amount of time

   One point of data on a 386/25, Linux 2.0.12:  619 kb/s to 127.0.0.1
*/

#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/signal.h>
#include <sys/time.h>

#define BLKS 65536

extern char *optarg;
extern int optind;

void usage(char *s)
{
    fprintf(stderr, "usage: %s [-t secs] [-u] [-n] addr\n", s);
    exit(1);
}

long rb=0, wb=0;
time_t st;
char rbuf[BLKS], wbuf[BLKS];
int pr=SOCK_STREAM;

void shand(int s)
{
    long t=time(0);
    printf("Sent %ld bytes, received %ld bytes in %ld seconds\n", wb, rb,
	   t-st);
    if (pr==SOCK_DGRAM)
	wb=rb; /* disregard lost datagrams */
    if (t-st)
	printf("%ld bytes/sec\n", (wb+rb)/(t-st));
    exit(0);
}

int main(int argc, char *argv[])
{
    int sec=60;
    int max=BLKS;
    int nb=0;
    int fd, i;
    struct sockaddr_in sa;
    fd_set rs, ws;

    while((i=getopt(argc, argv, "t:un"))!=EOF) {
	switch(i) {
	case 't': sec=atoi(optarg); break;
	case 'u': pr=SOCK_DGRAM; max=512; break;
	case 'n': ++nb; break;
	default: usage(argv[0]);
	}
    }
    if (optind>=argc)
	usage(argv[0]);
    if ((fd=socket(AF_INET, pr, 0))<0) {
	perror("socket"); exit(1);
    }
    sa.sin_family=AF_INET;
    sa.sin_port=htons(7); /* echo */
    sa.sin_addr.s_addr=inet_addr(argv[optind]);
    if (connect(fd, (struct sockaddr *)&sa, sizeof(sa))<0) {
	perror("connect"); exit(1);
    }
    if (nb) {
	if ((i=fcntl(fd, F_GETFL))<0)
	    perror("F_GETFL");
	else if (fcntl(fd, F_SETFL, i|O_NONBLOCK)<0)
	    perror("F_SETFL");
    }

    for (i=0; i<BLKS; ++i)
	wbuf[i]=(i&0x3F)+0x20;
    st=time(0);
    signal(SIGINT, shand);
    signal(SIGTERM, shand);
    signal(SIGALRM, shand);
    signal(SIGHUP, shand);
    alarm(sec);

    while(1) {
	FD_ZERO(&rs); FD_ZERO(&ws);
	FD_SET(fd, &rs); FD_SET(fd, &ws);
	if (select(fd+1, &rs, &ws, NULL, NULL)>0) {
	    if (FD_ISSET(fd, &rs))
		rb+=recv(fd, rbuf, max, 0);
	    if (FD_ISSET(fd, &ws))
		wb+=send(fd, wbuf, max, 0);
	}
    }
    return 0;
}
