#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/un.h>
#ifdef __linux__
#include <net/if.h>
#endif
#include <netinet/in.h>
#include <netdb.h>
#include <errno.h>

#define ME "tcperf: "

#define strncpyz(d,s,l) *(strncpy(d,s,l)+l)=0
#define BadPtr(ptr) (!ptr || !*ptr)

#define PS 8188

static struct sockaddr_in bind_addr;

static char isstrdec(char *s,size_t l)
{
while(l--)
  {
  if(*s<48 || *s>57)
    return 0;
  }
return 1;
}

static char isstrhex(char *s,size_t l)
{
while(l--)
  {
  if(*s<48 || (*s>57 && *s<65) || (*s>70 && *s<97) || *s>102)
    return 0;
  }
return 1;
}

static char bad_inaddr(void *in)
{
int i;

for(i=0;i<4;i++)
  {
  if(*((u_char *)in+i)!=255)
    return 0;
  }
return 1;
}

static char any_inaddr(void *in)
{
int i;

for(i=0;i<4;i++)
  {
  if(*((u_char *)in+i)!=0)
    return 0;
  }
return 1;
}

static char lo_inaddr(void *in)
{
if(*(u_char *)in==127)
  return 1;
return 0;
}

static char *inetntoa(void *in)
{
static char buf[16];
u_char *s=(u_char *)in;
int a,b,c,d;

a=(int)*s++;
b=(int)*s++;
c=(int)*s++;
d=(int)*s++;
sprintf(buf,"%d.%d.%d.%d",a,b,c,d);
return buf;
}

static char inetaton(char *s,void *d)
{
char *p;
int i=0;

memset(d,-1,4);
p=s;
while(*p && i<4)
  {
  if(!(p=strchr(s,'.')))
    {
    if(*s=='*')
      {
      for(;i<4;i++)
        *((u_char *)d+i)=0;
      return 1;
      }
    p=strchr(s,0);
    }
  if((p-s)>3 || !isstrdec(s,p-s))
    {
    memset(d,-1,4);
    return 0;
    }
  *((u_char *)d+i)=(u_char)atoi(s);
  i++;
  s=p+1;
  }
if(*p || i!=4)
  {
  memset(d,-1,4);
  return 0;
  }
return 1;
}

static int mk_conn(char *host,int port,int win)
{
struct sockaddr_in addr;
int s;

if(!inetaton(host,&addr.sin_addr))
  {
  struct hostent *he;

  if(!(he=gethostbyname(host)))
    {
    fprintf(stderr,ME "Unable to resolve %s\n",host);
    exit(1);
    }
  if(he->h_length!=4)
    {
    fprintf(stderr,ME "Error resolving %s\n",host);
    exit(1);
    }
  memcpy(&addr.sin_addr,he->h_addr,he->h_length);
  }
addr.sin_family=AF_INET;
addr.sin_port=htons(port);
if((s=socket(AF_INET,SOCK_STREAM,0))<0)
  {
  fprintf(stderr,ME "socket() failed\n");
  exit(1);
  }
if(win && setsockopt(s,SOL_SOCKET,SO_SNDBUF,(void *)&win,sizeof(int))<0)
  {
  close(s);
  fprintf(stderr,ME "setsockopt(SO_SNDBUF) failed\n");
  exit(1);
  }
if(!any_inaddr(&bind_addr.sin_addr))
  {
  if(bind(s,(struct sockaddr *)&bind_addr,sizeof(struct sockaddr_in))<0)
    {
    close(s);
    fprintf(stderr,ME "bind() failed\n");
    exit(1);
    }
  }
if(connect(s,(struct sockaddr *)&addr,sizeof(struct sockaddr_in))<0)
  {
  close(s);
  fprintf(stderr,ME "connect() failed\n");
  exit(1);
  }
return s;
}

