///////////////////////////////////////////////////////////////////////////////
//
// Copyright (c) 1997-1999 by Bitek System Inc.
// All rights reserved.
//
// ZXSocket.cpp     : implementation of the CZXSocket class.
//

////////////////////////////////////////////////////////////////////////////////
//
// DESCRIPTION
//
// CZXSocket class
//

#include "stdafx.h"
#include "ZXSocket.h"

////////////////////////////////////////////////////////////////////////////////
//
// CZXSocket
//

////////////////////////////////////////
// Public Members :
//

//
// Constructor & Destructor
//

CZXSocket::CZXSocket(int nConnectTimeout /*= 10*/, int nReadTimeout /*= 10*/, int nWriteTimeout /*= 10*/)
{
    m_nSocket = INVALID_SOCKET;

    m_szConnectedServer[0] = '\0';
    m_nConnectedPort = 0;

    m_nConnectTimeout = nConnectTimeout;
    m_nReadTimeout = nReadTimeout;
    m_nWriteTimeout = nWriteTimeout;

    m_bCancel = FALSE;
	m_bConnected = FALSE;
}

CZXSocket::~CZXSocket()
{
    Close();
}

void CZXSocket::SetConnectTimeout(int nConnectTimeout)
{
    m_nConnectTimeout = nConnectTimeout;
}

void CZXSocket::SetReadTimeout(int nReadTimeout)
{
    m_nReadTimeout = nReadTimeout;
}

void CZXSocket::SetWriteTimeout(int nWriteTimeout)
{
    m_nWriteTimeout = nWriteTimeout;
}


int CZXSocket::Open(BOOL bServerMode)
{
    long        lAsyncIO = 1;
    BOOL        bReuse = TRUE; 

    if (m_nSocket != INVALID_SOCKET)
	{
        Close();
	}

    //
    // socket is blocking at creation time.
    //
    if ((m_nSocket=socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) == INVALID_SOCKET)
        goto _fail;

    if (setsockopt(m_nSocket, SOL_SOCKET, SO_REUSEADDR, (char FAR *)&bReuse, sizeof(bReuse)) == SOCKET_ERROR)
        goto _fail;

    //
    // Set non-blocking IO
    //
    // ioctl(m_nSocket, FIONBIO, 1);
    // ioctlsocket(m_nSocket, FIONBIO, (unsigned long *)&lAsyncIO);
    //
    if (ioctlsocket(m_nSocket, FIONBIO, (LPDWORD)&lAsyncIO) == SOCKET_ERROR)
        goto _fail;

	if( bServerMode )
	{
		if (Bind(NULL, 0) == ZXFAILURE)
			goto _fail;
	}

    return ZXSUCCESS;

_fail:
    Close();

    return ZXFAILURE;
}

int CZXSocket::Close()
{
    int     nResult = ZXSUCCESS;

    if (m_nSocket == INVALID_SOCKET)
        return nResult;

    if (WSAIsBlocking())
        WSACancelBlockingCall();

    shutdown(m_nSocket, SD_BOTH);

    if (closesocket(m_nSocket) == SOCKET_ERROR)
        nResult = ZXFAILURE;

    m_nSocket = INVALID_SOCKET;
	m_bConnected = FALSE;
    return nResult;
}

int CZXSocket::Bind(LPCTSTR lpszServer, int nPort)
{
    SOCKADDR_IN sockAddr;

    if (m_nSocket == INVALID_SOCKET)
        return ZXFAILURE;

    //
    // Check if address is IP address format of hostname format.
    //
    memset(&sockAddr, 0, sizeof(sockAddr));
    sockAddr.sin_family = AF_INET;
    sockAddr.sin_addr.s_addr = (lpszServer == NULL) ? htonl(INADDR_ANY) : inet_addr(lpszServer);
    sockAddr.sin_port = htons((u_short)nPort);

    //
    // copy network address of destination to structure
    //
    if (sockAddr.sin_addr.s_addr == INADDR_NONE) {
        LPHOSTENT   lpHost;

        lpHost = gethostbyname(lpszServer);
        if (lpHost != NULL)
            sockAddr.sin_addr.s_addr = ((LPIN_ADDR)lpHost->h_addr)->s_addr;
        else {
            // WSASetLastError(WSAEINVAL);
            return ZXFAILURE;
        }
    }

    //
    // Bind the socket to the port which will system choose.
    //
    if (bind(m_nSocket, (LPSOCKADDR)&sockAddr, sizeof(SOCKADDR_IN)) == SOCKET_ERROR)
        return ZXFAILURE;

    return ZXSUCCESS;
}

