#include "StdAfx.h"
#include "FileTransfer.h"

CFileTransfer::CFileTransfer(void)
{
	m_hSocket = INVALID_SOCKET;
}

CFileTransfer::~CFileTransfer(void)
{
	Close();
}

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

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

int CFileTransfer::GetSocketType() const
{
	int type = -1;
	if ( INVALID_SOCKET != m_hSocket ) 
	{
		socklen_t optlen = sizeof(int);
		if ( getsockopt(GetSocket(), SOL_SOCKET, SO_TYPE, reinterpret_cast<char*>(&type),&optlen) == SOCKET_ERROR)
		{
			SetLastError( WSAGetLastError() );
		}
	}
	return type;
}

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

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

bool CFileTransfer::GetSockName(SockAddrIn& saddr_in) const
{
	_ASSERTE( IsOpen() );
	if (IsOpen()) 
	{
		socklen_t namelen = (socklen_t)saddr_in.Size();
		if (SOCKET_ERROR != getsockname(GetSocket(), saddr_in, &namelen))
		{
			return true;
		}
		SetLastError( WSAGetLastError() );
	}
	return false;
}

bool CFileTransfer::GetPeerName(SockAddrIn& saddr_in) const
{
	_ASSERTE( IsOpen() );
	if (IsOpen()) 
	{
		socklen_t namelen = (socklen_t)saddr_in.Size();
		if (SOCKET_ERROR != getpeername(GetSocket(), saddr_in, &namelen))
		{
			return true;
		}
		SetLastError( WSAGetLastError() );
	}
	return false;
}

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

	SOCKET sock = socket(nFamily, nType, IPPROTO_IP);
	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
		SockAddrIn sockAddr;
		sockAddr.CreateFrom(pszHostName, pszServiceName, nFamily);

		if ( SOCKET_ERROR == bind(sock, sockAddr, (int)sockAddr.Size()))
		{
			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);
}

bool CFileTransfer::ConnectTo(LPCTSTR pszHostName, LPCTSTR pszRemote,LPCTSTR pszServiceName, int nFamily, int nType)
{
	// Socket is already opened
	if ( IsOpen() ) {
		SetLastError(ERROR_ACCESS_DENIED);
		return false;
	}

	SOCKET sock = socket(nFamily, nType, IPPROTO_IP);
	if (INVALID_SOCKET != sock)
	{
		// Associate a local address with the socket but let provider assign a port number
		SockAddrIn sockAddr;
		if (false == sockAddr.CreateFrom(pszHostName, TEXT("0"), nFamily))
		{
			SetLastError( WSAGetLastError() );
			closesocket( sock );
			return false;
		}

		if ( SOCKET_ERROR == bind(sock, sockAddr, (int)sockAddr.Size()))
		{
			SetLastError( WSAGetLastError() );
			closesocket( sock );
			return false;
		}

#ifdef SOCKHANDLE_CONFIGBUF
		if (nType == SOCK_DGRAM)
		{
			// 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
		// Now get destination address & port
		sockAddr.CreateFrom( pszRemote, pszServiceName, nFamily );

		// try to connect - if fail, server not ready
		if (SOCKET_ERROR == connect( sock, sockAddr, (int)sockAddr.Size()))
		{
			SetLastError( WSAGetLastError() );
			closesocket( sock );
			return false;
		}

		// Success, now we may save this socket
		m_hSocket = sock;
	}
	return (INVALID_SOCKET != sock);
}

DWORD CFileTransfer::Read(LPBYTE lpBuffer, DWORD dwSize, LPSOCKADDR lpAddrIn,DWORD dwTimeout)
{
	_ASSERTE( IsOpen() );
	_ASSERTE( lpBuffer != NULL );

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

	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
	DWORD dwBytesRead = 0L;
	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 = (DWORD)((res >= 0)?(res) : (-1));

	return dwBytesRead;
}

DWORD CFileTransfer::Write(const LPBYTE lpBuffer, DWORD dwCount,const LPSOCKADDR lpAddrIn, DWORD dwTimeout)
{
	_ASSERTE( IsOpen() );
	_ASSERTE( NULL != lpBuffer );

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

	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
	DWORD dwBytesWritten = 0L;
	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 = (DWORD)((res >= 0)?(res) : (-1));

	return dwBytesWritten;
}

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

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

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

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