// Copyright (c) 1999  Microsoft Corporation

#include "speedy.h"

BOOL		SendFlag		= FALSE;

PCHAR		HostIp,
			RandomBuffer,
			Buffer;

int			SockReceiveBuf	= 256 * KILO,
			SockSendBuf		= 0 * KILO;

long		LengthOfBuffer	= 256 * KILO,
			NumberOfBuffers	= 20  * KILO;

SOCKET		sd				= INVALID_SOCKET;
PERF_INFO	PerfInformation	= {0};

// FUNC: pexit
void pexit(PCHAR r, PCHAR s) {

    printf ("%s: %s ERROR %d\n", r, s, (int) GetLastError()); exit(1);
}

// FUNC: WSAexit
void WSAexit(PCHAR r, PCHAR s) {

    LPTSTR  pbuf;
        
    if (0 < FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM|FORMAT_MESSAGE_ALLOCATE_BUFFER|FORMAT_MESSAGE_MAX_WIDTH_MASK,
        NULL, WSAGetLastError(), 0, (LPTSTR)&pbuf, 0, NULL)) {

        printf ("%s:%s \n%s\n(error: %ld)\n", r, s, pbuf, WSAGetLastError());
	    LocalFree (pbuf);

    } else {

        printf ("%s:%s Error: %ld\n", r, s, WSAGetLastError());
    }

	WSACleanup();
    exit(1);
}

// FUNC: process_args
void process_args(int argc, char *argv[]) {

	int	n = 0,
		i = 1;

	BOOL ReceiveFlag = FALSE;

	if (argc < 4) {

		printf ("\nspeedy [-s/-r] <-l buffer length> <-n buffer number> [-ip receiver IP]\n");
		exit(1);
    }

	while ((i+1) <= argc) {

        if (0 == _stricmp(argv[i], "-s")) {

				SendFlag = TRUE;
				i++;

		} else if (0 == _stricmp(argv[i], "-r")) {

			    ReceiveFlag = TRUE;
				i++;

		} else if (0 == _stricmp(argv[i], "-l")) {

			if (NULL != argv[i+1]) {

				LengthOfBuffer = atoi(argv[i+1]);
			}

			if ((8 > LengthOfBuffer) || (LengthOfBuffer % 8 != 0) || (NULL == argv[i+1])) {

				pexit("process_args", "-l option has to be higher than 8 and multiple of 8");
			}

			i += 2;

		} else if (0 == _stricmp(argv[i], "-n")) {

			if (NULL != argv[i+1]) {

				NumberOfBuffers = atoi(argv[i+1]);
			}

			if (0 >= NumberOfBuffers || NULL == argv[i+1]) {

				pexit("process_args", "-n option");
			}

			i += 2;

			} else if (0 == _stricmp(argv[i], "-sb")) {

			if (NULL != argv[i+1]) {

				SockSendBuf = atoi(argv[i+1]);
			}

			if (0 > SockSendBuf || NULL == argv[i+1]) {

				pexit("process_args", "-sb option");
			}

			i += 2;

			} else if (0 == _stricmp(argv[i], "-rb")) {

			if (NULL != argv[i+1]) {

				SockReceiveBuf = atoi(argv[i+1]);
			}

			if (0 > SockReceiveBuf || NULL == argv[i+1]) {

				pexit("process_args", "-rb option");
			}

			i += 2;

		} else if (0 == _stricmp(argv[i], "-ip")) {

			if (NULL == argv[i+1]) {
				pexit("process_args", "-ip option");
			}

			HostIp = argv[i+1];

			i += 2;

		} else {

			pexit("process_args", "unknown option");
		}
	}

	if ((TRUE == ReceiveFlag) && (TRUE == SendFlag)) {

		pexit("process_args", "cannot be both send and receive");
	}

	if ((FALSE == ReceiveFlag) && (FALSE == SendFlag)) {

		pexit("process_args", "need to specify either send or receive");
	}


	return;

}

