#include "Test.h"

struct sockaddr_in sinme;/*client socket*/
struct sockaddr_in sinhim;/*server socket*/
struct sockaddr_in frominet;
SOCKET testSocket_T, testSocket_U;/*TCP&UDP*/
SOCKET temp_tcpsocket;
SOCKET fd; /* fd of network socket */

struct hostent *addr;
extern int errno;
char *optarg;
char *Talk;

int buflen = 1024; /* length of buffer */
int delay = 0; /*milliseconds of delay before each write*/
int num = 0; /*total of packets*/
int trans = 1; /* 0 = receive, 1 = transmit mode */
int cmd = 0; /* 0 = other type, !0 = cmd*/
int performance = 0; /* */
int udp = 0; /* 0 = tcp, 1 = udp */
int ping = 0; /* 0 = other type, 1 = ping*/
int options = 0; /* socket options */
int one = 1; /* for 4.3 BSD style setsockopt() */
short port = 20001; /* port number */
char *host; /* ptr to name of host */
int optind = 1;
int offset = 0; /*0 = option+optarg , 1 = option+__optarg*/
int print = 0; /*0 = no print , 1 = print*/
int talk = 0; /*1 = talk*/
int cmdflag = 0; /* 1 = client */

char Usage[] = "\
HaoXu test v1.1\n\
Usage: test [-options] host [CLIENT]\n\
test -r [SERVER])\n\
Common options:\n\
-p# ping other host with ICMP\n\
-u use UDP instead of TCP (default TCP)\n\
-n number of packets written to network (default 1)\n\
-w milliseconds of delay before each write (default 0)\n\
-t talk to other peer\n\
format for rate:\n\
-M length ( MBytes ) of bufs read from or written to network (default 1024 Bytes)\n\
-l length ( Bytes ) of bufs read from or written to network (default 1024 Bytes)\n\
-c# command to the server\n\
format for command:\n\
-c passive send packets from the server\n\
e.g.:\n\
test -M3 X.X.X.X (send TCP 3 MBytes)\n\
test -u -l 1000 -n1000 -w500 X.X.X.X (send UDP)\n\
test -t X.X.X.X (talk to X.X.X.X)\n\
test -c passive -u -l 1000 -n1000 -w500 X.X.X.X (receive UDP from X.X.X.X)\n\
";

static BOOL UdpSend();
static BOOL UdpRecv();
static BOOL server();
static BOOL CreateTcpSock();
static BOOL CreateUdpSock();
static void mes(char* s);
static void SendTime(int byte, DWORD time);
static int getopt(int argc, char *argv[]);
static BOOL SendTcp();
static BOOL RecvTcp();
static BOOL Ping(int argc, char* argv[]);
static void fill_icmp_data(char *, int);
static USHORT checksum(USHORT *, int);
static int decode_resp(char *,int ,struct sockaddr_in *);

