tcpconnwrapper.cpp

mysee网络直播源代码Mysee Lite是Mysee独立研发的网络视频流媒体播放系统。在应有了P2P技术和一系列先进流媒体技术之后

#include "stdafx.h"
#include "TCPConnWrapper.h"

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

namespace gaov
{
	namespace chaos
	{

		int MySelect(SOCKET sock)
		{
			int n = (int)sock + 1;
			fd_set	rfds, wfds;
			timeval tv = {10, 0}; // 设置超时时间为10s
			int		err;

			while(true)
			{
				FD_ZERO(&rfds); //clear ReadSockSet to Zero
				FD_ZERO(&wfds); //clear WriteSockSet to Zero
				FD_SET(sock, &rfds); //add sock to ReadSockSet
				//FD_SET(sock, &wfds); //add sock to WriteSockSet

				err = select(n, &rfds, &wfds, NULL, &tv);
				switch (err)
				{
				case 0: // time limit expired
					return 0;
				case -1:// socket error.
					//error("select: %s\n", strerror(errno));
					return -1;
				}

				if (FD_ISSET(sock, &rfds))
					return 1;
				else
					return 0;
			}
		}

		////////////////////////////////
		/////////// TCPConn ////////////
		////////////////////////////////
		bool TCPConn::InitNetwork()
		{
			WSADATA WsaData;

			if(WSAStartup(MAKEWORD(2, 2), &WsaData) != 0)
				return false;

			if(LOBYTE(WsaData.wVersion) != 2 || HIBYTE(WsaData.wVersion) != 2){
				WSACleanup();
				return false;
			}

			return true;
		}

		PSTREAM TCPConn::InitData(UINT maxlen)
		{
			if( !m_bNetworkInited ){
				if( !(m_bNetworkInited = InitNetwork()) )
					return NULL;
			}

			if (maxlen > out.size)
			{
				out.data = (UINT8 *) realloc(out.data, maxlen);
				if( ! out.data ) return NULL;

				out.size = maxlen;
			}

			out.p = out.data;
			out.end = out.data + out.size;
			return &out;
		}

		/* Send TCP transport data packet */
		int TCPConn::Send(PSTREAM s)
		{
			int length = int(s->end - s->data);
			int sent, total = 0;

			while (total < length)
			{
				sent = send(sock, (const char*)(s->data + total), length - total, 0);
				if (sent < 0) // socket error
				{
					closesocket(sock);
					return -1;
				}

				total += sent;
			}

			return 1;
		}

		/* Receive a message on the TCP layer */
		PSTREAM TCPConn::Recv(PSTREAM s, UINT length)
		{
			unsigned int new_length, end_offset, p_offset;
			int rcvd = 0;
			int err;

			if (s == NULL)		/* read into "new" stream */
			{
				if (length > in.size)
				{
					in.data = (UINT8 *) realloc(in.data, length);
					if( ! in.data ) return NULL;

					in.size = length;
				}
				in.end = in.p = in.data;
				s = &in;
			}
			else	/* append to existing stream */
			{
				new_length = int(s->end - s->data) + length;
				if (new_length > s->size)
				{
					p_offset = int(s->p - s->data);
					end_offset = int(s->end - s->data);
					s->data = (UINT8 *) realloc(s->data, new_length);
					if( ! s->data ) return NULL;

					s->size = new_length;
					s->p = s->data + p_offset;
					s->end = s->data + end_offset;
				}
			}

			while (length > 0)
			{
				err = MySelect(sock);
				if (err > 0)
				{
					rcvd = recv(sock, (char*)(s->end), length, 0);
					if (rcvd > 0)
					{
						s->end += rcvd;
						length -= rcvd;
					}
					else if(rcvd < 0) // socket error
					{
						closesocket(sock);
						sock = INVALID_SOCKET;
						return NULL;
					}
					else // Connection closed.
						return NULL;
				}
				else if( err < 0 ) // socket error
				{
					closesocket(sock);
					sock = INVALID_SOCKET;
					return NULL;
				}
				else
					return NULL;
			}

			return s;
		}

		////////////////////////////////////////////////////
		/////////////////// TCPClient //////////////////////
		////////////////////////////////////////////////////
		bool TCPClient::SendEx(UINT8 msgType, UINT8 subMsgType, PBYTE data, UINT dataLen)
		{	
			PSTREAM s = InitData( dataLen + 6 );
			if( !s )
				return false;

			// Header: 6 bytes
			out_uint8(s, msgType);
			out_uint8(s, subMsgType);
			out_uint32_be(s, dataLen);
			// Data
			if( dataLen>0 && data ) 
				out_uint8p(s, data, dataLen);
			s_mark_end(s);
			if( Send(s) <= 0 )
				return false;

			return true;
		}