static int mk_listen(int port,int win)
{
int opt=1;
int s,s2;

bind_addr.sin_family=AF_INET;
bind_addr.sin_port=htons(port);
if((s=socket(AF_INET,SOCK_STREAM,0))<0)
  {
  fprintf(stderr,ME "socket() failed\n");
  exit(1);
  }
if(setsockopt(s,SOL_SOCKET,SO_REUSEADDR,(void *)&opt,sizeof(int))<0)
  {
  close(s);
  fprintf(stderr,ME "setsockopt(SO_REUSEADDR) failed\n");
  exit(1);
  }
if(win && setsockopt(s,SOL_SOCKET,SO_RCVBUF,(void *)&win,sizeof(int))<0)
  {
  close(s);
  fprintf(stderr,ME "setsockopt(SO_RCVBUF) failed\n");
  exit(1);
  }
if(bind(s,(struct sockaddr *)&bind_addr,sizeof(struct sockaddr_in))<0)
  {
  close(s);
  fprintf(stderr,ME "bind() failed\n");
  exit(1);
  }
if(listen(s,1)<0)
  {
  close(s);
  fprintf(stderr,ME "listen() failed\n");
  exit(1);
  }
if((s2=accept(s,0,0))<0)
  {
  close(s);
  fprintf(stderr,ME "accept() failed\n");
  exit(1);
  }
close(s);
return s2;
}

static void do_test(int s,long ctm)
{
static char buf[8192];
struct timeval tc,to;
unsigned long long tbytes;
unsigned long tms,ms,bytes;
int r;
char srv=(!ctm);

tms=0,tbytes=0;
ms=0,bytes=0;
gettimeofday(&to,0);
printf("Transferred: 0 kB  (0 kbit/s)       \n");
while(1)
  {
  if(ms>=500000)
    {
    ms=ms/1000;
    tbytes+=bytes;
    tms+=ms;
    printf("\033[1ATransferred: %llu kB  (%lu kbit/s)       \n",tbytes/1000,bytes/ms*8);
    ms=0,bytes=0;
    }

  if(srv)
    {
    if((r=recv(s,buf,PS,0))<0)
      {
      close(s);
      fprintf(stderr,ME "recv() failed\n");
      exit(1);
      }
    if(r>0)
      bytes+=r;
    else
      break;
    }
  else
    {
    r=send(s,buf,PS,0);
    if(r<=0)
      {
      close(s);
      fprintf(stderr,ME "send() failed\n");
      exit(1);
      }
    bytes+=r;
    }

  gettimeofday(&tc,0);
  ms+=((tc.tv_sec-to.tv_sec)*1000000)+(tc.tv_usec-to.tv_usec);
  to=tc;
  if(ms>=500000)
    ctm-=ms/1000;
  if(!srv && ctm<=0)
    break;
  }
close(s);
ms=ms/1000;
tbytes+=bytes;
tms+=ms;
if(tms<10)
  tms=10;
printf("\033[1ATransferred: %llu kB  (%lu kbit/s)       \n",tbytes/1000,tbytes/tms*8);
}

#ifdef __linux__
static char getifaddr(char *ifn,void *addr)
{
struct ifreq ifr;
struct sockaddr_in *ifa;
int s;

if((s=socket(AF_INET,SOCK_DGRAM,0))<0)
  return 0;
memset(&ifr,0,sizeof(ifr));
strncpyz(ifr.ifr_name,ifn,IFNAMSIZ-1);
if(ioctl(s,SIOCGIFADDR,&ifr)<0)
  {
  close(s);
  return 0;
  }
close(s);
ifa=(struct sockaddr_in *)&ifr.ifr_addr;
if(ifa->sin_family!=AF_INET || any_inaddr(&ifa->sin_addr))
  return 0;
memcpy(addr,&ifa->sin_addr,4);
return 1;
}
#endif

static void usage(void)
{
printf("Usage: tcperf [-b <host>] [-t <secs>] -c <host> [<port>]\n");
printf("       tcperf [-b <host>] -s [<port>]\n");
exit(0);
}

