mysocket.cpp

用于检测自行开发的tcp/ip协议栈的小程序

// MySocket.cpp: implementation of the CMySocket class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "MySocket.h"

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

CMySocket::CMySocket(HWND hWnd,UINT uMsg)
{
	m_Socket=INVALID_SOCKET;
	m_hWnd=hWnd;
	m_uMsg=uMsg;
}

CMySocket::~CMySocket()
{
	Close();
}

BOOL CMySocket::InitSocket()
{
	if(m_Socket!=INVALID_SOCKET)
		closesocket(m_Socket);

	m_Socket=socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);
	if(m_Socket==INVALID_SOCKET)
		return FALSE;

	return TRUE;
}

BOOL CMySocket::Connect(CString strServerIP, u_short uPort)
{
	int nIPLen=strServerIP.GetLength();
	if(nIPLen<=0)
		return FALSE;

	if(!InitSocket())
		return FALSE;

	char *szServer=new char[nIPLen+1];
	memset(szServer,0,nIPLen+1);
	struct sockaddr_in server;
	struct hostent *host=NULL;
	
	memcpy(szServer,(char *)LPCTSTR(strServerIP),nIPLen);
	
	server.sin_family=AF_INET;
	server.sin_port=htons(uPort);
	server.sin_addr.s_addr=inet_addr(szServer);
	
	if(server.sin_addr.s_addr==INADDR_NONE)
	{
		host=gethostbyname(szServer);
		if(host==NULL)
		{
			delete[] szServer;
			closesocket(m_Socket);
			m_Socket=INVALID_SOCKET;
			return FALSE;
		}
		CopyMemory(&server.sin_addr,host->h_addr_list[0],host->h_length);
	}
	delete[] szServer;
	
	if(WSAAsyncSelect(m_Socket,m_hWnd,m_uMsg,FD_READ | FD_CLOSE | FD_CONNECT)==SOCKET_ERROR)
	{
		closesocket(m_Socket);
		m_Socket=INVALID_SOCKET;
		return FALSE;
	}
	
	connect(m_Socket,(struct sockaddr *)&server,sizeof(server));

	return TRUE;
}

BOOL CMySocket::Listen(u_short uPort)
{
	if(!InitSocket())
		return FALSE;

	struct sockaddr_in local;
	local.sin_addr.s_addr=htonl(INADDR_ANY);
	local.sin_family=AF_INET;
	local.sin_port=htons(uPort);
	
	if(bind(m_Socket,(struct sockaddr *)&local,sizeof(local))==SOCKET_ERROR)
	{
		closesocket(m_Socket);
		m_Socket=INVALID_SOCKET;
		return FALSE;
	}
	
	if(WSAAsyncSelect(m_Socket,m_hWnd,m_uMsg,FD_ACCEPT)==SOCKET_ERROR)
	{
		closesocket(m_Socket);
		m_Socket=INVALID_SOCKET;
		return FALSE;
	}
	
	if(listen(m_Socket,SOMAXCONN)==SOCKET_ERROR)
	{
		closesocket(m_Socket);
		m_Socket=INVALID_SOCKET;
		return FALSE;
	}

	return TRUE;
}

CMySocket *CMySocket::Accept()
{
	struct sockaddr_in client;
	int nAddrSize = sizeof(client);
	SOCKET socket=::accept(m_Socket,(struct sockaddr *)&client,&nAddrSize);
	if(socket==INVALID_SOCKET)
		return NULL;

	CMySocket *pMySocket=new CMySocket(m_hWnd,m_uMsg);
	if(pMySocket==NULL)
		return NULL;
	pMySocket->m_Socket=socket;

	if(WSAAsyncSelect(pMySocket->m_Socket,pMySocket->m_hWnd,pMySocket->m_uMsg,FD_READ | FD_CLOSE)==SOCKET_ERROR)
	{
		delete pMySocket;
		return NULL;
	}

	return pMySocket;
}

int CMySocket::Send(BYTE *pBuf, int nBufLen)
{
	if(m_Socket==INVALID_SOCKET)
		return SOCKET_ERROR;

	if((pBuf==NULL) || (nBufLen<=0))
		return 0;

	int nPos=0;
	int nRet=0;
	while((nPos<nBufLen) && (nRet!=SOCKET_ERROR))
	{
		nRet=send(m_Socket,(char *)pBuf+nPos,nBufLen-nPos,0);
		if(nRet!=SOCKET_ERROR)
			nPos+=nRet;
	}

	return nPos;
}

int CMySocket::Recv(BYTE *pBuf, int nBufLen)
{
	if(m_Socket==INVALID_SOCKET)
		return SOCKET_ERROR;

	if((pBuf==NULL) || (nBufLen<=0))
		return 0;

	return recv(m_Socket,(char *)pBuf,nBufLen,0);
}

BOOL CMySocket::Close()
{
	if(m_Socket!=INVALID_SOCKET)
		closesocket(m_Socket);
	m_Socket=INVALID_SOCKET;

	return TRUE;
}

int CMySocket::GetRemoteInfo(CString &strIP, u_short &uPort)
{
	if(m_Socket==INVALID_SOCKET)
		return SOCKET_ERROR;

	struct sockaddr_in sa;
	int saLen=sizeof(sa);

	int nRet=getpeername(m_Socket,(struct sockaddr *)&sa,&saLen);
	if(nRet==SOCKET_ERROR)
		return SOCKET_ERROR;

	strIP=CString(inet_ntoa(sa.sin_addr));
	uPort=sa.sin_port;

	return nRet;
}