lyfz_repos/经典版源码/kernelcode/NetWork/IocpServer/TCPSrv.cpp

#include "stdafx.h"
#include "TCPSrv.h"

// 加载SOCK DLL
bool MyTcpSrv::LoadSock()
{
	WSADATA wsaData;
	return  WSAStartup(MAKEWORD(2, 2), &wsaData) == 0 ? true : false;
}

// 卸载SOCK DLL
void MyTcpSrv::UnloadSock()
{
	WSACleanup();
}

// 设置本地终端
bool MyTcpSrv::SetLocalAddr(WORD nPort, const char *szAddr)
{
	memset(&m_LocalAddr, 0, sizeof(SOCKADDR_IN));
	m_LocalAddr.sin_family = AF_INET;
	m_LocalAddr.sin_port = htons(nPort);
	if(szAddr == NULL)
		m_LocalAddr.sin_addr.s_addr = htonl(INADDR_ANY);
	else
	{
		if((m_LocalAddr.sin_addr.s_addr = inet_addr(szAddr)) == INADDR_NONE)
		{
			HOSTENT *host = NULL;
			host = gethostbyname(szAddr);
			if(host == NULL)
				return false;
			CopyMemory(&m_LocalAddr.sin_addr, host->h_addr, host->h_length);
		}
	}
	return true;
}

// 初始化完成端口
bool MyTcpSrv::InitialIOCP()
{
	// 创建完成端口，其中第三个参数为零，表示完成端口最大同时支持CPU个线程
	if ((m_hIOCP = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0)) == NULL)
		return false;

	// 获取系统信息，计算所需工作线程数，创建工作线程
	SYSTEM_INFO sysInfo;
	GetSystemInfo(&sysInfo);
	UINT ThreadID;
	for(USHORT i = 0; i < sysInfo.dwNumberOfProcessors * 2 + 2; i++)
	{
		HANDLE hThread;
		if((hThread = (HANDLE)_beginthreadex(NULL, 0, ServerWorkerThread, this, 0, &ThreadID)) == NULL)
			return false;//创建工作线程失败			
		CloseHandle(hThread);
	}
	return true;
}

// 设置侦听套接字
bool MyTcpSrv::SetupListenSocket()
{
	// 创建侦听套接字
	if ((m_sLinstenSocket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED)) == INVALID_SOCKET)
		return false; // 创建侦听套接字失败

	// 绑定端口
	if (bind(m_sLinstenSocket, (PSOCKADDR) &m_LocalAddr, sizeof(SOCKADDR_IN)) == SOCKET_ERROR)
		return false; // 绑定端口失败

	// 进入侦听被动模式，准备接收
	if (listen(m_sLinstenSocket, SOMAXCONN) == SOCKET_ERROR)
		return false; // 侦听失败

	return true;
}

// 启动侦听线程
bool MyTcpSrv::RunListenThread()
{
	// 初始化信号
	m_hListenEventExit = CreateEvent(NULL, TRUE, FALSE, NULL);
	if((m_hListenThread = (HANDLE)_beginthreadex(NULL, 0, TCPListenThreadProc, this, 0, &m_nListenThreadId)) == NULL)
	{
		closesocket(m_sLinstenSocket);
		return false; // 启动侦听线程失败
	}
	return true;
}