int main(int argc,char *argv[])
{
static char bind_host[64];
static char host[64];
long ctm=5000;
int port=5001;
int i;
int win=0;
char opt_c=0;
char opt_s=0;

*bind_host=0;
*host=0;
memset(&bind_addr.sin_addr,0,4);

for(i=1;i<argc;i++)
  {
  if(BadPtr(argv[i]))
    {
    fprintf(stderr,ME "Panic: Null argument\n");
    exit(2);
    }
  if(!strcmp(argv[i],"--"))
    {
    i++;
    break;
    }
  if(*argv[i]!='-')
    break;

  if(!strcmp(argv[i],"--help"))
    usage();

  if(!strcmp(argv[i],"-t"))
    {
    i++;
    if(argc<i || BadPtr(argv[i]) || *argv[i]=='-')
      {
      fprintf(stderr,ME "Not enough arguments\n");
      exit(1);
      }
    if(!(ctm=atoi(argv[i])*1000))
      {
      fprintf(stderr,ME "Invalid timer specified\n");
      exit(1);
      }
    continue;
    }

  if(!strcmp(argv[i],"-w"))
    {
    i++;
    if(argc<i || BadPtr(argv[i]) || *argv[i]=='-')
      {
      fprintf(stderr,ME "Not enough arguments\n");
      exit(1);
      }
    if(!(win=atoi(argv[i])*1024))
      {
      fprintf(stderr,ME "Invalid TCP Window specified\n");
      exit(1);
      }
    continue;
    }

  if(!strcmp(argv[i],"-i"))
    {
    #ifdef __linux__
    i++;
    if(argc<i || BadPtr(argv[i]) || *argv[i]=='-')
      {
      fprintf(stderr,ME "Not enough arguments\n");
      exit(1);
      }
    if(!getifaddr(argv[i],&bind_addr.sin_addr))
      {
      fprintf(stderr,ME "Invalid interface specified\n");
      exit(1);
      }
    bind_addr.sin_family=AF_INET;
    bind_addr.sin_port=0;
    continue;
    #else
    fprintf(stderr,ME "Not supported\n");
    exit(1);
    #endif
    }

  if(!strcmp(argv[i],"-b"))
    {
    struct hostent *he;

    i++;
    if(argc<i || BadPtr(argv[i]) || *argv[i]=='-')
      {
      fprintf(stderr,ME "Not enough arguments\n");
      exit(1);
      }
    strncpyz(bind_host,argv[i],63);
    bind_addr.sin_family=AF_INET;
    bind_addr.sin_port=0;
    if(inetaton(bind_host,&bind_addr.sin_addr))
      continue;
    if(!(he=gethostbyname(bind_host)))
      {
      fprintf(stderr,ME "Unable to resolve %s\n",bind_host);
      exit(1);
      }
    if(he->h_length!=4)
      {
      fprintf(stderr,ME "Error resolving %s\n",bind_host);
      exit(1);
      }
    memcpy(&bind_addr.sin_addr,he->h_addr,4);
    continue;
    }

  if(!strcmp(argv[i],"-c"))
    {
    i++;
    opt_c=1;
    if(argc<i || BadPtr(argv[i]) || *argv[i]=='-')
      {
      fprintf(stderr,ME "Not enough arguments\n");
      exit(1);
      }
    strncpyz(host,argv[i],63);
    if(argc>i+1 && !BadPtr(argv[i+1]) && *argv[i+1]!='-')
      {
      i++;
      if(!(port=atoi(argv[i])))
        {
        fprintf(stderr,ME "Invalid port specified\n");
        exit(1);
        }
      }
    continue;
    }

  if(!strcmp(argv[i],"-s"))
    {
    opt_s=1;
    if(argc>i+1 && !BadPtr(argv[i+1]) && *argv[i+1]!='-')
      {
      i++;
      if(!(port=atoi(argv[i])))
        {
        fprintf(stderr,ME "Invalid port specified\n");
        exit(1);
        }
      }
    continue;
    }

  fprintf(stderr,ME "Unknown option\n");
  exit(1);
  }

if(opt_c && opt_s)
  {
  fprintf(stderr,ME "Cannot mix '-c' and '-s' options\n");
  exit(1);
  }

if(opt_s)
  {
  do_test(mk_listen(port,win),0);
  return 0;
  }
if(opt_c)
  {
  do_test(mk_conn(host,port,win),ctm);
  return 0;
  }
usage();
return 0;
}