int main(int argc, char* argv[])
{
unsigned long addr_tmp;
int c;
WORD wVersionRequested;
WSADATA wsaData;
int error;
BOOL Bool;
char passive[] = "passive";

if (argc < 2) goto usage;

while ((c = getopt(argc, argv)) != -1)
{
switch (c)
{
case 'l':
{
buflen = atoi(optarg);
if(offset)
optind++;
break;
}

case 'u':
{
udp = 1;
break;
}

case 'n':
{
num = atoi(optarg);
if(offset)
optind++;
break;
}

case 'w':
{
delay = atoi(optarg);
if(offset)
optind++;
break;
}

case 'M':
{
buflen = atoi(optarg) * 1024 *1024;
if(offset)
optind++;
break;
}

case 'p':
{
ping = 1;
break;
}

case 'c':
{
if(strcmp(passive, optarg))
{
printf("++++++++++Command parameter error.+++++++++++++++\n");
goto usage;
}
else
{
cmd = 1;
}
if(offset)
optind++;
break;
}

case 'x':
{
performance = atoi(optarg);
if(offset)
optind++;
break;
}

case 't':
{
udp = 1;
talk = 1;
Talk = (char*)malloc(buflen-HEADLEN);
memset(Talk, 0 , buflen-HEADLEN);
break;
}

case 's':
{
print = 1;
break;
}

case 'r':
{
trans = 0;
break;
}

default:
{
goto usage;
break;
}
}
}

wVersionRequested = MAKEWORD(1, 1);

/* Initialize the WINSOCK.DLL */
error = WSAStartup(wVersionRequested, &wsaData);

/* Tell the user that we couldn't find a useable */
/* winsock.dll. */
if(error != 0)
return 0;

/* Confirm that the Windows Sockets DLL supports 1.1.*/
/* Note that if the DLL supports versions greater */
/* than 1.1 in addition to 1.1, it will still return */
/* 1.1 in wVersion since that is the version we */
/* requested. */

if ( LOBYTE( wsaData.wVersion ) != 1 || HIBYTE( wsaData.wVersion ) != 1 )
{
/* Tell the user that we couldn't find a useable */
/* winsock.dll. */
WSACleanup();
return 0;
}
/* The Windows Sockets DLL is acceptable. Proceed. */

memset((struct sockaddr*)&sinme, 0, sizeof(struct sockaddr));
sinme.sin_family = AF_INET;
#if defined(cray)
addr_tmp = inet_addr(INADDR_ANY);
sinme.sin_addr = addr_tmp;
#else
sinme.sin_addr.s_addr = htonl(INADDR_ANY);
#endif
if(trans)
sinme.sin_port = 0;
else
sinme.sin_port = htons(port);

if(trans)
{
if (optind == argc)
goto usage;
memset((char *)&sinhim, 0, sizeof(sinhim));
host = argv[optind];

if (atoi(host) > 0 )
{
/* Numeric */
sinhim.sin_family = AF_INET;
#if defined(cray)
addr_tmp = inet_addr(host);
sinhim.sin_addr = addr_tmp;
#else
sinhim.sin_addr.s_addr = inet_addr(host);
#endif
}
else
{
if ((addr=gethostbyname(host)) == NULL)
mes("addr");
sinhim.sin_family = addr->h_addrtype;
CopyMemory(&addr_tmp,(const char*)addr->h_addr, sizeof(addr_tmp));
#if defined(cray)
sinhim.sin_addr = addr_tmp;
#else
sinhim.sin_addr.s_addr = addr_tmp;
#endif /* cray */
}
sinhim.sin_port = htons(port);
}

/*transmit*/
if(trans)
{
if(ping)
{
Bool = Ping(argc, argv);
if(Bool == FALSE)
{
printf("ping failed.\n");
}
}

if(udp || cmd || talk)
{
Bool = UdpSend();
if(Bool == FALSE)
{
printf("Send UDP packets failed.\n");
}
}

if(!udp && !cmd && !ping && !talk)
{
Bool = SendTcp();
if(Bool == FALSE)
{
printf("Send TCP failed.\n");
}
}
}

/*receive*/
if(!trans)
{
Bool = server();
if(Bool == FALSE)
{
printf("receive error.\n");
}
}

WSACleanup();
exit(0);

usage:
fprintf(stderr,Usage);
return (0);
}

