///////////////////////////////////////////////////////////////////////
//  sockUtils.cpp-sockets helper functions                           //
//  ver 1.1                                                          //
//                                                                   //
//  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) {

    //////////////////////////////////////////////////////////////
    // a SOCKET is a handle referencing the new socket 
    // created by accept for the connecting client.  The 
    // threadProc passed to _beginthreadex, below, needs
    // to make a copy of this descriptor, as it will be
    // overwritten the next time a client connects

    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 sendAll(SOCKET &sock, const char *block) {

  const int sendRetries = 100;

  int blockLen = static_cast<int>(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;
}

//

//----< get system error message string >----------------------

string sysError::getLastMsg() {

// ask system what type of error occurred

  DWORD errorCode = GetLastError();

// map errorCode into a system defined error string
    
  DWORD dwFlags = 
    FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_ALLOCATE_BUFFER;
  LPCVOID lpSource = NULL;
  DWORD dwMessageID = errorCode;
  DWORD dwLanguageId = MAKELANGID(LANG_ENGLISH,SUBLANG_ENGLISH_US);
  LPSTR lpBuffer;
  DWORD nSize = 0;
  va_list *Arguments = NULL;

  FormatMessage(
    dwFlags,lpSource,dwMessageID,dwLanguageId, 
    (LPTSTR)&lpBuffer,nSize,Arguments
  );

  _msg = lpBuffer;
  LocalFree(lpBuffer);
  return _msg;
}

//

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

#include <iostream>

unsigned _stdcall threadProc(void *pArgs) {

  /////////////////////////////////////////////////////////////
  // Make a copy of socket handle as it will be overwritten
  // the next time a new client is connected.

  SOCKET sock = *((SOCKET*)pArgs);

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

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

  do {
    if(sock == INVALID_SOCKET)
      break;
    try
    {
      bytes = recv(sock,buffer,bufSize-1);
      if(bytes==SOCKET_ERROR)
        throw exception("connection to client broken");
    }
    catch(exception& e)
    {
      cout << "\n  " << e.what() << "\n\n";
      break;
    }
    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 << " aborting this client\n\n";
    LeaveCriticalSection(&cs);
    return errType;
  }
  EnterCriticalSection(&cs);
  cout << "\n  returning from threadProc\n";
  LeaveCriticalSection(&cs);
  closesocket(sock);
  return 0;
}

//

#ifdef TEST_SOCKUTILS

//----< 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,2048,threadProc);
    closesocket(server);
  #endif

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

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

    for(int i=0; i<500; ++i)
    {
      SYSTEMTIME st;
      ::GetSystemTime(&st);
      ostringstream systime;
      systime << "client time is " << st.wHour << ":" 
        << st.wMinute << ":" << st.wSecond;
      cout << "\n  sending \"" << systime.str().c_str() << "\"\n";
      try
      {
        sendAll(client,systime.str().c_str());
      }
      catch(exception& e)
      {
        cout << "\n  " << e.what();
        break;
      }
      ::Sleep(100);
    }
    shutdown(client,SD_SEND);
    closesocket(client);
  #endif

    stopSockets();
  }
  catch(exception& e) { std::cout << "\n  " << e.what() << "\n\n"; }
  catch(const string &msg) { std::cout << msg; }

  return 0;
}

#endif