gsocket.cpp

一个非常有用的开源代码

		m_s = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
		if(m_s < 0)
			continue;

		if(connect(m_s, res->ai_addr, res->ai_addrlen) < 0)
		{
			CloseSocket(m_s);
			m_s = INVALID_SOCKET;
//printf("Failed.\n");
			continue;
		}

//printf("OK\n");
		break;  // we got a connection
	}
	freeaddrinfo(res0);
	if(m_s == INVALID_SOCKET)
	{
		// todo: handle the error
		return false;
	}

	// Spawn the listener thread
	m_hListenThread = GThread::SpawnThread(ListenThread, this);
	if(m_hListenThread == BAD_HANDLE)
	{
		GAssert(false, "Failed to spawn listening thread\n");
		gsocket_LogError();
		return false;
	}
	GThread::sleep(0);
	
	return true;
}

void GSocketClientBase::Disconnect()
{
	JoinListenThread();

	// Disconnect the connection
	if(m_s != INVALID_SOCKET)
	{
		shutdown(m_s, 2);
		CloseSocket(m_s);
		m_s = INVALID_SOCKET;
	}
}

bool GSocketClientBase::Send(const unsigned char *pBuff, int len)
{
	if(send(m_s, (const char*)pBuff, len, 0) == SOCKET_ERROR)
	{
#ifdef WIN32		
		int n = WSAGetLastError();
		switch(n)
		{
			case WSAECONNABORTED:	break;
			case WSAECONNRESET:		break;
			default:	gsocket_LogError();		break;
		}
#endif // WIN32
		return false;
	}

	return true;
}

bool GSocketClientBase::IsConnected()
{
	if(m_hListenThread == BAD_HANDLE)
		return false;
	else
		return true;
}

SOCKET GSocketClientBase::GetSocketHandle()
{
	return m_s;
}

in_addr GSocketClientBase::GetMyIPAddr()
{
	struct sockaddr sAddr;
	socklen_t l;
	l = sizeof(SOCKADDR);
	if(getsockname(m_s, &sAddr, &l))
	{
		gsocket_LogError();
	}
	if(sAddr.sa_family != AF_INET)
		GAssert(false, "Error, family is not AF_INET\n");
	SOCKADDR_IN* pInfo = (SOCKADDR_IN*)&sAddr;
	return pInfo->sin_addr;
}

char* GSocketClientBase::GetMyIPAddr(char* szBuff, int nBuffSize)
{
	GString::StrCpy(szBuff, inet_ntoa(GetMyIPAddr()), nBuffSize);
	return szBuff;
}

u_short GSocketClientBase::GetMyPort()
{
	SOCKADDR sAddr;
	socklen_t l;
	l = sizeof(SOCKADDR);
	if(getsockname(m_s, &sAddr, &l))
	{
		gsocket_LogError();
	}
	if(sAddr.sa_family != AF_INET)
		GAssert(false, "Error, family is not AF_INET\n");
	SOCKADDR_IN* pInfo = (SOCKADDR_IN*)&sAddr;
	return htons(pInfo->sin_port);
}

char* GSocketClientBase::GetMyName(char* szBuff, int nBuffSize)
{
	SOCKADDR sAddr;
	socklen_t l;
	l = sizeof(SOCKADDR);
	if(getsockname(m_s, &sAddr, &l))
	{
		gsocket_LogError();
	}
	if(sAddr.sa_family != AF_INET)
		GAssert(false, "Error, family is not AF_INET\n");
	SOCKADDR_IN* pInfo = (SOCKADDR_IN*)&sAddr;
	HOSTENT* namestruct = gethostbyaddr((const char*)&pInfo->sin_addr, 4, pInfo->sin_family);
	if(!namestruct)
	{
		GAssert(false, "Error calling gethostbyaddr\n");
	}
	GString::StrCpy(szBuff, namestruct->h_name, nBuffSize);
	return(szBuff);
}

in_addr GSocketClientBase::GetTheirIPAddr()
{
	struct sockaddr sAddr;
	socklen_t l;
	l = sizeof(SOCKADDR);
//	if(nConnectionNumber == 0)
//	{
		if(getpeername(m_s, &sAddr, &l))
			gsocket_LogError();
//	}
//	else
//	{
//		if(getpeername(m_pHostSockets->GetSocket(nConnectionNumber - 1), &sAddr, &l))
//			gsocket_LogError();
//	}
	if(sAddr.sa_family != AF_INET)
		GAssert(false, "Error, family is not AF_INET\n");
	SOCKADDR_IN* pInfo = (SOCKADDR_IN*)&sAddr;
	return pInfo->sin_addr;
}