int CZXSocket::Listen(int nBackLog)
{
    // Queue the this listening socket
    if (listen(m_nSocket, nBackLog) == SOCKET_ERROR)
        return ZXFAILURE;

    return ZXSUCCESS;
}

int CZXSocket::Connect(LPCTSTR lpszServer, int nPort)
{
    SOCKADDR_IN sockAddr;
    int         nState;

    if (m_nSocket == INVALID_SOCKET) {
        return ZXFAILURE;
	}

	if (lpszServer[0] == '[' 
			&& lpszServer[strlen(lpszServer) -1] == ']') {
		strcpy(m_szConnectedServer, lpszServer+1);
		m_szConnectedServer[strlen(m_szConnectedServer) -1] = '\0';
	}
	else {
		strcpy(m_szConnectedServer, lpszServer);		
	}

	m_nConnectedPort = nPort;
	
    //
    // Check if address is IP address format of hostname format.
    //
    memset(&sockAddr, 0, sizeof(sockAddr));
    sockAddr.sin_family = AF_INET;
    sockAddr.sin_addr.s_addr = inet_addr(lpszServer);
    sockAddr.sin_port = htons((u_short)nPort);

    //
    // copy network address of destination to structure
    //
    if (sockAddr.sin_addr.s_addr == INADDR_NONE) {
        LPHOSTENT   lpHost;

        lpHost = gethostbyname(lpszServer);
        if (lpHost != NULL)
            sockAddr.sin_addr.s_addr = ((LPIN_ADDR)lpHost->h_addr)->s_addr;
        else {
            // WSASetLastError(WSAEINVAL);
            return ZXFAILURE;
        }
    }

    nState = connect(m_nSocket, (LPSOCKADDR)&sockAddr, sizeof(sockAddr));
    if (m_bCancel == TRUE) {
        return ZXFAILURE;
    }

    if (nState == SOCKET_ERROR) {
        if (WSAGetLastError() == WSAEWOULDBLOCK) {
            TIMEVAL     tv = { 1, 0 /*TIMEVAL_USEC*/ };
            FD_SET      fds, exceptfds;
            int         nTimeoutRemain = 0;

            while (1) {
                FD_ZERO(&fds);
                FD_ZERO(&exceptfds);
                FD_SET(m_nSocket, &fds);
                FD_SET(m_nSocket, &exceptfds);

                nState = select(m_nSocket+1, NULL, &fds, &exceptfds, &tv);
                if (m_bCancel == TRUE) {
                    return ZXFAILURE;
                }
                if (nState == 0) { // Timeout
                    if (++nTimeoutRemain == m_nConnectTimeout)
                        return ZXFAILURE;

                    continue;
                }
                else if (nState == SOCKET_ERROR) {
                    return ZXFAILURE;
                }
                if (FD_ISSET(m_nSocket, &exceptfds)) { // Failure
                    WSASetLastError(WSAECONNREFUSED);
                    return ZXFAILURE;
                }
                if (FD_ISSET(m_nSocket, &fds)) { // Success
                    break;
                }
            } // while
        }
        else {
            return ZXFAILURE;
        }
    }

	m_bConnected = TRUE;
    return ZXSUCCESS;
}

