////////////////////////////////////////////////////////////////////////////////
//	Copyright (C) 2000-2001 Softelf Inc. All rights reserved.
////////////////////////////////////////////////////////////////////////////////
//
//	Author	: Telan 
//	Date	: 2000-10-04
//	Purpose	: Encapsulate winsock2 functions to make it more easily used 
//	History	: 
//		1.0	: 2000-03-10 - First Edition of this source code ( called:FE_SOCKET )
//		2.0	: 2000-06-25 - Second Edition ( name changed to TE_SOCKET )
//						 - Add Error Control
//		3.0	: 2000-09-21 - Third Edition ( name changed to TE_SOCKET )
//						 - Change the errors' process mechanism
//						 - Add BufSocket Model
//						 - Add TE_ConnectEx(...)
//						 - Add TE_BSocketGetData(...) for specail usage
//		3.1 : 2000-10-04 - Add TE_AcceptEx(...)
//						 - Add TE_GetIP(...) to fix NT DNS resolve cache problem
//						 - Modify TE_ConnectEx
//						 - Fix several bugs in NetEvent process
//
//	Mailto	: telan@263.net ( Bugs' Report or Comments )
//	Notes	: This source code may be used in any form in any way you desire. It is
//			  provided "as is" without express or implied warranty. Use it at your own
//			  risk! The author accepts no liability for any damage/loss of business 
//			  that this product may cause.
//
////////////////////////////////////////////////////////////////////////////////
#include "stdafx.h"
#include "TE_Socket.h"

//需要ws2_32.lib庫(kù)
#pragma comment(lib,"ws2_32.lib")

//錯(cuò)誤代號(hào)
static __declspec(thread) int _iThreadedErrorNo = 0;

// 輔助函數(shù)
//獲取最近一次操作的錯(cuò)誤代碼
int TE_GetLastError()
{
	return (_iThreadedErrorNo);
}

//設(shè)置錯(cuò)誤代碼
void TE_SetLastError(int iErrorCode)
{
	_iThreadedErrorNo = iErrorCode;
}

/////////////////////////////////////////////////////////////////////////
//		Winsock2 函數(shù)
/////////////////////////////////////////////////////////////////////////
//初始化Winsock2動(dòng)態(tài)連接庫(kù)
int TE_InitLibrary()
{
	WSADATA         WSD;
	WORD wVersionRequired = MAKEWORD( TE_SOCKET_MAJOR_VERSION,TE_SOCKET_MINOR_VERSION );
	
	ZeroMemory(&WSD,sizeof(WSADATA));

	int nErrorNo = WSAStartup(wVersionRequired, &WSD);

	if ( SOCKET_SUCCESS != nErrorNo )
	{
		TE_SetLastError( nErrorNo );
		return ( SOCKET_ERROR );
	}

	if ( LOBYTE( WSD.wVersion ) != TE_SOCKET_MINOR_VERSION ||
		 HIBYTE( WSD.wVersion ) != TE_SOCKET_MAJOR_VERSION ) 
	{
		WSACleanup( );
		TE_SetLastError( WSAVERNOTSUPPORTED );
		return (SOCKET_ERROR); 
	}
 
	//成功初始化
	return (SOCKET_SUCCESS);
}


//釋放Winsock2動(dòng)態(tài)連接庫(kù)
void TE_CleanupLibrary(void)
{
    WSACleanup();
}


//設(shè)置套接字屬性
int TE_SetSocketOption(SOCKET hSocket)
{

    int             iActivate = 1;

	//允許地址重用
    if (setsockopt(hSocket, SOL_SOCKET, SO_REUSEADDR, (const char *) &iActivate,sizeof(iActivate)) == SOCKET_ERROR )
    {
        TE_SetLastError( WSAGetLastError() );
        return (SOCKET_ERROR);
    }

	//  如果支持，設(shè)置KEEPALIVE屬性 (這樣做會(huì)帶來其他不良后果)
	//setsockopt(hSocket, SOL_SOCKET, SO_KEEPALIVE, (const char *) &iActivate,sizeof(iActivate));

	return (SOCKET_SUCCESS);
}