char* GSocketClientBase::GetTheirIPAddr(char* szBuff, int nBuffSize)
{
	GString::StrCpy(szBuff, inet_ntoa(GetTheirIPAddr()), nBuffSize);
	return szBuff;
}

u_short GSocketClientBase::GetTheirPort()
{
	SOCKADDR sAddr;
	socklen_t l;
	l = sizeof(SOCKADDR);
//	if(nConnectionNumber == 0)
//	{
		if(getpeername(m_s, &sAddr, &l))
			gsocket_LogError();
//	}
//	else
//	{
//		if(getpeername(m_pHostSockets->GetSocket(nConnectionNumber - 1), &sAddr, &l))
//			gsocket_LogError();
//	}
	if(sAddr.sa_family != AF_INET)
		GAssert(false, "Error, family is not AF_INET\n");
	SOCKADDR_IN* pInfo = (SOCKADDR_IN*)&sAddr;
	return htons(pInfo->sin_port);
}

char* GSocketClientBase::GetTheirName(char* szBuff, int nBuffSize)
{
	SOCKADDR sAddr;
	socklen_t l;
	l = sizeof(SOCKADDR);
//	if(nConnectionNumber == 0)
//	{
		if(getpeername(m_s, &sAddr, &l))
			gsocket_LogError();
//	}
//	else
//	{
//		if(getpeername(m_pHostSockets->GetSocket(nConnectionNumber - 1), &sAddr, &l))
//			gsocket_LogError();
//	}
	if(sAddr.sa_family != AF_INET)
		GAssert(false, "Error, family is not AF_INET\n");
	SOCKADDR_IN* pInfo = (SOCKADDR_IN*)&sAddr;
	HOSTENT* namestruct = gethostbyaddr((const char*)&pInfo->sin_addr, 4, pInfo->sin_family);
	if(!namestruct)
	{
		GAssert(false, "Error calling gethostbyaddr\n");
	}
	GString::StrCpy(szBuff, namestruct->h_name, nBuffSize);
	return(szBuff);
}

// ------------------------------------------------------------------------------

GSocketServerBase::GSocketServerBase(bool bUDP, int nPort, int nMaxConnections)
{
	m_hWorkerThread = BAD_HANDLE;
	m_socketConnectionListener = INVALID_SOCKET;
	m_bKeepWorking = true;
	m_bUDP = false;
	m_szReceiveBuffer = new char[2048];
	m_pConnections = new GSocketArray(16);
	m_pConnectionsLock = new GSpinLock();
	Init(bUDP, nPort, nMaxConnections);
}

GSocketServerBase::~GSocketServerBase()
{
	JoinWorkerThread();

	// Disconnect the connection
	if(m_socketConnectionListener != INVALID_SOCKET)
	{
		shutdown(m_socketConnectionListener, 2);
		CloseSocket(m_socketConnectionListener);
		m_socketConnectionListener = INVALID_SOCKET;
	}

	// Disconnect all the connections
	{
		GSpinLockHolder hLock(m_pConnectionsLock, "~GSocketServerBase");
		int nCount = m_pConnections->GetSize();
		int n;
		SOCKET Sock;
		for(n = 0; n < nCount; n++)
		{
			Sock = m_pConnections->GetSocket(n);
			if(Sock != INVALID_SOCKET)
			{
				shutdown(Sock, 2);
				CloseSocket(Sock);
				if(m_pConnections->GetSize() > n)
					m_pConnections->SetSocket(n, INVALID_SOCKET);
			}
		}

		// Delete the Host Arrays
		delete(m_pConnections);
	}
	delete(m_szReceiveBuffer);
	delete(m_pConnectionsLock);
}

void GSocketServerBase::JoinWorkerThread()
{
	m_bKeepWorking = false;
	time_t tStart;
	time_t tNow;
	time(&tStart);
	while(m_hWorkerThread != BAD_HANDLE)
	{
		GThread::sleep(0);
		time(&tNow);
		if(tNow - tStart > 4)
		{
			GAssert(false, "Error, took too long for the worker thread to exit");
			break;
		}
	}
}

