/*
** Copyright 2003-2009, Ernest Laurentin (http://www.ernzo.com/)
**
** Licensed under the Apache License, Version 2.0 (the "License"); 
** you may not use this file except in compliance with the License. 
** You may obtain a copy of the License at 
**
**     http://www.apache.org/licenses/LICENSE-2.0 
**
** Unless required by applicable law or agreed to in writing, software 
** distributed under the License is distributed on an "AS IS" BASIS, 
** WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
** See the License for the specific language governing permissions and 
** limitations under the License.
**
** File:        SocketHandle.cpp
** Version:     1.4 - IPv6 support
**              1.3 - Update for Asynchronous mode / Linux port
**              1.2 - Update interface for TCP remote connection
**              1.1 - Added multicast support
*/
#include <crtdbg.h>
#include "SocketHandle.h"

#ifndef BUFFER_SIZE
#define BUFFER_SIZE     64*1024
#endif
#ifndef SOCKHANDLE_TTL
#define SOCKHANDLE_TTL 5
#endif
#ifndef SOCKHANDLE_HOPS
#define SOCKHANDLE_HOPS 10
#endif
#define HOSTNAME_SIZE   MAX_PATH
#define STRING_LENGTH   40


///////////////////////////////////////////////////////////////////////////////
// CSocketHandle
CSocketHandle::CSocketHandle()
: m_hSocket(INVALID_SOCKET)
{
}

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

///////////////////////////////////////////////////////////////////////////////
// IsOpen
bool CSocketHandle::IsOpen() const 
{
    return ( INVALID_SOCKET != m_hSocket );
}

///////////////////////////////////////////////////////////////////////////////
// GetSocket
SOCKET CSocketHandle::GetSocket() const
{
    return m_hSocket;
}


///////////////////////////////////////////////////////////////////////////////
// Attach
bool CSocketHandle::Attach(SOCKET sock)
{
    if ( INVALID_SOCKET == m_hSocket )
    {
        m_hSocket = sock;
        return true;
    }
    return false;
}

///////////////////////////////////////////////////////////////////////////////
// Detach
SOCKET CSocketHandle::Detach()
{
    SOCKET sock = m_hSocket;
    ::InterlockedExchange(reinterpret_cast<long*>(&m_hSocket), INVALID_SOCKET);
    return sock;
}


///////////////////////////////////////////////////////////////////////////////
// Close
void CSocketHandle::Close()
{
    if ( IsOpen() )
    {
        ShutdownConnection(static_cast<SOCKET>(
            ::InterlockedExchange((LONG*)&m_hSocket, INVALID_SOCKET)
            ));
    }
	m_hSocket = INVALID_SOCKET;
}

///////////////////////////////////////////////////////////////////////////////
// CreateSocket
bool CSocketHandle::CreateSocket(LPCTSTR pszHostName, int port,
                                int nFamily, int nType, UINT uOptions /* = 0 */)
{
    // Socket is already opened
    if ( IsOpen() ) {
        SetLastError(ERROR_ACCESS_DENIED);
        return false;
    }

    // Create a Socket that is bound to a specific service provider
    // nFamily: (AF_INET, AF_INET6)
    // nType: (SOCK_STREAM, SOCK_DGRAM)
#ifdef SOCKHANDLE_USE_OVERLAPPED
    SOCKET sock = WSASocket(nFamily, nType, IPPROTO_IP, NULL, 0, WSA_FLAG_OVERLAPPED);
#else
    SOCKET sock = socket(nFamily, nType, IPPROTO_IP);
#endif
    if (INVALID_SOCKET != sock)
    {
        if (uOptions & SO_REUSEADDR)
        {
            // Inform Windows Sockets provider that a bind on a socket should not be disallowed
            // because the desired address is already in use by another socket
            BOOL optval = TRUE;
            if ( SOCKET_ERROR == setsockopt( sock, SOL_SOCKET, SO_REUSEADDR, (char *) &optval, sizeof( BOOL ) ) )
            {
                SetLastError( WSAGetLastError() );
                closesocket( sock );
                return false;
            }
        }

        if (nType == SOCK_DGRAM)
        {
            if ((uOptions & SO_BROADCAST) && (nFamily == AF_INET))
            {
                // Inform Windows Sockets provider that broadcast messages are allowed
                BOOL optval = TRUE;
                if ( SOCKET_ERROR == setsockopt( sock, SOL_SOCKET, SO_BROADCAST, (char *) &optval, sizeof( BOOL ) ) )
                {
                    SetLastError( WSAGetLastError() );
                    closesocket( sock );
                    return false;
                }
            }
#ifdef SOCKHANDLE_CONFIGBUF
            // configure buffer size
            socklen_t rcvbuf = BUFFER_SIZE;
            if ( SOCKET_ERROR == setsockopt( sock, SOL_SOCKET, SO_RCVBUF, (char *) &rcvbuf, sizeof( int ) ) )
            {
                SetLastError( WSAGetLastError() );
                closesocket( sock );
                return false;
            }
#endif
        }

        // Associate a local address with the socket
		SOCKADDR_IN sockAddr;
		memset((char *)&sockAddr, 0,sizeof(sockAddr));
		sockAddr.sin_family = AF_INET;
		sockAddr.sin_port = htons(port);   
		sockAddr.sin_addr.s_addr = htonl(INADDR_ANY);

        if (SOCKET_ERROR == bind(sock, (struct sockaddr *)&sockAddr, sizeof(sockAddr)))
        {
            SetLastError( WSAGetLastError() );
            closesocket( sock );
            return false;
        }

        // Listen to the socket, only valid for connection socket (TCP)
        if (SOCK_STREAM == nType)
        {
            if ( SOCKET_ERROR == listen(sock, SOMAXCONN))
            {
                SetLastError( WSAGetLastError() );
                closesocket( sock );
                return false;
            }
        }

        // Success, now we may save this socket
        m_hSocket = sock;
    }

    return (INVALID_SOCKET != sock);
}