BOOL UdpSend()
{
char* Pkt;
char* temp_ack;
pSign server_ack, pkt;
int result;
fd_set readFds;
struct timeval sockTimeout;
int i , j, maxnum, FromLen, seq;
DWORD currentTime, time;
DWORD Average = 0;
BOOL bool;

if(!cmd)
{
printf("\nUDP Transmit Test:\n");
printf(" Statistics : UDP -> %s:%d\n", inet_ntoa(sinhim.sin_addr), port);
printf(" Buffer Size = %d\n", buflen);
}

if ((fd = socket(AF_INET, SOCK_DGRAM, 0)) == INVALID_SOCKET)
mes("socket");

if (bind(fd, (struct sockaddr *)&sinme, sizeof(sinme)) == INVALID_SOCKET)
mes("bind");

/* We are the client if transmitting */
if(cmd)/*if command, send format of cmd packet */
{
Pkt = (char*)malloc(HEADLEN + sizeof(Cmd));
memset(Pkt, 0, (HEADLEN + sizeof(Cmd)));
}
else/*normal format packets*/
{
Pkt = (char*)malloc(buflen);
memset(Pkt, 0, buflen);
}
pkt = (pSign)Pkt;

pkt->mType = DIAG_MSG_TYPE_REQ;
pkt->Seq = 1;
pkt->Buflen = (buflen - HEADLEN);

if(talk)
{
printf(" Send Mode : Read from STDIN\n");
pkt->cType = DIAG_CMD_TYPE_TALK;
}
else if(cmd && !talk)
{
pkt->cType = DIAG_CMD_TYPE_CMD;
((pCmd)(pkt+1))->Buflen = buflen;
((pCmd)(pkt+1))->Delay = delay;
((pCmd)(pkt+1))->Num = num;

if(udp)
{
((pCmd)(pkt+1))->Protocol = DIAG_CMD_TYPE_UDP;
}
else
{
((pCmd)(pkt+1))->Protocol = DIAG_CMD_TYPE_TCP;
}
}
else
{
pkt->cType = DIAG_CMD_TYPE_UDP;

for(j=HEADLEN,i=65;j<buflen;j++,i++)
{
if(i==91)
{
i = 48;
*(Pkt+j) = i;
}
else if(i==58)
{
i = 97;
*(Pkt+j) = i;
}
else if(i==123)
{
i = 65;
*(Pkt+j) = i;
}
else
*(Pkt+j) = i;
}
}

seq = 0;

if( num == 0 || talk || cmd )/* talk or cmd */
{
maxnum = 1;
}
else if( !talk && !cmd )/*send UDP packets*/
{
maxnum = num;
}

for(i=0; i<maxnum ;i++)
{
pkt->Time = GetTickCount();
pkt->Sendnum = maxnum;

if(talk)
{
while(1)
{
scanf("%s", Talk);
memcpy(((char*)pkt+HEADLEN), Talk, buflen - HEADLEN);
if(sendto(fd, (char*)pkt, buflen, 0, (struct sockaddr*)&sinhim, sizeof(struct sockaddr)) == INVALID_SOCKET)
{
mes("send");
}
}
}

if(sendto(fd, (char*)pkt, cmd?(HEADLEN + sizeof(Cmd)):buflen, 0, (struct sockaddr*)&sinhim, sizeof(struct sockaddr)) == INVALID_SOCKET)
{
mes("send");
}

if(cmd)
{
cmdflag = 1;
sinme.sin_port = htons(port);
free(Pkt);
bool = server();
if(bool == FALSE)
return (FALSE);
}

if(!talk && !cmd)
{
sockTimeout.tv_sec = TIMEOUT;
sockTimeout.tv_usec = 0;

/* Wait for reply at the ephemeral port selected by the sendto () call. */
FD_ZERO (&readFds);
FD_SET (fd, &readFds);

result = select (0, &readFds, NULL, NULL, &sockTimeout);

if (result < 0)
{
return (FALSE);
}

if (result == 0) /* TIMEOUT interval expired. */
{
printf("Request timed out.\n");
}

FromLen = sizeof(sinhim);
temp_ack = malloc(HEADLEN);
if(result > 0)
{
result = recvfrom(fd, (char*)temp_ack, HEADLEN, 0, (struct sockaddr*)&sinhim, &FromLen);
if(result == INVALID_SOCKET)
{
mes("recvfrom");
}

server_ack = (pSign)temp_ack;
currentTime = GetTickCount();

if((server_ack->cType == DIAG_CMD_TYPE_UDP) && (server_ack->mType == DIAG_MSG_TYPE_ACK))
{
time = currentTime - server_ack->Time;
seq += 1;
}
pkt->Seq += 1;
Average += time;

free(temp_ack);
Sleep(delay);
}
}
}

free(Pkt);
SendTime(buflen*seq, Average);
if(!talk && !cmd)
{
printf("+++ %d bytes in %2.2f real seconds = %2.2f KB/sec +++\n", buflen*seq, (float)Average/CLOCKS_PER_SEC, (float)Average/CLOCKS_PER_SEC == 0?0:((float)(buflen*seq)/(float)(Average*1024))*CLOCKS_PER_SEC);
}

return (TRUE);
}

