te_socket.cpp

基于Winsock2的支持断点续传和SOCKS代理的HTTP下载类 CHttpDownload

////////////////////////////////////////////////////////////////////////////////
//	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	: [e-mail] ( 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库
#pragma comment(lib,"ws2_32.lib")

//错误代号
static __declspec(thread) int _iThreadedErrorNo = 0;

// 辅助函数
//获取最近一次操作的错误代码
int TE_GetLastError()
{
	return (_iThreadedErrorNo);
}

//设置错误代码
void TE_SetLastError(int iErrorCode)
{
	_iThreadedErrorNo = iErrorCode;
}

/////////////////////////////////////////////////////////////////////////
//		Winsock2 函数
/////////////////////////////////////////////////////////////////////////
//初始化Winsock2动态连接库
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动态连接库
void TE_CleanupLibrary(void)
{
    WSACleanup();
}


//设置套接字属性
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);
    }

	//  如果支持，设置KEEPALIVE属性 (这样做会带来其他不良后果)
	//setsockopt(hSocket, SOL_SOCKET, SO_KEEPALIVE, (const char *) &iActivate,sizeof(iActivate));

	return (SOCKET_SUCCESS);
}

//创建具有重叠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);
	}

	//设置套接字选项
	if ( SOCKET_ERROR == TE_SetSocketOption(hSocket) )	//设置属性失败
	{
		TE_CloseSocket(hSocket, TRUE);
		return (INVALID_SOCKET);
	}
	return (hSocket);
}


//关闭套接字
void TE_CloseSocket(SOCKET hSocket, BOOL bHardClose)
{
	// 不需要捕获错误
	if (!bHardClose) // 优雅关闭 Graceful close
	{
		// 不再发送数据，对于TCP套接字，在所有的数据都发送完毕之后，
		// 将发送一个 FIN ，通知接收方所有数据已经发送完毕。
		shutdown(hSocket, SD_SEND);

		// 接收缓冲区有可能还有未接收的数据，在关闭套接字之前应该先
		// 读取残留的数据。
		int		iRecvResult;
		HANDLE	hSocketEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
		char	szBuffer[256];
		do
		{
			if (hSocketEvent != NULL)
			{
				//注册网络事件
				WSAEventSelect(hSocket, (WSAEVENT) hSocketEvent, FD_READ | FD_CLOSE);
				WSAWaitForMultipleEvents(1, &hSocketEvent, TRUE,TE_SHUTDOWN_RECV_TIMEOUT, TRUE);
				//清除网络事件
				WSAEventSelect(hSocket, (WSAEVENT) hSocketEvent, 0);
			}
			ZeroMemory(szBuffer,256);
			iRecvResult = TE_RecvLL(hSocket, szBuffer, sizeof(szBuffer));
		} while (iRecvResult > 0);

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

		//不再允许接收和发送
		shutdown(hSocket, SD_BOTH);
	}
	
	// 关闭套接字
	closesocket(hSocket);
}


// 一次性接收数据(重叠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());
}

// 接收数据(阻塞直至收到数据为止)
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 (;;)
	{
		// 注册FD_READ | FD_CLOSE 事件 
		// (因为可能在等待FD_READ事件中，对方关闭套接字，所以要关注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事件的发生
		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 ，也应该 
		///			进一步检查网络是否发生错误
		///////////////////////////////////////////////////////////////
		WSANETWORKEVENTS NetEvent;
		if(WSAEnumNetworkEvents(hSocket,(WSAEVENT)hReadEvent,&NetEvent) == SOCKET_ERROR)
		{
			// 清除事件
			WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, 0);
			CloseHandle(hReadEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		//判断发生了什么事件 FD_READ 或 FD_CLOSE	
		if( ( NetEvent.lNetworkEvents == FD_CLOSE ) ||
			 ( NetEvent.lNetworkEvents	== FD_READ && 
			   NetEvent.iErrorCode[FD_READ_BIT] !=0 ) )	// 发生错误
		{
			// 清除事件
			WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, 0);
			CloseHandle(hReadEvent);
			TE_SetLastError(WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		////////////////////////////////////////////////////////////////
		// 清除事件
		WSAEventSelect(hSocket, (WSAEVENT) hReadEvent, 0);
	
		// 接收数据
		if ((iRecvBytes = TE_RecvLL(hSocket, pszBuffer, iBufferSize)) >= 0)
			break;	// 跳出循环

		int iErrorCode = -iRecvBytes;

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

		//阻塞住了
		////////////////////////////////////////////////////////////////////////
		//  如果发生阻塞,就等待一定时间后重试,以免CPU轮询浪费时间
		////////////////////////////////////////////////////////////////////////
		Sleep(TE_BLOCKED_SNDRCV_SLEEP);
    }

    CloseHandle(hReadEvent);
    return (iRecvBytes);
}


// 接收数据直至达到规定的长度(缓冲区满)或超时或没有数据可读取时
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)	//没有数据可读取
			return (iRtxBytes);

		iRtxBytes += iRtxCurrent;
	}

	return (iRtxBytes);
}


//一次性发送数据(重叠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());

}

// 发送数据
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 (;;)
	{
		////////////////////////////////////////////////////////////////
		// 发送数据成功
		if ((iSendBytes = TE_SendLL(hSocket, pszBuffer, iBufferSize)) >= 0)
			break;

		int iErrorCode = -iSendBytes;

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

		
		///////////////////////////////////////////////////////////////////////////////
		//  睡眠一段时间
		///////////////////////////////////////////////////////////////////////////////
		Sleep(TE_BLOCKED_SNDRCV_SLEEP);


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

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

		if (dwWaitResult != WSA_WAIT_EVENT_0)
		{
			// 清除网络事件
			WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, 0);
			CloseHandle(hWriteEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}

		////////////////////////////////////////////////////////////// 
		///	注意：即使 dwWaitResult == WSA_WAIT_EVENT0 ，也应该 
		///			进一步检查网络是否发生错误
		///////////////////////////////////////////////////////////////
		WSANETWORKEVENTS NetEvent;
		if(WSAEnumNetworkEvents(hSocket,(WSAEVENT)hWriteEvent,&NetEvent) == SOCKET_ERROR)
		{
  			// 清除网络事件
			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 ) )	// 发生错误
		{
			// 清除网络事件
			WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, 0);
			CloseHandle(hWriteEvent);
			TE_SetLastError( WSAGetLastError() );
			return (SOCKET_ERROR);
		}
		// 清除网络事件
		WSAEventSelect(hSocket, (WSAEVENT) hWriteEvent, 0);
   }

    CloseHandle(hWriteEvent);
    return (iSendBytes);
}

//发送完所有数据或超时
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);