///////////////////////////////////////////////////////////////////////////////
// ConnectTo
bool CSocketHandle::ConnectTo(LPCTSTR pszRemote,int port, int nFamily, int nType,long timeout)
{
    // Socket is already opened
    if ( IsOpen() ) {
        SetLastError(ERROR_ACCESS_DENIED);
        return false;
    }
	
   SOCKET sock = socket(nFamily, nType, IPPROTO_TCP);
   if (sock < 0) {
	   return false;
   }
   bool rel = false;

        // Associate a local address with the socket but let provider assign a port number
	sockaddr_in  serAddr;
	serAddr.sin_family = AF_INET;
	serAddr.sin_port = htons(port);
	//serAddr.sin_addr.s_addr = inet_addr(pszRemote);
	inet_pton(AF_INET, pszRemote, &serAddr.sin_addr);
	memset(serAddr.sin_zero, 0x00, sizeof(serAddr.sin_zero));

	int error = -1;
	int len = sizeof(int);
	timeval tm;
	fd_set set;
	unsigned long ul = 1;
	ioctlsocket(sock, FIONBIO, &ul);
    // try to connect - if fail, server not ready
    if (SOCKET_ERROR == connect( sock, (sockaddr *)&serAddr, sizeof(serAddr)))
    {
		tm.tv_sec = timeout / 1000.0;
		tm.tv_usec = (timeout % 1000) * 1000;
		FD_ZERO(&set);
		FD_SET(sock, &set);
		if (select(sock + 1, NULL, &set, NULL, &tm) > 0)
		{
			getsockopt(sock, SOL_SOCKET, SO_ERROR, (char *)&error, /*(socklen_t *)*/&len);
			if (error == 0) {
				rel=true;
			}
			else {
				rel=false;
			}
		}
		else {
			rel=false;
		}
	}
	else {
		rel = true;
	}
	
	ul = 0;
	ioctlsocket(sock, FIONBIO, &ul);
	if (!rel) {
		closesocket(sock);
		return false;
	}
	else {
		m_hSocket = sock;
		return true;
	}
}