BOOL UdpRecv()
{
int rlt = 0;
int FromLen;
Sign udp_ack;
char *TALK;
short cli_port;
char* cli_addr;
pSign client_req;
char temp_req[TEMPBUF];

FromLen = sizeof(sinhim);
rlt = recvfrom(testSocket_U, (char*)&temp_req, TEMPBUF, 0, (struct sockaddr*)&sinhim, &FromLen);

if (rlt < 0)
{
printf("Receive UDP packets failed.\n");
return (FALSE);
}

client_req = (pSign)temp_req;

if((client_req->cType == DIAG_CMD_TYPE_TALK) && (client_req->mType == DIAG_MSG_TYPE_REQ))
{
TALK = (char*)malloc(buflen);
memset(TALK, 0, buflen);
memcpy(TALK, (client_req+1), buflen-HEADLEN);
cli_port = (short)ntohs(sinhim.sin_port);
cli_addr = inet_ntoa(sinhim.sin_addr);
printf("Writing Received Data from %s:%d: ", cli_addr, cli_port);
printf("\t%s\n", TALK);
}

if((client_req->cType == DIAG_CMD_TYPE_CMD) && (client_req->mType == DIAG_MSG_TYPE_REQ))
{
Sleep(200);
sinme.sin_port = 0;
sinhim.sin_port = htons(port);
buflen = ((pCmd)(client_req+1))->Buflen;
delay = ((pCmd)(client_req+1))->Delay;
num = ((pCmd)(client_req+1))->Num;

if(((pCmd)(client_req+1))->Protocol == DIAG_CMD_TYPE_UDP)
{
udp = 1;
UdpSend();
}
else if(((pCmd)(client_req+1))->Protocol == DIAG_CMD_TYPE_TCP)
{
udp = 0;
SendTcp();
}
sinme.sin_port = htons(port);
}

if((client_req->cType == DIAG_CMD_TYPE_TIME) && (client_req->mType == DIAG_MSG_TYPE_REQ))
{
printf(" in %2.2f real seconds = %2.2f KB/sec +++\n", (float)client_req->Time/CLOCKS_PER_SEC, (float)client_req->Time/CLOCKS_PER_SEC == 0?0:((float)client_req->Totalbuflen/(float)(client_req->Time*1024))*CLOCKS_PER_SEC);
if(cmdflag)
{
exit (0);
}
}

if((client_req->cType == DIAG_CMD_TYPE_UDP) && (client_req->mType == DIAG_MSG_TYPE_REQ))
{
memset(&udp_ack, 0 , HEADLEN);
udp_ack.cType = DIAG_CMD_TYPE_UDP;
udp_ack.mType = DIAG_MSG_TYPE_ACK;
udp_ack.Seq = client_req->Seq;
udp_ack.Time = client_req->Time;
udp_ack.Buflen= 0;
if(udp == 0 )
{
udp = 1;
}
else
udp += 1;

if(sendto(testSocket_U, (char*)&udp_ack, HEADLEN, 0, (struct sockaddr*)&sinhim, sizeof(struct sockaddr)) == INVALID_SOCKET)
{
printf("Reply failed.\n");
return (FALSE);
}

if(client_req->Seq == client_req->Sendnum)
{
cli_port = (short)ntohs(sinhim.sin_port);
cli_addr = inet_ntoa(sinhim.sin_addr);
printf("+++ Received UDP %d Bytes from %s:%d", cmd==1?(client_req->Buflen+HEADLEN)*(udp-1):(client_req->Buflen+HEADLEN)*udp, cli_addr, cli_port);
udp = 0;
}
}
return (TRUE);
}