//創(chuàng)建具有重疊IO能力的套接字
SOCKET TE_CreateSocket(int iAddressFamily /*= AF_INET*/, int iType/*= SOCK_STREAM*/,int iProtocol/*= 0*/)
{
	SOCKET hSocket = WSASocket(iAddressFamily, iType, iProtocol, NULL,0,WSA_FLAG_OVERLAPPED);
	if ( hSocket == INVALID_SOCKET )
	{
		TE_SetLastError( WSAGetLastError() );
		return (INVALID_SOCKET);
	}

	//設(shè)置套接字選項(xiàng)
	if ( SOCKET_ERROR == TE_SetSocketOption(hSocket) )	//設(shè)置屬性失敗
	{
		TE_CloseSocket(hSocket, TRUE);
		return (INVALID_SOCKET);
	}
	return (hSocket);
}


//關(guān)閉套接字
void TE_CloseSocket(SOCKET hSocket, BOOL bHardClose)
{
	// 不需要捕獲錯(cuò)誤
	if (!bHardClose) // 優(yōu)雅關(guān)閉 Graceful close
	{
		// 不再發(fā)送數(shù)據(jù)，對(duì)于TCP套接字，在所有的數(shù)據(jù)都發(fā)送完畢之后，
		// 將發(fā)送一個(gè) FIN ，通知接收方所有數(shù)據(jù)已經(jīng)發(fā)送完畢。
		shutdown(hSocket, SD_SEND);

		// 接收緩沖區(qū)有可能還有未接收的數(shù)據(jù)，在關(guān)閉套接字之前應(yīng)該先
		// 讀取殘留的數(shù)據(jù)。
		int		iRecvResult;
		HANDLE	hSocketEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
		char	szBuffer[256];
		do
		{
			if (hSocketEvent != NULL)
			{
				//注冊(cè)網(wǎng)絡(luò)事件
				WSAEventSelect(hSocket, (WSAEVENT) hSocketEvent, FD_READ | FD_CLOSE);
				WSAWaitForMultipleEvents(1, &hSocketEvent, TRUE,TE_SHUTDOWN_RECV_TIMEOUT, TRUE);
				//清除網(wǎng)絡(luò)事件
				WSAEventSelect(hSocket, (WSAEVENT) hSocketEvent, 0);
			}
			ZeroMemory(szBuffer,256);
			iRecvResult = TE_RecvLL(hSocket, szBuffer, sizeof(szBuffer));
		} while (iRecvResult > 0);

		if (hSocketEvent != NULL)
			CloseHandle(hSocketEvent);

		//不再允許接收和發(fā)送
		shutdown(hSocket, SD_BOTH);
	}
	
	// 關(guān)閉套接字
	closesocket(hSocket);
}


// 一次性接收數(shù)據(jù)(重疊IO)
int TE_RecvLL(SOCKET hSocket, char *pszBuffer, int iBufferSize)
{
	DWORD		dwRtxBytes = 0,
				dwRtxFlags = 0;
	WSABUF		WSABuff;

	//清空緩沖
	ZeroMemory(&WSABuff,sizeof(WSABUF));

	WSABuff.len = iBufferSize;
	WSABuff.buf = pszBuffer;

	return ((WSARecv(hSocket, &WSABuff, 1, &dwRtxBytes, &dwRtxFlags,NULL, NULL) == SOCKET_SUCCESS) ? (int) dwRtxBytes : -WSAGetLastError());
}