///////////////////////////////////////////////////////////////////////////////
// Read
int CSocketHandle::Read(LPBYTE lpBuffer, int dwSize, LPSOCKADDR lpAddrIn,
                         long dwTimeout)
{
    _ASSERTE( IsOpen() );
    _ASSERTE( lpBuffer != NULL );

    if (!IsOpen() || lpBuffer == NULL || dwSize < 1L)
        return -1;

    fd_set  fdRead  = { 0 };
    TIMEVAL stTime;
    TIMEVAL *pstTime = NULL;

    if ( INFINITE != dwTimeout ) {
        stTime.tv_sec = dwTimeout/1000;
        stTime.tv_usec = (dwTimeout%1000)*1000;
        pstTime = &stTime;
    }

    SOCKET s = GetSocket();

    // Set Descriptor
    FD_SET( s, &fdRead );

    // Select function set read timeout
    int dwBytesRead = 0;
    int res = 1;
    if ( pstTime != NULL )
        res = select((int)s, &fdRead, NULL, NULL, pstTime );
    if ( res > 0)
    {
        if (lpAddrIn)
        {
            // UDP
            socklen_t fromlen = sizeof(SOCKADDR_STORAGE);
            res = recvfrom(s, reinterpret_cast<LPSTR>(lpBuffer), dwSize, 0, lpAddrIn, &fromlen);
        }
        else
        {
            // TCP
            res = recv(s, reinterpret_cast<LPSTR>(lpBuffer), dwSize, 0);
        }
        if ( res == 0 ) {
            WSASetLastError(WSAECONNRESET);
            res = SOCKET_ERROR;
        }
    }
    if ( res == SOCKET_ERROR )
    {
        SetLastError( WSAGetLastError() );
    }
    dwBytesRead = ((res >= 0)?(res) : (-1));

    return dwBytesRead;
}

#ifdef WIN32
///////////////////////////////////////////////////////////////////////////////
// ReadEx
DWORD CSocketHandle::ReadEx(LPBYTE lpBuffer, DWORD dwSize, LPSOCKADDR lpAddrIn,
                         LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
    _ASSERTE( IsOpen() );
    _ASSERTE( lpBuffer != NULL );

    if (!IsOpen() || lpBuffer == NULL || dwSize < 1L)
        return (DWORD)-1L;

    SOCKET s = GetSocket();
    // Send message to peer
    WSABUF wsabuf;
    wsabuf.buf = (char FAR*)lpBuffer;
    wsabuf.len = dwSize;

    // Select function set read timeout
    DWORD dwBytesRead = 0L;
    DWORD dwFlags = 0L;
    int res = 0;
    if (lpAddrIn)
    {
        // UDP
        socklen_t fromlen = sizeof(SOCKADDR_STORAGE);
        res = WSARecvFrom( s, &wsabuf, 1, &dwBytesRead, &dwFlags, lpAddrIn, &fromlen, lpOverlapped, lpCompletionRoutine);
    }
    else
    {
        // TCP
        res = WSARecv( s, &wsabuf, 1, &dwBytesRead, &dwFlags, lpOverlapped, lpCompletionRoutine);
    }
    if ( res == SOCKET_ERROR )
    {
        res = WSAGetLastError();
        if ( res != WSA_IO_PENDING )
        {
            dwBytesRead = (DWORD)-1L;
            SetLastError( res );
        }
    }

    return dwBytesRead;
}
#endif

///////////////////////////////////////////////////////////////////////////////
// Write
int CSocketHandle::Write(const LPBYTE lpBuffer, int dwCount,
                          const LPSOCKADDR lpAddrIn, long dwTimeout)
{
    _ASSERTE( IsOpen() );
    _ASSERTE( NULL != lpBuffer );

    // validate params
    if (!IsOpen() || NULL == lpBuffer)
        return -1;

    fd_set  fdWrite  = { 0 };
    TIMEVAL stTime;
    TIMEVAL *pstTime = NULL;

    if ( INFINITE != dwTimeout ) {
        stTime.tv_sec = dwTimeout/1000;
        stTime.tv_usec = (dwTimeout%1000)*1000;
        pstTime = &stTime;
    }

    SOCKET s = GetSocket();

    // Set Descriptor
    FD_SET( s, &fdWrite );

    // Select function set write timeout
    int dwBytesWritten = 0;
    int res = 1;
    if ( pstTime != NULL )
        res = select((int)s, NULL, &fdWrite, NULL, pstTime );
    if ( res > 0)
    {
        // Send message to peer
        if (lpAddrIn)
        {
            // UDP
            res = sendto( s, reinterpret_cast<LPCSTR>(lpBuffer), dwCount, 0, lpAddrIn, sizeof(SOCKADDR_STORAGE));
        }
        else
        {
            // TCP
            res = send( s, reinterpret_cast<LPCSTR>(lpBuffer), dwCount, 0);
        }
    }
    if ( res == SOCKET_ERROR )
    {
        SetLastError( WSAGetLastError() );
    }
    dwBytesWritten =((res >= 0)?(res) : (-1));

    return dwBytesWritten;
}

