tcp-udp socket 简单工具类

#ifndef TCP_UDP_SOCKET_HPP
#define TCP_UDP_SOCKET_HPP

#include <winsock2.h>
#include <string>
#include <vector>

#pragma comment (lib,"ws2_32.lib")

typedef struct st_DataHead
{
	char szID[4];
	int option;
	int dataLen;
}DataHead,*PDataHead;

class TcpUdpSocket
{
public:
	TcpUdpSocket() { StartupSockect();}
	~TcpUdpSocket() { CleanupSocket();}
	TcpUdpSocket(const TcpUdpSocket& rhs);
	TcpUdpSocket& operator=(const TcpUdpSocket& rhs);

public:

	BOOL StartupSockect()
	{
		WORD wVersion=MAKEWORD(2,2);
		WSADATA wsData;
		int nResult= WSAStartup(wVersion,&wsData);
		if(nResult !=0)
		{
			printf("WSAStartup falied %d",nResult);
			WSACleanup();
			return FALSE;
		}

		udpSocket_ = INVALID_SOCKET;
		tcpSocket_ = INVALID_SOCKET;

		return TRUE;
	}

	void CleanupSocket()
	{
		if(INVALID_SOCKET != udpSocket_)
			closesocket(udpSocket_);
		if(INVALID_SOCKET != tcpSocket_)
			closesocket(tcpSocket_);

		WSACleanup();
	}

	//
	// UDP
	//
	BOOL  UdpInitSock(const std::string& serverIP, unsigned short serverPort )
	{
		udpSocket_ = socket(AF_INET,SOCK_DGRAM,0);
		if(udpSocket_==INVALID_SOCKET)
		{
			CleanupSocket();
			return FALSE;
		}
		serverSockAddr_.sin_family = AF_INET;
		serverSockAddr_.sin_port = htons(serverPort);
		serverSockAddr_.sin_addr.S_un.S_addr=inet_addr(serverIP.c_str());
		return TRUE;
	}

	int UdpSendto(const char *sendData, int sendLen)
	{
		if(udpSocket_ == INVALID_SOCKET)
		{
			printf("udpSocket_ should be initialized first");
			return -1;
		}

		int nSend=sendto(udpSocket_, sendData, sendLen,0,(sockaddr*)&serverSockAddr_,sizeof(serverSockAddr_));
		if(nSend <= 0)
		{
			CleanupSocket();
			return -1;
		}
		return nSend;
	}

	int UpdRecvFrom(char *recvBuf, int recvLen)
	{
		if(udpSocket_ == INVALID_SOCKET)
		{
			printf("udpSocket_ should be initialized first");
			return -1;
		}

		struct sockaddr_in recvSockAddr;
		int recvSockaddrLen = sizeof(sockaddr);
		int ret = recvfrom(udpSocket_, recvBuf, recvLen, 0, (sockaddr*)&recvSockAddr, &recvSockaddrLen);
		if (ret <= 0 )
		{
			CleanupSocket();
			return -1;
		}
		return ret;
	}

	//
	//TCP
	//

	BOOL  TcpInitSock(const std::string& serverIP, unsigned short serverPort )
	{
		if(serverIP.empty())
			return FALSE;

		tcpSocket_ = socket(AF_INET,SOCK_STREAM,0);
		if(tcpSocket_==INVALID_SOCKET)
		{
			int nErr = WSAGetLastError();
			return FALSE;
		}
		serverSockAddr_.sin_family = AF_INET;
		serverSockAddr_.sin_port = htons(serverPort);
		serverSockAddr_.sin_addr.S_un.S_addr=inet_addr(serverIP.c_str());
		// serverSockAddr_.sin_addr.s_addr = inet_addr(serverIP.c_str());
		return TRUE;
	}

	int TcpSendWithHeader(const char *sendData, int sendLen)
	{
		if( !sendData || sendLen <=0) 
			return -1;

		int nErr = -1;

		if(tcpSocket_ == INVALID_SOCKET)
		{
			printf("tcpSocket_ should be initialized first");
			nErr = WSAGetLastError();
			return nErr;
		}

		if (connect(tcpSocket_,(SOCKADDR*)&serverSockAddr_,sizeof(serverSockAddr_)) == SOCKET_ERROR)
		{
			nErr = WSAGetLastError();
			CleanupSocket();
			return nErr;
		}

		DataHead dataHead;
		memcpy(dataHead.szID, "CTSR", 4);
		dataHead.option = 0;
		dataHead.dataLen = sendLen;

		int nTotalLen = sizeof(dataHead) + sendLen;
		std::vector<char> finalSendBuff;
		finalSendBuff.resize(nTotalLen);

		char *pSendBuf = &finalSendBuff[0];
		int nLeftLen = nTotalLen;
		int nSentLen = 0;
		memcpy(pSendBuf, &dataHead, sizeof(dataHead));
		memcpy(pSendBuf + sizeof(dataHead), sendData, sendLen);
		char *ptr = pSendBuf;
		while(nLeftLen > 0)
		{
			nSentLen = send(tcpSocket_, ptr, nLeftLen, 0);
			if (SOCKET_ERROR == nSentLen)
			{
				nErr = WSAGetLastError();
				closesocket(tcpSocket_);
				
				return nErr;
			}
			nLeftLen = nTotalLen - nSentLen;
			ptr+=nSentLen;
		}

		if (tcpSocket_ != INVALID_SOCKET)
			closesocket(tcpSocket_);

		return nSentLen;
	}

private:

	SOCKET udpSocket_;
	SOCKET tcpSocket_;
	struct sockaddr_in serverSockAddr_;
};



#endif//TCP_UDP_SOCKET_HPP