int GSocketServerBase::GetFirstAvailableConnectionNumber()
{
	// Find the first empty Handle slot for the listening thread
	GAssert(m_pConnectionsLock->IsLocked(), "The connections lock should be held when this is called");
	int nSize = m_pConnections->GetSize();
	int nSocketNumber = -1;
	int n;
	for(n = 0; n < nSize; n++)
	{
		if(m_pConnections->GetSocket(n) == INVALID_SOCKET)
		{
			nSocketNumber = n;
			break;
		}
	}

	// Add a new slot if we couldn't find one
	if(nSocketNumber < 0 && m_pConnections->GetSize() < m_nMaxConnections)
	{
		nSocketNumber = nSize;
		m_pConnections->AddSocket(INVALID_SOCKET);
	}

	return nSocketNumber;
}

SOCKET GSocketServerBase::RefreshSocketSet()
{
	// Clear the set
	FD_ZERO(&m_socketSet);

	// Add the connection listener socket so that select() will return if a new connection comes in
	SOCKET highSocket = m_socketConnectionListener;
	FD_SET(m_socketConnectionListener, &m_socketSet);

	// Add all the current connections to the set
	m_pConnectionsLock->Lock("RefreshSocketSet");
	int nCount = m_pConnections->GetSize();
	SOCKET s;
	int n;
	for(n = 0; n < nCount; n++)
	{
		s = m_pConnections->GetSocket(n);
		if(s != INVALID_SOCKET)
		{
			FD_SET(s, &m_socketSet);
			if(s > highSocket)
				highSocket = s;
		}
	}
	m_pConnectionsLock->Unlock();
	return highSocket;
}

void GSocketServerBase::HandleNewConnection()
{
	// Accept the connection
	SOCKET s;
	SOCKADDR_IN sHostAddrIn;
	socklen_t nStructSize = sizeof(struct sockaddr);
	s = accept(m_socketConnectionListener, (struct sockaddr*)&sHostAddrIn, &nStructSize);

	// Set the connection to non-blocking mode
	SetSocketToNonBlockingMode(s);

	// Find a place for the new socket
	int nConnection;
	{
		GSpinLockHolder hLock(m_pConnectionsLock, "HandleNewConnection");
		nConnection = GetFirstAvailableConnectionNumber();
		if(nConnection < 0)
		{
			GAssert(false, "no room for this connection");
	
			// Why accepted the connection even though we didn't have room for it so
			// we could close it so it won't keep bugging us about accepting it.
			CloseSocket(s);
			return;
		}
		m_pConnections->SetSocket(nConnection, s);
	}
	// WARNING: the accept function will return as soon as it gets
	//         an ACK packet back from the client, but the connection
	//         isn't actually established until more data is
	//         received.  Therefore, if you try to send data immediately
	//         (which someone might want to do in OnAcceptConnetion, the
	//         data might be lost since the connection might not be
	//         fully open.
	OnAcceptConnection(nConnection);
}

unsigned int ServerWorkerThread(void* pData)
{
	((GSocketServerBase*)pData)->ServerWorker();
	return 0;
}

void GSocketServerBase::ServerWorker()
{
#ifdef WIN32
	GWindows::YieldToWindows();
#endif // else WIN32
	int n, nCount, nBytes;
	struct timeval timeout;
	int nReadySocketCount; // the number of sockets ready for reading
	SOCKET s, highSocket;
	while(m_bKeepWorking)
	{
		// We need to refresh the socket set each time we loop because select() changes the set
		highSocket = RefreshSocketSet();

		// Check which sockets are ready for reading
		timeout.tv_sec = 1;
		timeout.tv_usec = 0;
		nReadySocketCount = select(highSocket + 1, &m_socketSet, NULL, NULL, &timeout);
		// Handle errors
		if(nReadySocketCount < 0)
		{
			const wchar_t* wszError = gsocket_GetLastError();
			fprintf(stderr, "*** Socket error: %ls\n", wszError);
			break;
		}

		// Read from the ready sockets
		if(nReadySocketCount > 0)
		{
			// Check the connection listener socket for incoming connections
			if(FD_ISSET(m_socketConnectionListener, &m_socketSet))
			{
				HandleNewConnection();
			}

			// Check each connection socket for incoming data
			nCount = m_pConnections->GetSize();
			for(n = 0; n < nCount; n++)
			{
				s = m_pConnections->GetSocket(n);
				if(s != INVALID_SOCKET && FD_ISSET(s, &m_socketSet))
				{
					// The recv() function blocks until there is some to read or connection is closed, but we already know there is something to read
					nBytes = recv(s, m_szReceiveBuffer, 2048, 0);
					if(nBytes > 0)
					{
						Receive((unsigned char*)m_szReceiveBuffer, nBytes, n);
					}
					else
					{
						// The socket was closed or an error occurred. Either way, close the socket
						OnCloseConnection(n);
						CloseSocket(s);
						m_pConnections->SetSocket(n, INVALID_SOCKET);
						ReduceConnectionList();
					}
				}
			}
		}
		else
			GThread::sleep(100);
	}
	m_hWorkerThread = BAD_HANDLE;
}