#ifdef WIN32
///////////////////////////////////////////////////////////////////////////////
// WriteEx
DWORD CSocketHandle::WriteEx(const LPBYTE lpBuffer, DWORD dwCount,
                            const LPSOCKADDR lpAddrIn,
                            LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
    _ASSERTE( IsOpen() );
    _ASSERTE( NULL != lpBuffer );

    // validate params
    if (!IsOpen() || NULL == lpBuffer)
        return (DWORD)-1L;

    SOCKET s = GetSocket();

    // Select function set write timeout
    DWORD dwBytesWritten = 0L;
    int res = 0;
    // Send message to peer
    WSABUF wsabuf;
    wsabuf.buf = (char FAR*) lpBuffer;
    wsabuf.len = dwCount;
    if (lpAddrIn)
    {
        // UDP
        res = WSASendTo( s, &wsabuf, 1, &dwBytesWritten, 0, lpAddrIn, sizeof(SOCKADDR_STORAGE),
            lpOverlapped, lpCompletionRoutine);
    }
    else // TCP
        res = WSASend( s, &wsabuf, 1, &dwBytesWritten, 0, lpOverlapped, lpCompletionRoutine);

    if ( res == SOCKET_ERROR )
    {
        res = WSAGetLastError();
        if ( res != WSA_IO_PENDING )
        {
            dwBytesWritten = (DWORD)-1L;
            SetLastError( res );
        }
    }

    return dwBytesWritten;
}
#endif

#ifdef WIN32
///////////////////////////////////////////////////////////////////////////////
// IOControl
bool CSocketHandle::IOControl(DWORD dwIoCode, LPBYTE lpInBuffer, DWORD cbInBuffer,
                              LPBYTE lpOutBuffer, DWORD cbOutBuffer,
                              LPDWORD lpcbBytesReturned, LPWSAOVERLAPPED lpOverlapped,
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
    _ASSERTE( IsOpen() );
    // validate params
    if ( !IsOpen() ) {
        SetLastError(ERROR_INVALID_HANDLE);
        return false;
    }
    int res;
    SOCKET s = GetSocket();
    res = WSAIoctl(s, dwIoCode, lpInBuffer, cbInBuffer, lpOutBuffer, cbOutBuffer,
                   lpcbBytesReturned, lpOverlapped, lpCompletionRoutine);
    if ( res == SOCKET_ERROR )
    {
        SetLastError( WSAGetLastError() );
    }
    return ( res != SOCKET_ERROR );
}

///////////////////////////////////////////////////////////////////////////////
// GetTransferOverlappedResult
bool CSocketHandle::GetTransferOverlappedResult(LPWSAOVERLAPPED lpOverlapped, LPDWORD lpcbTransfer,
                                 bool bWait /*= true*/, LPDWORD lpdwFlags /*= NULL*/)
{
    _ASSERTE( IsOpen() );
    _ASSERTE( NULL != lpOverlapped );

    // validate params
    if (!IsOpen() || NULL == lpOverlapped) {
        SetLastError(ERROR_INVALID_HANDLE);
        return false;
    }

    SOCKET s = GetSocket();
    DWORD dwFlags = 0;
    if ( lpdwFlags == NULL )
        lpdwFlags = &dwFlags;
    BOOL bRet = WSAGetOverlappedResult( s, lpOverlapped, lpcbTransfer, bWait, lpdwFlags );
    if ( !bRet )
    {
        SetLastError( WSAGetLastError() );
    }
    return (bRet != FALSE);
}
#endif


///////////////////////////////////////////////////////////////////////////////
// Utility functions

///////////////////////////////////////////////////////////////////////////////
// InitLibrary
bool CSocketHandle::InitLibrary(WORD wVersion)
{
#ifdef WIN32
    WSADATA WSAData = { 0 };
    return ( 0 == WSAStartup( wVersion, &WSAData ) );
#else
    return true;
#endif
}

///////////////////////////////////////////////////////////////////////////////
// ReleaseLibrary
bool CSocketHandle::ReleaseLibrary()
{
#ifdef WIN32
    return ( 0 == WSACleanup() );
#else
    return true;
#endif
}

///////////////////////////////////////////////////////////////////////////////
// WaitForConnection
SOCKET CSocketHandle::WaitForConnection(SOCKET sock)
{
    return accept(sock, 0, 0);
}

///////////////////////////////////////////////////////////////////////////////
// ShutdownConnection
bool CSocketHandle::ShutdownConnection(SOCKET sock)
{
    shutdown(sock, SD_BOTH);
    return ( 0 == closesocket( sock ));
}