// 侦听执行体
UINT WINAPI MyTcpSrv::TCPListenThreadProc(LPVOID pParam)
{	
	// 接收套接字
	TS_CLIENT_ELEMENT		AcceptElement;
	WSAEVENT				wsaListenEvent;
	OVERLAPPED				ListenOverlapped;
	CHAR					szAcceptBuffer [2 * (sizeof(SOCKADDR_IN) + 16)];
	DWORD					dwBytesReceived;
	LPSOCKADDR				lpServerSockInfo;
	LPSOCKADDR				lpClientSockInfo;
	int						nServerAddrLen;
	int						nClientAddrLen;
	LPSOCKET_INFORMATION	lpSocketInfo;
	MyTcpSrv*				pthis;

	pthis = (MyTcpSrv*)pParam;
	wsaListenEvent = WSACreateEvent();

	// 不断接收连接接入
	while(true)
	{
		ZeroMemory(&AcceptElement, sizeof(TS_CLIENT_ELEMENT));
		if ((AcceptElement.s_socket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED)) == INVALID_SOCKET) 
			return 0L;//创建套接字失败

		ZeroMemory(&ListenOverlapped, sizeof(WSAOVERLAPPED));
		WSAResetEvent(wsaListenEvent);
		ListenOverlapped.hEvent = wsaListenEvent;

		// 异步接收套接字连接
		if (!AcceptEx(pthis->m_sLinstenSocket, AcceptElement.s_socket, (PVOID)szAcceptBuffer, 0, sizeof(SOCKADDR_IN) + 16, sizeof(SOCKADDR_IN) + 16, &dwBytesReceived, &ListenOverlapped))
		{
			if (WSAGetLastError() != ERROR_IO_PENDING)
				return 0L;//接受套接字失败
		}

		// 等待AcceptEx返回
		while(TRUE)
		{
			HANDLE hEvents[2];
			hEvents[0]  = ListenOverlapped.hEvent;
			hEvents[1]  = pthis->m_hListenEventExit;

			DWORD dwEventIndex = WSAWaitForMultipleEvents(2, hEvents, FALSE, WSA_INFINITE, TRUE);
			if (dwEventIndex == WSA_WAIT_FAILED)
				return 0L; // 异常
			if ((dwEventIndex-WAIT_OBJECT_0) == 1)
				return 0L; // 接到退出线程通知

			if (dwEventIndex != WAIT_IO_COMPLETION)
				break;
		}

		// 获取客户端详细信息，并登记
		GetAcceptExSockaddrs(szAcceptBuffer, 0, sizeof(SOCKADDR_IN) + 16,  sizeof(SOCKADDR_IN) + 16, &lpServerSockInfo, &nServerAddrLen, &lpClientSockInfo, &nClientAddrLen);

		memcpy(&AcceptElement.s_clientInfo, lpClientSockInfo, sizeof(sockaddr_in));
		time(&AcceptElement.s_tVisitTM);// 获取访问时间
		AcceptElement.s_bFirstFrame = TRUE;
		pthis->m_clntList.push_back(AcceptElement);

		lpSocketInfo = new SOCKET_INFORMATION;// 创建通信载体
		lpSocketInfo->s_lpClientElement = &(pthis->m_clntList.back());
		lpSocketInfo->s_dataBuf.len = SOCKET_BUF_SIZE;
		lpSocketInfo->s_dataBuf.buf = lpSocketInfo->s_buffer;
		lpSocketInfo->s_nBufLen = 0;
		lpSocketInfo->s_pThis = pthis;
		

		// 将获取到的SOCKET绑定到完成端口
		if (CreateIoCompletionPort((HANDLE)lpSocketInfo->s_lpClientElement->s_socket, lpSocketInfo->s_pThis->m_hIOCP, (DWORD)lpSocketInfo, 0) == NULL)
			return 0L; // 绑定完成端口失败

		// 递交异步接收请求
		lpSocketInfo->s_pThis->DeliverRecv(lpSocketInfo);
	}
	return 1L;
}

// 递交异步接收请求
void MyTcpSrv::DeliverRecv(LPSOCKET_INFORMATION lpSocketInfo)
{
	DWORD		dwRecvBytes = 0;
	DWORD		dwFlags = 0;

	// 整理通信载体
	ZeroMemory(&(lpSocketInfo->s_overlapped), sizeof(WSAOVERLAPPED));
	lpSocketInfo->s_dataBuf.len = SOCKET_BUF_SIZE;
	lpSocketInfo->s_dataBuf.buf = lpSocketInfo->s_buffer;
	lpSocketInfo->s_nBufLen = 0;
	lpSocketInfo->s_pThis = this;
	lpSocketInfo->s_uDeliverType = OVERLAPPED_DELIVER_RECV;

	if (WSARecv(lpSocketInfo->s_lpClientElement->s_socket, &(lpSocketInfo->s_dataBuf), 1, &dwRecvBytes, \
		&dwFlags, &(lpSocketInfo->s_overlapped), NULL) == SOCKET_ERROR)
	{
		// 异步投递失败
		if (WSAGetLastError() != WSA_IO_PENDING)
			Erase(lpSocketInfo);
	}
}