void GSocketServerBase::Init(bool bUDP, int nPort, int nMaxConnections)
{
	m_nMaxConnections = nMaxConnections;
#ifdef WIN32
	if(!gsocket_InitWinSock())
		throw "failed to init WinSock";
#endif // WIN32

	GAssert(nPort > 0, "invalid port number");
	if(m_bUDP)
	{
		GAssert(false, "UDP not implemented yet");
	}
	m_bUDP = bUDP;

	// Make the Socket
	m_socketConnectionListener = socket(AF_INET, m_bUDP ? SOCK_DGRAM : SOCK_STREAM, 0);
	if(m_socketConnectionListener == INVALID_SOCKET)
	{
		gsocket_LogError();
		throw "faled to make a socket";
	}

	// Put the socket into non-blocking mode (so the call to "accept" will return immediately
	// if there are no connections in the queue ready to be accepted)
	SetSocketToNonBlockingMode(m_socketConnectionListener);

	// Tell the socket that it's okay to reuse an old crashed socket that hasn't timed out yet
	int flag = 1;
	setsockopt(m_socketConnectionListener, SOL_SOCKET, SO_REUSEADDR, (const char*)&flag, sizeof(flag)); 

	// Prepare the socket for accepting
	memset(&m_sHostAddrIn, '\0', sizeof(SOCKADDR_IN));
	m_sHostAddrIn.sin_family = AF_INET;
	m_sHostAddrIn.sin_port = htons((u_short)nPort);
	m_sHostAddrIn.sin_addr.s_addr = htonl(INADDR_ANY);
	if(bind(m_socketConnectionListener, (struct sockaddr*)&m_sHostAddrIn, sizeof(SOCKADDR)))
	{
		gsocket_LogError();
		throw "failed to bind a socket";
	}

	// Start listening for connections
	if(listen(m_socketConnectionListener, nMaxConnections))
	{
		gsocket_LogError();
		throw "Failed to listen on a socket";
	}

	// Spawn the worker thread
	m_hWorkerThread = GThread::SpawnThread(ServerWorkerThread, this);
	if(m_hWorkerThread == BAD_HANDLE)
	{
		GAssert(false, "Error spawning server worker thread");
		throw "Failed to spawn worker thread";
	}

	// Give the worker thread a chance to awake
	GThread::sleep(0);
}

void GSocketServerBase::ReduceConnectionList()
{
	GSpinLockHolder hLock(m_pConnectionsLock, "ReduceConnectionList");
	while(true)
	{
		int n = m_pConnections->GetSize();
		if(n <= 0)
			break;
		if(m_pConnections->GetSocket(n - 1) != INVALID_SOCKET)
			break;
		m_pConnections->DeleteCell(n - 1);
	}
}

void GSocketServerBase::Disconnect(int nConnectionNumber)
{
	m_pConnectionsLock->Lock("Disconnect");
	if(nConnectionNumber < 0 || nConnectionNumber >= m_pConnections->GetSize())
	{
		m_pConnectionsLock->Unlock();
		return;
	}
	SOCKET s = m_pConnections->GetSocket(nConnectionNumber);
	if(s != INVALID_SOCKET)
	{
		OnCloseConnection(nConnectionNumber);
		shutdown(s, 2);
		m_pConnections->SetSocket(nConnectionNumber, INVALID_SOCKET);