// 接收數(shù)據(jù)(阻塞直至收到數(shù)據(jù)為止)
int TE_RecvData(SOCKET hSocket, char *pszBuffer, int iBufferSize, DWORD dwTimeout)
{
	HANDLE hReadEvent = CreateEvent(NULL, FALSE, FALSE, NULL);

	if (hReadEvent == NULL)
	{
		TE_SetLastError( (int)GetLastError() );
		return ( SOCKET_ERROR );
	}

	int		iRecvBytes = 0;
	DWORD	dwWaitResult ;
	for (;;)
	{
		// 注冊(cè)FD_READ | FD_CLOSE 事件 
		// (因?yàn)榭赡茉诘却鼺D_READ事件中，對(duì)方關(guān)閉套接字，所以要關(guān)注FD_CLOSE)
		if( WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, FD_READ | FD_CLOSE) == SOCKET_ERROR)
		{
			CloseHandle(hReadEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}
	
		// 等等FD_READ | FD_CLOSE事件的發(fā)生
		dwWaitResult = WSAWaitForMultipleEvents(1, &hReadEvent, TRUE,dwTimeout, TRUE);

		if (dwWaitResult != WSA_WAIT_EVENT_0)
		{
			// 清除事件
			WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, 0);
			CloseHandle(hReadEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		
		////////////////////////////////////////////////////////////// 
		///	注意：即使 dwWaitResult == WSA_WAIT_EVENT0 ，也應(yīng)該 
		///			進(jìn)一步檢查網(wǎng)絡(luò)是否發(fā)生錯(cuò)誤
		///////////////////////////////////////////////////////////////
		WSANETWORKEVENTS NetEvent;
		if(WSAEnumNetworkEvents(hSocket,(WSAEVENT)hReadEvent,&NetEvent) == SOCKET_ERROR)
		{
			// 清除事件
			WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, 0);
			CloseHandle(hReadEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		//判斷發(fā)生了什么事件 FD_READ 或 FD_CLOSE	
		if( ( NetEvent.lNetworkEvents == FD_CLOSE ) ||
			 ( NetEvent.lNetworkEvents	== FD_READ && 
			   NetEvent.iErrorCode[FD_READ_BIT] !=0 ) )	// 發(fā)生錯(cuò)誤
		{
			// 清除事件
			WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, 0);
			CloseHandle(hReadEvent);
			TE_SetLastError(WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		////////////////////////////////////////////////////////////////
		// 清除事件
		WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, 0);
	
		// 接收數(shù)據(jù)
		if ((iRecvBytes = TE_RecvLL(hSocket, pszBuffer, iBufferSize)) >= 0)
			break;	// 跳出循環(huán)

		int iErrorCode = -iRecvBytes;

		if ( iErrorCode != WSAEWOULDBLOCK )	//太多的未完成重疊操作
		{
			CloseHandle(hReadEvent);
			TE_SetLastError( iErrorCode );
			return (SOCKET_ERROR);
		}

		//阻塞住了
		////////////////////////////////////////////////////////////////////////
		//  如果發(fā)生阻塞,就等待一定時(shí)間后重試,以免CPU輪詢浪費(fèi)時(shí)間
		////////////////////////////////////////////////////////////////////////
		Sleep(TE_BLOCKED_SNDRCV_SLEEP);
    }

    CloseHandle(hReadEvent);
    return (iRecvBytes);
}


// 接收數(shù)據(jù)直至達(dá)到規(guī)定的長(zhǎng)度(緩沖區(qū)滿)或超時(shí)或沒有數(shù)據(jù)可讀取時(shí)
int TE_Recv(SOCKET hSocket, char *pszBuffer, int iBufferSize, DWORD dwTimeout)
{
	int	iRtxBytes	= 0;
	int iRtxCurrent	= 0;

	while (iRtxBytes < iBufferSize)
	{
		iRtxCurrent = TE_RecvData(hSocket, (pszBuffer + iRtxBytes),(iBufferSize - iRtxBytes), dwTimeout);

		if (iRtxCurrent < 0)	//沒有數(shù)據(jù)可讀取
			return (iRtxBytes);

		iRtxBytes += iRtxCurrent;
	}

	return (iRtxBytes);
}


//一次性發(fā)送數(shù)據(jù)(重疊IO)
int TE_SendLL(SOCKET hSocket, char const * pszBuffer, int iBufferSize)
{
	DWORD	dwRtxBytes = 0;
	WSABUF	WSABuff;

	ZeroMemory(&WSABuff,sizeof(WSABUF));
	WSABuff.len = iBufferSize;
	WSABuff.buf = (char *) pszBuffer;

	return ((WSASend(hSocket, &WSABuff, 1, &dwRtxBytes, 0,NULL, NULL) == SOCKET_SUCCESS) ? (int) dwRtxBytes : -WSAGetLastError());

}

// 發(fā)送數(shù)據(jù)
int TE_SendData(SOCKET hSocket, char const * pszBuffer, int iBufferSize, DWORD dwTimeout)
{

	HANDLE	hWriteEvent = CreateEvent(NULL, FALSE, FALSE, NULL);

	if (hWriteEvent == NULL)
	{
		TE_SetLastError( (int)GetLastError() );
		return (SOCKET_ERROR);
	}

	int iSendBytes = 0;

	for (;;)
	{
		////////////////////////////////////////////////////////////////
		// 發(fā)送數(shù)據(jù)成功
		if ((iSendBytes = TE_SendLL(hSocket, pszBuffer, iBufferSize)) >= 0)
			break;

		int iErrorCode = -iSendBytes;

		if (iErrorCode != WSAEWOULDBLOCK)
		{
			CloseHandle(hWriteEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}

		
		///////////////////////////////////////////////////////////////////////////////
		//  睡眠一段時(shí)間
		///////////////////////////////////////////////////////////////////////////////
		Sleep(TE_BLOCKED_SNDRCV_SLEEP);


		// 注冊(cè)FD_WRITE | FD_CLOSE 事件
		if( WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, FD_WRITE|FD_CLOSE) == SOCKET_ERROR)
		{
			CloseHandle(hWriteEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}

		// 等待事件發(fā)生
		DWORD dwWaitResult = WSAWaitForMultipleEvents(1, &hWriteEvent, TRUE,dwTimeout, TRUE);

		if (dwWaitResult != WSA_WAIT_EVENT_0)
		{
			// 清除網(wǎng)絡(luò)事件
			WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, 0);
			CloseHandle(hWriteEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}

		////////////////////////////////////////////////////////////// 
		///	注意：即使 dwWaitResult == WSA_WAIT_EVENT0 ，也應(yīng)該 
		///			進(jìn)一步檢查網(wǎng)絡(luò)是否發(fā)生錯(cuò)誤
		///////////////////////////////////////////////////////////////
		WSANETWORKEVENTS NetEvent;
		if(WSAEnumNetworkEvents(hSocket,(WSAEVENT)hWriteEvent,&NetEvent) == SOCKET_ERROR)
		{
  			// 清除網(wǎng)絡(luò)事件
			WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, 0);
			CloseHandle(hWriteEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		if( ( NetEvent.lNetworkEvents == FD_CLOSE ) ||
			 ( NetEvent.lNetworkEvents == FD_WRITE   &&
				NetEvent.iErrorCode[FD_WRITE_BIT] !=0 ) )	// 發(fā)生錯(cuò)誤
		{
			// 清除網(wǎng)絡(luò)事件
			WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, 0);
			CloseHandle(hWriteEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		// 清除網(wǎng)絡(luò)事件
		WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, 0);
   }

    CloseHandle(hWriteEvent);
    return (iSendBytes);
}

//發(fā)送完所有數(shù)據(jù)或超時(shí)
int TE_Send(SOCKET hSocket, char const * pszBuffer, int iBufferSize, DWORD dwTimeout)
{

	int iRtxBytes	= 0;
	int iRtxCurrent = 0;

	while (iRtxBytes < iBufferSize)
	{
		iRtxCurrent = TE_SendData(hSocket, (pszBuffer + iRtxBytes),(iBufferSize - iRtxBytes), dwTimeout);
		if (iRtxCurrent < 0)
			return (iRtxBytes);

		iRtxBytes += iRtxCurrent;
	}
	return (iRtxBytes);
}

// 建立連接
int TE_Connect(SOCKET hSocket, const struct sockaddr * pSocketAddress, int iNameLen,DWORD dwTimeout)
{

	HANDLE hConnectEvent = CreateEvent(NULL, FALSE, FALSE, NULL);

	if (hConnectEvent == NULL)
	{
		TE_SetLastError( (int)GetLastError() );
		return (SOCKET_ERROR);