		bool TCPClient::SendEx2(UINT8 msgType, UINT8 subMsgType, PBYTE data1, UINT dataLen1, PBYTE data2, UINT dataLen2)
		{
			UINT dataLen = dataLen1 + dataLen2;
			PSTREAM s = InitData( dataLen + 4*2 + 6 );
			if( !s )
				return false;

			// Header: 6 bytes
			out_uint8(s, msgType);
			out_uint8(s, subMsgType);
			out_uint32_be(s, dataLen); // dataLen = dataLen1 + dataLen2

			if( dataLen > 0 )
			{
				// Data1
				out_uint32_be(s, dataLen1);
				if(dataLen1>0 && data1) 
					out_uint8p(s, data1, dataLen1);

				// Data2
				out_uint32_be(s, dataLen2);
				if(dataLen2>0 && data2) 
					out_uint8p(s, data2, dataLen2);
			}

			s_mark_end(s);

			if( Send(s) <= 0 )
				return false;

			return true;
		}

		bool TCPClient::RecvEx(UINT8& msgType, UINT8& subMsgType, PBYTE& data, UINT& dataLen, const bool bRecvData/*=true*/)
		{
			PSTREAM s = Recv(NULL, 6);
			if( ! s )
				return false;

			in_uint8(s, msgType);
			in_uint8(s, subMsgType);
			in_uint32_be(s, dataLen);

			if( dataLen>0 && bRecvData )
			{
				s = Recv(s, dataLen);
				if( ! s )
					return false;

				in_uint8p(s, data, dataLen);
			}

			return true;
		}

		bool TCPClient::RecvEx2(UINT8& msgType, UINT8& subMsgType, PBYTE& data1, UINT& dataLen1, 
			PBYTE& data2, UINT& dataLen2, const bool bRecvData/*=true*/ )
		{
			PSTREAM s = Recv(NULL, 6);
			if( ! s ) 
				return false;

			in_uint8(s, msgType);
			in_uint8(s, subMsgType);
			UINT totalLen;
			in_uint32_be(s, totalLen);  // totalLen: dataLen1 + dataLen2

			if( totalLen > 0 )
			{
				s = Recv(s, totalLen + 4*2);
				if( ! s )
					return false;

				// Data1
				in_uint32_be(s, dataLen1);
				if(dataLen1>0 && bRecvData) 
					in_uint8p(s, data1, dataLen1);

				// Data2
				in_uint32_be(s, dataLen2);
				if(dataLen2>0 && bRecvData) 
					in_uint8p(s, data2, dataLen2);
			}

			return true;
		}

		/* Establish a connection on the TCP layer */
		// -3 : init network failed
		// -2 : connection exists; 
		// -1 : sock error
		// 0 : mem error
		// 1 : success
		int TCPClient::Connect(const char* sIP, USHORT port)
		{
			if( !m_bNetworkInited ){
				if( !(m_bNetworkInited = InitNetwork()) )
					return -3;
			}

			if( sock!=INVALID_SOCKET ) // exists
				return -2;

			if ( (sock = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET ){
				int lastErr = WSAGetLastError();
				//error("socket: %s\n", strerror(errno));
				return -1;
			}

			sockaddr_in addr;
			memset(&addr, 0, sizeof(addr));

			addr.sin_family = AF_INET;
			addr.sin_addr.s_addr = inet_addr(sIP);
			addr.sin_port = htons(port);

			if (connect(sock, (struct sockaddr *) &addr, sizeof(struct sockaddr)) < 0){
				//error("connect: %s\n", strerror(errno));
				int lastErr = WSAGetLastError();
				closesocket(sock);
				sock = INVALID_SOCKET;
				return -1;
			}

			int true_value = 1;
			setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (const char *) &true_value, sizeof(true_value));

			in.size = 4096;
			in.data = (UINT8 *) malloc(in.size);
			if( ! in.data ) return 0;

			out.size = 4096;
			out.data = (UINT8 *) malloc(out.size);
			if( ! out.data ) return 0;

			return 1;
		}


		////////////////////////////////////////////////////
		/////////////////// TCPServer //////////////////////
		////////////////////////////////////////////////////

		bool TCPServer::Listen(const char* sIP, USHORT port)
		{
			if( !m_bNetworkInited ){
				if( !(m_bNetworkInited = InitNetwork()) )
					return false;
			}

			// already exists
			if( sock!=INVALID_SOCKET )
				return false;

			sock = socket(AF_INET, SOCK_STREAM, 0);
			if( INVALID_SOCKET==sock )
				return false;

			struct sockaddr_in addr;
			int addrLen = sizeof(addr);
			memset(&addr, 0, addrLen);

			addr.sin_family = AF_INET;
			addr.sin_addr.s_addr = inet_addr(sIP);
			addr.sin_port = htons(port);

			if( bind(sock, (sockaddr *)&addr, addrLen) == SOCKET_ERROR )
				return false;

			if( listen(sock, 5) == SOCKET_ERROR )
				return false;

			return true;
		}

		SOCKET TCPServer::Accept()
		{
			if( !m_bNetworkInited || sock==INVALID_SOCKET )
				return INVALID_SOCKET;

			return accept( sock, NULL, NULL );
		}

	};
};