int CZXSocket::Accept(CZXSocket* lpConnectedSocket)
{
    SOCKADDR_IN sockAddr;
    int         nLength;

    TIMEVAL     tv = { 1, 0 /*TIMEVAL_USEC*/ };
    FD_SET      fds;
    int         nTimeoutRemain = 0;

    int         nState;

    if (m_nSocket == INVALID_SOCKET)
        return ZXFAILURE;

    nLength = sizeof(sockAddr);
    while (1) {
        FD_ZERO(&fds);
        FD_SET(m_nSocket, &fds);

        nState = select(m_nSocket+1, &fds, NULL, NULL, &tv);
        if (m_bCancel == TRUE) {
            return ZXFAILURE;
        }
        if (nState == 0) { // Timeout
            if (++nTimeoutRemain == m_nReadTimeout)
                return ZXFAILURE;

            continue;
        }
        else if (nState == SOCKET_ERROR) {
            return ZXFAILURE;
        }

        if (!FD_ISSET(m_nSocket, &fds))
            continue;

        lpConnectedSocket->m_nSocket = accept(m_nSocket, (LPSOCKADDR)&sockAddr, &nLength);
        if (lpConnectedSocket->m_nSocket == INVALID_SOCKET) {
            if (WSAGetLastError() == WSAEWOULDBLOCK)
                continue;

            return ZXFAILURE;
        }
        else
            break;
    }

    _tcscpy(lpConnectedSocket->m_szConnectedServer, m_szConnectedServer);
    lpConnectedSocket->m_nConnectedPort = m_nConnectedPort;
    lpConnectedSocket->m_nConnectTimeout = m_nConnectTimeout;
    lpConnectedSocket->m_nReadTimeout = m_nReadTimeout;
    lpConnectedSocket->m_nWriteTimeout = m_nWriteTimeout;

    return ZXSUCCESS;
}

int CZXSocket::GetSockName(LPTSTR* lppszSocketAddress, int* lpnSocketPort)
{
    SOCKADDR_IN sockAddr;
    int         nLength;

    if (m_nSocket == INVALID_SOCKET)
        return ZXFAILURE;

    nLength = sizeof(sockAddr);
    if (getsockname(m_nSocket, (LPSOCKADDR)&sockAddr, &nLength) == SOCKET_ERROR)
        return ZXFAILURE;

    if (lppszSocketAddress != NULL)
        *lppszSocketAddress = inet_ntoa(sockAddr.sin_addr);
    if (lpnSocketPort != NULL)
        *lpnSocketPort = ntohs(sockAddr.sin_port);

    return ZXSUCCESS;
}

/***
int CZXSocket::GetSockAddr(LPSOCKADDR lpSockAddr, int* lpnSockAddrLen)
{
    if (m_nSocket == INVALID_SOCKET)
        return ZXFAILURE;

    if (getsockname(m_nSocket, (LPSOCKADDR)lpSockAddr, lpnSockAddrLen) == SOCKET_ERROR)
        return ZXFAILURE;

    return ZXSUCCESS;
}
***/

////////////////////////////////////////////////////////////////////////////////

int CZXSocket::ReadBytes(LPSTR lpszBuffer, int nLength)
{
    TIMEVAL     tv = { 1, 0 /*TIMEVAL_USEC*/ };
    FD_SET      fds;
    int         nTimeoutRemain = 0;

    int         nBufferOffset, nBufferRemain, nLen;
    int         nState;

    if (m_nSocket == INVALID_SOCKET)
        return ZXFAILURE;

    nBufferRemain = nLength;
    nBufferOffset = 0;

    while (nBufferRemain > 0) {
        FD_ZERO(&fds);
        FD_SET(m_nSocket, &fds);

        nState = select(m_nSocket+1, &fds, NULL, NULL, &tv);
        if (m_bCancel == TRUE) {
            return ZXFAILURE;
        }
        if (nState == 0) { // Timeout
            if (++nTimeoutRemain == m_nReadTimeout)
                return ZXFAILURE;

            continue;
        }
        else if (nState == SOCKET_ERROR) {
            return ZXFAILURE;
        }

        if (!FD_ISSET(m_nSocket, &fds))
            continue;

        nLen = recv(m_nSocket, lpszBuffer+nBufferOffset, nBufferRemain, 0);
        if (nLen == SOCKET_ERROR) {
            if (WSAGetLastError() == WSAEWOULDBLOCK)
                continue;

            return ZXFAILURE;
        }
        else if (nLen == 0)  {
            break;
        }

        nBufferRemain -= nLen;
        nBufferOffset += nLen;
    }

    return nBufferOffset;
}