BOOL SendTcp()
{
int i, j, k,result, templen, tcpbuf;
char *temp_send, *temp_recv, *temp_tcp;
pSign req_tcp;
struct timeval sockTimeout;
DWORD currentTime, time;
fd_set readFds;
BOOL SockOpt;

printf("\nTCP Transmit Test:\n");
printf(" Transmit : TCP -> %s:%d\n", inet_ntoa(sinhim.sin_addr), port);
printf(" Buffer Size : %d\n", buflen);

temp_send = malloc(buflen);
memset((char *)temp_send, 0, buflen);
req_tcp = (pSign)temp_send;
req_tcp->Seq = 1;

for(k=0;num==0?k<1:k<num;k++)
{
if ((fd = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET)
mes("socket");

SockOpt = TRUE;
result = setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (char*)&SockOpt, sizeof(SockOpt));

if (bind(fd, (struct sockaddr *)&sinme, sizeof(sinme)) == SOCKET_ERROR)
mes("bind");

if(connect(fd, (struct sockaddr*)&sinhim, sizeof(sinhim)) == SOCKET_ERROR)
mes("connect");

req_tcp->cType = DIAG_CMD_TYPE_TCP;
req_tcp->mType = DIAG_MSG_TYPE_REQ;
req_tcp->Buflen = (buflen - HEADLEN);
num==0?(req_tcp->Sendnum=1):(req_tcp->Sendnum=num);
if(num == 0)
req_tcp->Sendnum = 1;
else
req_tcp->Sendnum = num;

for(j=HEADLEN,i=65;j<buflen;j++,i++)
{
if(i==91)
{
i = 48;
*(temp_send+j) = i;
}
else if(i==58)
{
i = 97;
*(temp_send+j) = i;
}
else if(i==123)
{
i = 65;
*(temp_send+j) = i;
}
else
*(temp_send+j) = i;
}

req_tcp->Time = GetTickCount();

if(send(fd, (char *)req_tcp, buflen, 0) == INVALID_SOCKET)
mes("send");

temp_recv = (char*)malloc(HEADLEN);
temp_tcp = (char*)malloc(HEADLEN);
memset((char *)temp_recv, 0, HEADLEN);
memset((char *)temp_tcp, 0, HEADLEN);

sockTimeout.tv_sec = TIMEOUT*4;
sockTimeout.tv_usec = 0;

/* Wait for reply at the ephemeral port selected by the sendto () call. */
FD_ZERO (&readFds);
FD_SET (fd, &readFds);

result = select (0, &readFds, NULL, NULL, &sockTimeout);

if (result < 0)
{
printf("TCP select failed.\n");
free(temp_recv);
free(temp_tcp);
free(temp_send);
return (FALSE);
}

if (result == 0) /* TIMEOUT interval expired. */
{
printf("Request time out.\n");
free(temp_recv);
free(temp_tcp);
free(temp_send);
return (FALSE);
}

if (result > 0)
{
if(req_tcp->cType == DIAG_CMD_TYPE_TCP && req_tcp->mType == DIAG_MSG_TYPE_REQ && req_tcp->Cmd[0] == 0)
{
templen = 0;
tcpbuf = HEADLEN;

while(1)
{
result = recv(fd, (char*)temp_recv, tcpbuf, 0);
if(result == -1)
{
mes("recv");
}

if(result> 0 && result < tcpbuf)
{
tcpbuf -= result;
memcpy((char *)(temp_tcp + templen), temp_recv, result);
templen += result;
}
else if(result > 0 && result == tcpbuf)
{
memcpy((temp_tcp + templen), temp_recv, result);
break;
}
}
}

if(((pSign)temp_tcp)->cType == DIAG_CMD_TYPE_TCP && ((pSign)temp_tcp)->mType == DIAG_MSG_TYPE_ACK)
{
currentTime = GetTickCount();
time = currentTime - ((pSign)temp_tcp)->Time;
req_tcp->Seq += 1;
}

}
SendTime(buflen, time);
printf("+++ %d bytes in %2.2f real seconds = %2.2f KB/sec +++\n", buflen, (float)time/CLOCKS_PER_SEC, (float)time/CLOCKS_PER_SEC == 0?0:((float)(buflen)/(float)(time*1024))*CLOCKS_PER_SEC);
Sleep(delay);
free(temp_recv);
free(temp_tcp);
shutdown(fd, SD_BOTH);
closesocket(fd);
}

free(temp_send);
return (TRUE);
}