///////////////////////////////////////////////////////////////////////
//  sockUtils.cpp-sockets helper functions                           //
//  ver 1.0                                                          //
//                                                                   //
//  Language:     Visual C++, ver 6.0                                //
//  Platform:     Dell XBS1000r, Windows 2000 Professional           //
//  Application:  Demonstration for CSE691 - SW Mod & Anal           //
//  Author:       Jim Fawcett, CST 2-187, Syracuse Univ              //
//                (315) 443-3948, jfawcett@twcny.rr.com              //
///////////////////////////////////////////////////////////////////////

#include <process.h>
#include <sstream>
#include "sockUtils.h"
using namespace std;

CRITICAL_SECTION cs;

//----< throw exception string >---------------------------------------

void throwString(const char *msg, const char *file, int line) {

  int errType = WSAGetLastError();
  ostringstream collect;
  collect << msg << " " << errType
          << "\nfile: " << file 
          << "\nline: " << line << '\n' << ends;
  const char *debug = collect.str().c_str();
  throw collect.str();
}

//----< start sockets >------------------------------------------------

void startSockets() {

  ///////////////////////////////////////////////////////////////////
  // load version 1.1 of winsock dll

  WORD wVersionRequested = MAKEWORD(1,1);   // requesting version 1.1
  WSAData wData;                            // startup data filled
                                            // by WSAStartup

  int err = WSAStartup(wVersionRequested, &wData);

  if(err == SOCKET_ERROR)
    throwString("initialization error type: ",__FILE__,__LINE__);

  InitializeCriticalSection(&cs);
}

//

//----< shutdown sockets >---------------------------------------------

void stopSockets() {

  WSACleanup();
  DeleteCriticalSection(&cs);
}

//----< create client TCP/IP socket object >---------------------------

SOCKET createSocket() {

  return socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
}

//----< connect client socket to server >------------------------------

void connectToServer(SOCKET &sock, const char *IPaddress, int port) {
  
  SOCKADDR_IN tcpAddr;
  tcpAddr.sin_family = AF_INET;
  tcpAddr.sin_addr.s_addr = inet_addr(IPaddress);
  tcpAddr.sin_port = htons(port);

  int err = connect(sock, (sockaddr*)&tcpAddr, sizeof(tcpAddr));

  if(err == SOCKET_ERROR)
    throwString("connection error type: ",__FILE__,__LINE__);
}

//

//----< listen for and accept connections >----------------------------

typedef unsigned(_stdcall *fnPtr)(void*);

void startSocketServer(SOCKET &listener, int port, fnPtr threadProc) {

  SOCKADDR_IN tcpAddr;
  tcpAddr.sin_family = AF_INET;   // TCP/IP
  tcpAddr.sin_port = htons(port); // listening port
  tcpAddr.sin_addr.s_addr = INADDR_ANY;
                                  // listen over every network interface
  int err = bind(listener, (SOCKADDR*)&tcpAddr, sizeof(tcpAddr));

  if(err == SOCKET_ERROR)
    throwString("binding error type:",__FILE__,__LINE__);

  /////////////////////////////////////////////////////////////////
  // listen for incoming connection requests

  int backLog = 10;
  err = listen(listener, backLog);

  if(err == SOCKET_ERROR)
    throwString("listen error type:",__FILE__,__LINE__);

  /////////////////////////////////////////////////////////////////
  // block until a connection is established

  int size = sizeof(tcpAddr);

  while(1) {

    SOCKET toClient = accept(listener, (SOCKADDR*)&tcpAddr, &size);      

    if(toClient == INVALID_SOCKET)
      break;

  /////////////////////////////////////////////////////////////////
  // got connection so spawn thread to read transferred data

    unsigned int threadID;

  /////////////////////////////////////////////////////////////////
  // make sure you close toClient socket at end of threadProc

    HANDLE hThread = 
      (HANDLE)_beginthreadex(NULL,0,threadProc,&toClient,0,&threadID);

    CloseHandle(hThread);
  }
  closesocket(listener);
}

//

//----< keep sending until socket has eaten all bytes >----------------