// 递交异步发送请求
void MyTcpSrv::DeliverSend(LPSOCKET_INFORMATION lpSocketInfo)
{
	DWORD dwSendBytes = 0;
	DWORD dwFlags = 0;

	// 已经全部发送完毕
	if (lpSocketInfo->s_uAntiSendSize - lpSocketInfo->s_nBufLen <= 0)
	{
		delete lpSocketInfo;
		return;
	}

	// 整理通信载体
	ZeroMemory(&(lpSocketInfo->s_overlapped), sizeof(WSAOVERLAPPED));
	lpSocketInfo->s_dataBuf.len = lpSocketInfo->s_uAntiSendSize - lpSocketInfo->s_nBufLen;
	lpSocketInfo->s_dataBuf.buf = lpSocketInfo->s_buffer + lpSocketInfo->s_nBufLen;
	lpSocketInfo->s_pThis = this;
	lpSocketInfo->s_uDeliverType = OVERLAPPED_DELIVER_SEND;	// 发送
	
	// 异步发送请求
	if(WSASend(lpSocketInfo->s_lpClientElement->s_socket, &lpSocketInfo->s_dataBuf, 1, &dwSendBytes, dwFlags, &lpSocketInfo->s_overlapped, NULL) == SOCKET_ERROR)
	{
		// 异步发送请求失败
		if (WSAGetLastError() != WSA_IO_PENDING)
		{
			delete lpSocketInfo;// 释放通信载体
			return; 
		}		
	}
}

// 擦除记录
void MyTcpSrv::Erase(LPSOCKET_INFORMATION lpSocketInfo)
{
	// 关闭连接
	closesocket(lpSocketInfo->s_lpClientElement->s_socket);
	
	// 清除记录
	EnterCriticalSection(&m_csLock);
	VClnt::iterator ite;
	for(ite=m_clntList.begin();ite != m_clntList.end();++ite)
	{
		if (ite->s_socket == lpSocketInfo->s_lpClientElement->s_socket)
		{
			m_clntList.erase(ite);
			break;
		}
	}
	LeaveCriticalSection(&m_csLock);

	delete lpSocketInfo; // 释放通信载体
}

// 工作例程
UINT WINAPI MyTcpSrv::ServerWorkerThread(LPVOID lpMyTcpSrv)
{
	MyTcpSrv* pThis = (MyTcpSrv*)lpMyTcpSrv;
	DWORD BytesTransferred;
	DWORD dwInfo;
	LPSOCKET_INFORMATION lpSocketInfo;		// 通信载体
	//LPTS_CLIENT_ELEMENT lpClientInfo;
	//lpSocketInfo = new SOCKET_INFORMATION;
	while(TRUE)
	{
		BOOL bRes = GetQueuedCompletionStatus(pThis->m_hIOCP, &BytesTransferred, &dwInfo, (LPOVERLAPPED*)(&lpSocketInfo), INFINITE);

		switch(lpSocketInfo->s_uDeliverType)
		{
		case OVERLAPPED_DELIVER_SEND:
			// 完成端口异常或客户端关闭则不再尝试发送数据
			if (!bRes || BytesTransferred == 0)
			{
				delete lpSocketInfo;	// 释放通信载体
				return 1L;
			}
			// 解析获取的数据
			lpSocketInfo->s_nBufLen += BytesTransferred;	// 计算已发送总字节数
			// 再次投递
			lpSocketInfo->s_pThis->DeliverSend(lpSocketInfo);
			break;
		case OVERLAPPED_DELIVER_RECV:
			// 完成端口异常，退出线程
			if (!bRes)
			{
				pThis->Erase(lpSocketInfo);
				return 1L;
			}
			// 客户端退出
			if (BytesTransferred == 0)
			{
				pThis->Erase(lpSocketInfo);
				continue;
			}
			lpSocketInfo->s_nBufLen = BytesTransferred;		// 实际接收字节数
			// 解析获取的数据
			lpSocketInfo->s_pThis->OnRecvData(lpSocketInfo);
			// 再次投递
			lpSocketInfo->s_pThis->DeliverRecv(lpSocketInfo);
			break;
		default:
			break;
		}
	}
}

// 获取数据完成处理
void MyTcpSrv::OnRecvData(LPSOCKET_INFORMATION lpSocketInfo)
{
	SNetPara pNet;
	memset(pNet.s_szBuf, 0, SOCKET_BUF_SIZE);
	TERMINAL_INFO tInfo;

	tInfo.s_ip = inet_ntoa(lpSocketInfo->s_lpClientElement->s_clientInfo.sin_addr);
	tInfo.s_uPort = ntohs(lpSocketInfo->s_lpClientElement->s_clientInfo.sin_port);

	memcpy(pNet.s_szBuf, lpSocketInfo->s_buffer, min(SOCKET_BUF_SIZE, lpSocketInfo->s_nBufLen));
	pNet.s_wBufLen = lpSocketInfo->s_nBufLen;

	EnterCriticalSection(&m_csLock);
	std::pair<MNetDatagram::iterator, bool> res = m_netParaMap.insert(std::make_pair(tInfo, pNet));
	if (!res.second)
		m_netParaMap[tInfo].Append(&pNet);
	LeaveCriticalSection(&m_csLock);
}

