// WMWinSocket2.cpp: implementation of the WMWinSocket2 class.
//
//////////////////////////////////////////////////////////////////////
//
// Copyright (c) 2001 Robert Cunnings. All rights reserved.
//
// You may copy, modify, distribute or publish this code free of charge.
//
// NO WARRANTY: This software is provided AS IS, without warranty of any kind.
//
//////////////////////////////////////////////////////////////////////
//
// WMWinSocket2 class - a class for Winsock 2 client support.
// 
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "WMWinSocket2.h"
#include <fstream>
//////////////////////////////////////////////////////////////////////
// WMWinSocket2 Methods
//////////////////////////////////////////////////////////////////////

WMWinSocket2::WMWinSocket2(DWORD dwTimeOut)
{
	m_dwTimeOut = dwTimeOut;
	m_bInitialized = false;
	m_bConnected = false;
	m_hSocket = INVALID_SOCKET;

	// initialize winsock2
    WORD wVersionRequested = MAKEWORD(VERSION_MAJOR, VERSION_MINOR); 
	WSADATA wsaData;
	m_nLastError = WSAStartup(wVersionRequested, &wsaData);
	if (m_nLastError != 0)
	{
		return;
	}
    // Ensure that the WinSock 2 DLL supports the version we want
    if (LOBYTE(wsaData.wVersion) != VERSION_MAJOR || 
        HIBYTE(wsaData.wVersion) != VERSION_MINOR)
	{
		m_nLastError = WSAVERNOTSUPPORTED;
        WSACleanup(); 
        return; 
    }
	// create event
	m_hEvent = WSACreateEvent();
	if (m_hEvent == WSA_INVALID_EVENT)
	{
		m_nLastError = WSAGetLastError();
		WSACleanup();
		return;
	}
	m_bInitialized = true;
}

WMWinSocket2::~WMWinSocket2()
{
	Clear();
	if (m_bInitialized == true)
	{
		WSACloseEvent(m_hEvent);
		WSACleanup();
	}
}

void WMWinSocket2::Clear()
{
    // if socket is open, close it
	if (m_hSocket != INVALID_SOCKET)
    {
        closesocket(m_hSocket);
	    m_hSocket = INVALID_SOCKET;
    }
	m_bConnected = false;
	// release any unsent buffers
	deque<LPWSABUF>::iterator iD;
	for (iD = m_dequeDataOut.begin();
			iD != m_dequeDataOut.end();
			iD++)
	{
		LPWSABUF pBuf = *iD;
		if (pBuf != NULL)
		{
			if (pBuf->buf != NULL)
			{
				free(pBuf->buf);
			}
			free(pBuf);
		}
	}
}