// FUNC: setup
void setup(PCHAR host) {

	int		i,
			port			= 5001,
			iterations		= 0, 
            remainder		= 0,
			LengthOfca;
    
	struct	sockaddr_in  sa, 
                         ca;

	WSADATA              wd;
	

	if (0 == SetThreadAffinityMask(GetCurrentThread(), 1)) {
	
		pexit("Setup", "SetThreadAffinity");
	}

	RandomBuffer	= (PCHAR) VirtualAlloc(NULL, LengthOfBuffer, MEM_COMMIT, PAGE_READWRITE);

	if (NULL == RandomBuffer) {
			pexit("setup", "VirtualAlloc RandomBuffer");
	}

	if (TRUE == SendFlag) {

		for (i=0; i < LengthOfBuffer; i++) {

			 RandomBuffer[i] = rand();
		}
	}

	// Network Setup

    if (WSAStartup( MAKEWORD (2,2), &wd ) != 0) {

        WSAexit("setup", "WSAStartup");
    }
	
    sd = socket(AF_INET, SOCK_STREAM, 0);

	if (INVALID_SOCKET == sd) {

        WSAexit("setup", "create socket tcp");
    }

	if (SOCKET_ERROR == setsockopt(sd, SOL_SOCKET, SO_SNDBUF, (char *) &SockSendBuf, sizeof(SockSendBuf))) {

        WSAexit("setup", "WSAStartup SNDBUF");
    } 

	if (SOCKET_ERROR == setsockopt(sd, SOL_SOCKET, SO_RCVBUF, (char *) &SockReceiveBuf, sizeof(SockReceiveBuf))) {

        WSAexit("setup", "WSAStartup RCVBUF");
    }

	memset(&sa, 0, sizeof(struct sockaddr_in));
            
    sa.sin_family		= AF_INET;
    sa.sin_addr.s_addr	= inet_addr(host);
    sa.sin_port			= htons((unsigned short) port);

    if (TRUE == SendFlag) {

       if (SOCKET_ERROR == connect(sd, (const struct sockaddr FAR *) &sa, sizeof(struct sockaddr_in))) {
           WSAexit("setup", "connect_tcp");
       }

	   printf("connected\n");

    } else {

       if (SOCKET_ERROR == bind(sd, (struct sockaddr *) &sa, sizeof(struct sockaddr_in))) {
           WSAexit("setup", "bind tcp");
       }

	   printf("bound\n");

       if (SOCKET_ERROR == listen(sd, 0)) {
           WSAexit("setup", "listen tcp");
       }

	   printf("listening\n");

	   LengthOfca = sizeof(struct sockaddr_in);

       sd = accept(sd, (struct sockaddr *) &ca, &LengthOfca);

	   printf("accepted\n");
    }

	printf("Network Activity progressing...\n");
}