// 停止侦听
void MyTcpSrv::StopListening()
{
	SetEvent(m_hListenEventExit);
	WaitForSingleObject(m_hListenThread, INFINITE);	// 等待侦听线程退出
	DisconnectAllUsers();
	CloseHandle(m_hIOCP);
	m_hIOCP = INVALID_HANDLE_VALUE;
	closesocket(m_sLinstenSocket);
	m_sLinstenSocket = INVALID_SOCKET;
	CloseHandle(m_hListenThread);
	m_hListenThread = INVALID_HANDLE_VALUE;
	CloseHandle(m_hListenEventExit);
	m_hListenEventExit = INVALID_HANDLE_VALUE;
}

// 断开与指定用户的连接
void MyTcpSrv::DisconnectUser(SOCKADDR_IN clientAddr)
{
	SOCKET clientSock = INVALID_SOCKET;

	// 清除记录
	EnterCriticalSection(&m_csLock);
	VClnt::iterator ite = m_clntList.begin();
	for(; ite!=m_clntList.end(); ++ite)
	{
		if (ite->s_clientInfo.sin_addr.s_addr== clientAddr.sin_addr.s_addr && ite->s_clientInfo.sin_port == clientAddr.sin_port)
		{
			m_clntList.erase(ite);
			break;
		}
	}
	LeaveCriticalSection(&m_csLock);

	if(clientSock == INVALID_SOCKET)
		return;

	// 关闭套接字
	shutdown(clientSock, SD_BOTH);
	closesocket(clientSock);	
}

// 断开与所有用户的连接
void MyTcpSrv::DisconnectAllUsers()
{
	while (!m_clntList.empty())
		DisconnectUser(m_clntList.front().s_clientInfo);
}

// 创建通道
bool MyTcpSrv::Create(LPTERMINAL_INFO pInfo)
{
	return LoadSock() && SetLocalAddr(pInfo->s_uPort,pInfo->s_ip.c_str()) && InitialIOCP() && SetupListenSocket() && RunListenThread();
}

// 读取数据
int MyTcpSrv::Read(byte byRecv[], UINT uMaxSize, LPTERMINAL_INFO pInfo)
{
	UINT uSize = 0;
	EnterCriticalSection(&m_csLock);
	SNetPara& netPara = m_netParaMap[*pInfo];
	uSize = min(netPara.s_wBufLen, uMaxSize);
	memcpy(byRecv, netPara.s_szBuf , uSize);
	netPara.Release(uSize);
	LeaveCriticalSection(&m_csLock);
	return uSize;
}

// 写入数据
int MyTcpSrv::Write(byte bySend[], UINT uAntiSize, LPTERMINAL_INFO pInfo)
{
	// 创建通信载体
	LPSOCKET_INFORMATION lpSocketInfo = new SOCKET_INFORMATION;// 创建通信载体
	// 查找客户端信息
	EnterCriticalSection(&m_csLock);
	VClnt::iterator ite = m_clntList.begin();
	for (;ite!=m_clntList.end(); ++ite)
	{
		if (inet_ntoa(ite->s_clientInfo.sin_addr) == pInfo->s_ip &&  ntohs(ite->s_clientInfo.sin_port) == pInfo->s_uPort)
		{		
			lpSocketInfo->s_lpClientElement = &(*ite);
			break;
		}
	}
	LeaveCriticalSection(&m_csLock);
	
	if (lpSocketInfo->s_lpClientElement == NULL)
		return 0; // 不存在指定客户端

	UINT uSize = min(SOCKET_BUF_SIZE, uAntiSize);
	memcpy(lpSocketInfo->s_buffer, bySend, uSize);
	lpSocketInfo->s_nBufLen = 0;	
	lpSocketInfo->s_pThis = this;
	lpSocketInfo->s_uAntiSendSize = uSize;

	// 递交发送请求
	DeliverSend(lpSocketInfo);
	return uSize;
}

// 重置服务器
bool MyTcpSrv::Reset()
{
	StopListening();
	return InitialIOCP() && SetupListenSocket() && RunListenThread();
}

// 关闭服务器
void MyTcpSrv::Close()
{
	StopListening();
}