bool WMWinSocket2::Connect(SOCKADDR_IN *pAddress)
{
	if (m_bInitialized == false)
	{
		m_nLastError = WSANOTINITIALISED;
		return false;
	}
	if (m_bConnected == true)
	{
		m_nLastError = WSAEISCONN;
		return false;
	}
	// reset Last Error
	m_nLastError = 0;
	// reset flag
	m_bConnected = false;
	// create socket
	m_hSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (m_hSocket == INVALID_SOCKET)
	{
		m_nLastError = WSAGetLastError();
		return false;
	}
	// select events
	int nResult = WSAEventSelect(m_hSocket, m_hEvent, FD_CONNECT);
	if (nResult == SOCKET_ERROR)
	{
		m_nLastError = WSAGetLastError();
		closesocket(m_hSocket);
        m_hSocket = INVALID_SOCKET;
		return false;
	}
	// make connection
	nResult = connect(m_hSocket, (LPSOCKADDR)pAddress, sizeof(SOCKADDR_IN));
	if (nResult == SOCKET_ERROR)
	{
		nResult = WSAGetLastError();
		if (nResult != WSAEWOULDBLOCK)
		{
			m_nLastError = nResult;
			closesocket(m_hSocket);
            m_hSocket = INVALID_SOCKET;
			return false;
		}
	}
	// wait for events
	while (true)
	{
		DWORD dwResult =  WSAWaitForMultipleEvents(1, &m_hEvent, FALSE, m_dwTimeOut, FALSE);
		// done waiting
		if (dwResult == WSA_WAIT_TIMEOUT)
		{
			m_nLastError = WSA_WAIT_TIMEOUT;
			closesocket(m_hSocket);
            m_hSocket = INVALID_SOCKET;
			m_bConnected = false;
			return false;
		}
		// get event info
		WSANETWORKEVENTS netEvents;
		nResult = WSAEnumNetworkEvents(m_hSocket, m_hEvent, &netEvents);
		if (nResult == SOCKET_ERROR)
		{
			goto failure;
		}
		if (netEvents.lNetworkEvents & FD_CONNECT)
		{
			// connect event received
			if (netEvents.iErrorCode[FD_CONNECT_BIT] != 0)
			{
				goto failure;
			}
			int nResult = WSAEventSelect(m_hSocket, m_hEvent, FD_READ | FD_WRITE | FD_CLOSE);
			if (nResult == SOCKET_ERROR)
			{
				m_nLastError = WSAGetLastError();
				closesocket(m_hSocket);
                m_hSocket = INVALID_SOCKET;
				return false;
			}
			// set flag
			m_bConnected = true;
			// call virtual overide
			if (OnSockConnect() == false)
			{
				break;
			}
		}
		if (netEvents.lNetworkEvents & FD_WRITE)
		{
			// write event received
			if (netEvents.iErrorCode[FD_WRITE_BIT] == WSAENETDOWN)
			{
				goto failure;
			}
			while (m_dequeDataOut.empty() == false)
			{
				// send data
				LPWSABUF pBuf = m_dequeDataOut.front();
				unsigned long uLength = pBuf->len;
				char *pData = pBuf->buf;
				m_dequeDataOut.pop_front();
				// now try to send
				int nResult;
				do
				{
					nResult = send(m_hSocket, pData, uLength, 0);
					if (nResult == SOCKET_ERROR)
					{
						m_nLastError = WSAGetLastError();
						break;
					}
					pData += nResult;
					uLength -= nResult;
					Sleep(0);
				} while (uLength > 0);
				// did it all get sent?
				if (nResult == SOCKET_ERROR)
				{
					// would block?
					if (m_nLastError == WSAEWOULDBLOCK)
					{
						// add remaining data back to queue
						LPWSABUF pTemp = (LPWSABUF)malloc(sizeof(WSABUF));
						pTemp->buf = (char *)malloc(uLength);
						memcpy(pTemp->buf, pData, uLength);
						pTemp->len = uLength;
						m_dequeDataOut.push_front(pTemp);
						break;
					}
					else
					{
						// failure
						free (pBuf->buf);
						free (pBuf);
						goto failure;
					}
				}
				// done with buffer
				free (pBuf->buf);
				free (pBuf);
			}
			// call virtual overide
			if (OnSockWrite() == false)
			{
				break;
			}
		}
		if (netEvents.lNetworkEvents & FD_READ)
		{
			// read event received
			if (netEvents.iErrorCode[FD_READ_BIT] == WSAENETDOWN)
			{
				goto failure;
			}
			// call virtual overide
			if (OnSockRead() == false)
			{
				break;
			}
		}
		if (netEvents.lNetworkEvents & FD_CLOSE)
		{
			// close event received
			if (netEvents.iErrorCode[FD_CLOSE_BIT] == 0)
			{
				// no error... call virtual override
				OnSockClose();
				break;
			}
			goto failure;
		}
	}
	// closed
	closesocket(m_hSocket);
	m_hSocket = INVALID_SOCKET;
	m_bConnected = false;
	return true;

failure:
	m_nLastError = WSAGetLastError();
	closesocket(m_hSocket);
	m_hSocket = INVALID_SOCKET;
	m_bConnected = false;
	return false;
}

bool WMWinSocket2::Send(char *pszBuffer, int nBufLen)
{
	// are we able to send?
	if (m_dequeDataOut.empty() == false)
	{
		// must not be, so add to queue
		LPWSABUF pTemp = (LPWSABUF)malloc(sizeof(WSABUF));
		pTemp->buf = (char *)malloc(nBufLen);
		memcpy(pTemp->buf, pszBuffer, nBufLen);
		pTemp->len = nBufLen;
		m_dequeDataOut.push_back(pTemp);
		return true;
	}
	// send data
	int nResult;
	do
	{
		nResult = send(m_hSocket, pszBuffer, nBufLen, 0);
		if (nResult == SOCKET_ERROR)
		{
			m_nLastError = WSAGetLastError();
			break;
		}
		pszBuffer += nResult;
		nBufLen -= nResult;
		Sleep(0);
	} while (nBufLen > 0);
	// what happened?
	if (nResult == SOCKET_ERROR)
	{
		// would block?
		if (m_nLastError == WSAEWOULDBLOCK)
		{
			// add remaining data to queue
			LPWSABUF pTemp = (LPWSABUF)malloc(sizeof(WSABUF));
			pTemp->buf = (char *)malloc(nBufLen);
			memcpy(pTemp->buf, pszBuffer, nBufLen);
			pTemp->len = nBufLen;
			m_dequeDataOut.push_front(pTemp);
			return true;
		}
		// failure
		return false;
	}
	return true;
}

bool WMWinSocket2::Recv(char *pszBuffer, int nBufLen, int *pnNumBytes)
{
	bool bRtn = true;
	int nResult = recv(m_hSocket, pszBuffer, nBufLen, 0);
	if (nResult == SOCKET_ERROR)
	{
		m_nLastError = WSAGetLastError();
		if (m_nLastError != WSAEWOULDBLOCK)
		{
			bRtn = false;
		}
	}
	*pnNumBytes = nResult;
	return bRtn;
}