int send(SOCKET &sock, const char *block) {

  const int sendRetries = 100;

  int blockLen = strlen(block);
  int bytesSent;
  int bytesLeft = blockLen;
  int blockIndx = 0;
  int count = 0;
  while(bytesLeft > 0) {
    bytesSent = ::send(sock,&block[blockIndx],bytesLeft,0);
    if(bytesSent == SOCKET_ERROR) {
      Sleep(50);
      ++count;
    }
    if(count==sendRetries)
      throwString("sendAll failed after 100 retries",__FILE__,__LINE__);

    bytesLeft -= bytesSent;
    blockIndx += bytesSent;
  }
  return blockLen;
}

//----< receive on socket connection >---------------------------------

int recv(SOCKET &sock, char *block, int blockLen) {

  return ::recv(sock,block,blockLen,0);
}

//----< keep receiving until socket has released all bytes >-----------
//
//  don't use this unless you know how many bytes the client
//  will send otherwise the function may wait forever
//
int recvAll(SOCKET &sock, char *block, int blockLen) {

  const int recvRetries = 100;
  int bytesRecvd, bytesLeft = blockLen;
  int blockIndx = 0, count = 0;
  while(bytesLeft > 0) {
    bytesRecvd = recv(sock,&block[blockIndx],bytesLeft,0);
    if(bytesRecvd == SOCKET_ERROR) {
      Sleep(50);
      ++count;
    }
    if(count==recvRetries)
      throwString("recvAll failed after 100 retries",__FILE__,__LINE__);

    bytesLeft -= bytesRecvd;
    blockIndx += bytesRecvd;
  }
  return blockLen;
}

//

//----< test stub >----------------------------------------------------

#include <iostream>

unsigned _stdcall threadProc(void *pArgs) {

  SOCKET *pSock = (SOCKET*)pArgs;

#ifdef TEST_SERVER
  const int bufSize = 8;
#else
  const int bufSize = 1024;
#endif

  char buffer[bufSize];
  int bytes, count=0;
  string save;

  do {
    if(*pSock == INVALID_SOCKET)
      break;
    bytes = recv(*pSock,buffer,bufSize-1);
    buffer[bytes] = '\0';
    save = buffer;
    if(bytes > 0) {
      EnterCriticalSection(&cs);
      cout << "\n  block " << ++count << ": " << save.c_str();
      LeaveCriticalSection(&cs);
    }
  } while(bytes != 0 && bytes != SOCKET_ERROR);

  if(bytes == SOCKET_ERROR) {
    int errType = WSAGetLastError();
    EnterCriticalSection(&cs);
    cout << "\n  receive error type: " << errType << "\n\n";
    LeaveCriticalSection(&cs);
    return errType;
  }
  EnterCriticalSection(&cs);
  cout << "\n  returning from threadProc\n";
  LeaveCriticalSection(&cs);
  closesocket(*pSock);
  return 0;
}

//

//----< start testing here >-------------------------------------------

int main() {

#ifdef SOCKUTILS
  cout << "\n  Testing SockUtils "
       << "\n ===================\n";

  try {
    throwString("throw message",__FILE__,__LINE__);
  }
  catch(const string &msg) { std::cout << msg; }
#endif

  try {
    startSockets();

  #ifdef TEST_SERVER
    cout << "\n  Testing SockUtils - Server Side "
         << "\n =================================\n";

    SOCKET server = createSocket();
    startSocketServer(server,1024,threadProc);
    closesocket(server);
  #endif

  #ifdef TEST_CLIENT
    cout << "\n  Testing SockUtils - Client Side "
         << "\n =================================\n";

    SOCKET client = createSocket();    
    connectToServer(client,"127.0.0.1",1024);

    char *sendBuff = "Hello World via Sockets";
    cout << "\n  sending \"" << sendBuff << "\"\n";
    send(client,sendBuff);

    shutdown(client,SD_SEND);
    closesocket(client);
  #endif

    stopSockets();
  }
  catch(const string &msg) { std::cout << msg; }

  return 0;
}