int CZXSocket::WriteBytes(LPCSTR lpszBuffer, int nLength)
{
    TIMEVAL     tv = { 1, 0 /*TIMEVAL_USEC*/ };
    FD_SET      fds;
    int         nTimeoutRemain = 0;

    int         nBufferOffset, nBufferRemain, nLen;
    int         nState;

    if (m_nSocket == INVALID_SOCKET)
        return ZXFAILURE;

    nBufferRemain = nLength;
    nBufferOffset = 0;

    while (nBufferRemain > 0) {
        FD_ZERO(&fds);
        FD_SET(m_nSocket, &fds);

        nState = select(m_nSocket+1, NULL, &fds, NULL, &tv);
        if (m_bCancel == TRUE) {
            return ZXFAILURE;
        }
        if (nState == 0) { // Timeout
            if (++nTimeoutRemain == m_nWriteTimeout)
                return ZXFAILURE;

            continue;
        }
        else if (nState == SOCKET_ERROR) {
            return ZXFAILURE;
        }

        if (!FD_ISSET(m_nSocket, &fds))
            continue;

        nLen = send(m_nSocket, lpszBuffer+nBufferOffset, nBufferRemain, 0);
        if (nLen == SOCKET_ERROR) {
            if (WSAGetLastError() == WSAEWOULDBLOCK)
                continue;

            return ZXFAILURE;
        }

        nBufferRemain -= nLen;
        nBufferOffset += nLen;
    }

    return nBufferOffset;
}

////////////////////////////////////////////////////////////////////////////////
int CZXSocket::ReadLine(LPTSTR lpszLine, int nLength)
{
    BOOL        bFindReturn;
    int         nIndex;

    bFindReturn = FALSE;
    nIndex = 0;
    while (1) {
        if (nIndex >= nLength-1)
            nIndex = 0;

        if (ReadBytes(lpszLine+nIndex, 1) <= 0)
            return ZXFAILURE;

        if (!bFindReturn) {
            if (lpszLine[nIndex++] == '\r')
                bFindReturn = TRUE;
        }
        else {
            if (lpszLine[nIndex++] == '\n')
                break;
        }
    }

    if (--nIndex <= 0)
        return ZXFAILURE;

    while (lpszLine[nIndex] == '\r' || lpszLine[nIndex] == '\n') nIndex--;
    lpszLine[nIndex+1] = '\0';

    return nIndex;
}

int CZXSocket::WriteLine(LPCTSTR lpszLine)
{
    int         nLength;

    if ((nLength=_tcslen(lpszLine)) == 0)
        return ZXFAILURE;

    if (_tcscmp(&lpszLine[nLength-2], _T("\r\n"))) {
        LPTSTR      lpLine;

        lpLine = new TCHAR[nLength+3];
        if (lpLine == NULL)
            return ZXFAILURE;

        _stprintf(lpLine, _T("%s\r\n"), lpszLine);
        return WriteBytes(lpLine, _tcslen(lpLine));
    }

    return WriteBytes((LPTSTR)lpszLine, _tcslen(lpszLine));
}

BOOL CZXSocket::HasConnectionDropped( void )
{
    if (m_nSocket == INVALID_SOCKET)
        return FALSE;

	BOOL bConnDropped = FALSE;
	INT iRet = 0;
	BOOL bOK = TRUE;
	
	struct timeval timeout = { 0, 0 };
	fd_set readSocketSet;
	
	FD_ZERO( &readSocketSet );
	FD_SET( m_nSocket, &readSocketSet );
	
	iRet = ::select( 0, &readSocketSet, NULL, NULL, &timeout );
	bOK = ( iRet > 0 );
	
	if( bOK )
	{
		bOK = FD_ISSET( m_nSocket, &readSocketSet );
	}
	
	if( bOK )
	{
		CHAR szBuffer[1] = "";
		iRet = ::recv( m_nSocket, szBuffer, 1, MSG_PEEK );
		bOK = ( iRet > 0 );
		if( !bOK )
		{
			INT iError = ::WSAGetLastError();
			bConnDropped = ( ( iError == WSAENETRESET ) ||
				( iError == WSAECONNABORTED ) ||
				( iError == WSAECONNRESET ) ||
				( iError == WSAEINVAL ) ||
				( iRet == 0 ) ); //Graceful disconnect from other side.
		}
	}
	
    return( bConnDropped );
}