// FUNC: work
void work(void) {

	int				CheckSum		= 0,
                    BuffersLeft		= 0,
					SequenceNumber	= 0,
					BufferCheckSum  = 0,
					BytesTransferred,
					BytesLeft;		


	struct _timeb   time0;
	struct _timeb   time1;


	* (int *) &RandomBuffer[0] = GetCheckSum(&RandomBuffer[2*sizeof(int)], LengthOfBuffer - (2 * sizeof (int)));


    BytesTransferred  = 0;
    BytesLeft         = LengthOfBuffer;

    BuffersLeft   = NumberOfBuffers;

	if (TRUE == SendFlag) {

		_ftime( &time0 );

		while (0 < BuffersLeft--) {

				* (int *) &RandomBuffer[sizeof(int)]	= SequenceNumber++;
				
				BytesLeft = LengthOfBuffer;
				
				BytesTransferred = send (sd, RandomBuffer, LengthOfBuffer, 0);

				if (0 == BytesTransferred) {	
						goto Done;
				}

				PerfInformation.Bytes += BytesTransferred;

				++(PerfInformation.TransCount);

				BytesLeft -= BytesTransferred;

				if (0 != BytesLeft) {

							Buffer = RandomBuffer + BytesTransferred;
				}

				while (0 != BytesLeft) {

						BytesTransferred = send (sd, Buffer, BytesLeft, 0);

						if (0 == BytesTransferred) {	
								goto Done;
						}

						PerfInformation.Bytes += BytesTransferred;

						++(PerfInformation.TransCount);

						Buffer    += BytesTransferred;
						BytesLeft -= BytesTransferred;
				}				
		}

		_ftime( &time1 ); 

	} else {									// This is a receiver

		_ftime( &time0 );

		while (0 < BuffersLeft--) {

				BytesLeft = LengthOfBuffer;
				
				BytesTransferred = recv (sd, RandomBuffer, LengthOfBuffer, 0);

				if (0 == BytesTransferred) {	
						goto Done;
				}

				PerfInformation.Bytes += BytesTransferred;

				++(PerfInformation.TransCount);

				BytesLeft -= BytesTransferred;

				if (0 != BytesLeft) {

							Buffer = RandomBuffer + BytesTransferred;
				}

				while (0 != BytesLeft) {

						BytesTransferred = recv (sd, Buffer, BytesLeft, 0);

						if (0 == BytesTransferred) {	
								goto Done;
						}

						PerfInformation.Bytes += BytesTransferred;

						++(PerfInformation.TransCount);

						Buffer    += BytesTransferred;
						BytesLeft -= BytesTransferred;
				}

				if (*(int *) &RandomBuffer[sizeof(int)] != SequenceNumber++) {

					pexit("work", "sequence error");
				}

				if (*(int *) &RandomBuffer[0] != GetCheckSum(&RandomBuffer[2*sizeof(int)], LengthOfBuffer - (2 * sizeof (int)))) {

					pexit("work", "checksum error");
				}				
		}

		_ftime( &time1 ); 
	}
	
Done:
   
     // clean up
    
	PerfInformation.Realtime = (time1.time - time0.time) + ((time1.millitm - time0.millitm)/(double)MS2S);

	closesocket(sd);

	VirtualFree(RandomBuffer, 0, MEM_RELEASE);

	WSACleanup();
}

// FUNC: output 
void output(void) {

    int             i = 0; 

    long            TransCount			= 0;

    double          Throughput          = 0.0, 
                    TotalThroughput		= 0.0, 
                    TotalMbytes			= 0.0,
					FramesCount			= 0.0,
                    BytesTransferred	= 0.0, 
                    Realtime            = -1;


    fprintf(stdout, "\n\nThread Realtime(s) Throughput(KB/s) Throughput(Mbit/s)\n");
    fprintf(stdout, "====== =========== ================ ==================\n\n");

		Throughput = (PerfInformation.Bytes/1000.0) / PerfInformation.Realtime;

        TotalMbytes += (PerfInformation.Bytes)/(1000.0); 

        TotalThroughput += Throughput;

        if ((double)(PerfInformation.Realtime) > Realtime) {
            Realtime = PerfInformation.Realtime;
        }

        TransCount += PerfInformation.TransCount;

        fprintf(stdout, "%6d   %9.3f        %9.3f          %9.3f\n",
                i, PerfInformation.Realtime, Throughput, Throughput*8/1000);

    fprintf(stdout, "\n\nTotal Bytes(MEG) Realtime(s) Total Throughput(Trans/s) Total Throughput(Mbit/s)\n");
    fprintf(stdout, "================ =========== ========================= ========================\n\n");
    fprintf(stdout, "%16.6f   %9.3f                 %9.3f                %9.3f\n", 
                    TotalMbytes/1000, Realtime, TransCount/Realtime, TotalThroughput*8/1000);

	FramesCount =  (TotalThroughput*1000*Realtime)/(LengthOfBuffer);
}

// FUNC: main
void __cdecl main(int argc, char * argv[]) {

	process_args(argc, argv);

	setup(HostIp);

	work();

	if (FALSE == SendFlag) {
		output();
	}
}

// FUNC: GetCheckSum
__inline int GetCheckSum (PCHAR Buffer, int Length) {

		// assume that Length is a multiple of 8 bytes

		int * First;
		int * Second;
		int i;
		int CheckSum = 0;

		First  = (int *) &Buffer[0];
		Second = (int *) &Buffer[sizeof(int)];

		for (i=0; i < Length; i += 2 * sizeof(int)) {

					CheckSum += *First + *Second;

					First  += 2;
					Second += 2;
		}

		return(CheckSum);
}