myvtm/RVCComm/ClientComm.cpp

// ClientComm.cpp: implementation of the CClientComm class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "ErrorCode.h"

#include <WTypes.h>

#ifdef _INCLUDE_SPBASE_
#include "SpBase.h"
#else
#define  Dbg 
#endif 

#include "ClientComm.h"
#include "Package.h"

#include <process.h>
#include <assert.h>
#include <string.h>


DWORD CClientComm::s_dwSessionID =0;
DWORD CClientComm::s_dwTokenHash =0;
BYTE CClientComm::s_arrSessionKey[16] = {0};

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////
#define SafeCallbackOnSocketError() \
	try \
	{ \
		int nErrCode = WSAGetLastError(); \
		string errMsg = GetSysErrorMsg(nErrCode); \
		if (m_pCallback != NULL)\
			m_pCallback->OnError(Error_IO, 0, errMsg.c_str()); \
	} \
	catch(...) {}

#define SafeCallbackOnError(dwSysCode, dwUserCode, pErrMsg) \
	try \
	{ \
		if (m_pCallback != NULL)\
			m_pCallback->OnError(dwSysCode, dwUserCode, pErrMsg); \
	} \
	catch(...) {}

#define SafeCallbackOnClose() \
	try \
	{ \
		if (m_pCallback != NULL)\
			m_pCallback->OnClose();\
	} \
	catch (...) {}

CClientComm::CClientComm(CSecureClientBase *pCallback)
	:SOCKET_RECV_BUF_LEN(8096)
{
	assert(pCallback != NULL);

	m_pCallback = pCallback;
	m_hSocket = INVALID_SOCKET;
	m_hWorkThread = 0;
	m_eState = State_None;
	m_bNeedAuth = true;

	m_hRecvEvent = ::CreateEventA(NULL, FALSE, FALSE, NULL);
	m_nSyncWaitResult = 0;

	m_nRecvBufLen = SOCKET_RECV_BUF_LEN * 100;  //800K
	m_pRecvBuf = new BYTE[m_nRecvBufLen];	
	memset(m_pRecvBuf, 0, m_nRecvBufLen); 
	m_nHasRecvLen =0;

	m_dwSessionID =0;
	m_dwTokenHash =0;
	memset(m_arrSessionKey, 0, 16);
}

CClientComm::~CClientComm()
{
	Close();

	if (m_hRecvEvent != 0)
	{
		CloseHandle(m_hRecvEvent);
		m_hRecvEvent = 0;
	}

	delete[] m_pRecvBuf;
	m_pRecvBuf = NULL;
	m_nRecvBufLen = 0;
}


// 创建连接，@option：1、重新鉴权新建会话密钥；2、通过握手使用缓存会话密钥；
//	3、不使用会话密钥，即非安全通道  4、不协商，直接使用共享会话密钥
bool CClientComm::Connect(const char *pServerAddr, int nPort, int nOption)
{
	m_eState = State_Connecting;
	m_hSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	
	if (m_hSocket == INVALID_SOCKET)
	{		
		m_eState = State_Error;
		SafeCallbackOnSocketError();
		return false;
	}

	//SetSocketOption();

	sockaddr_in addr;
	ZeroMemory(&addr, sizeof(sockaddr_in));

	addr.sin_family = AF_INET;
	addr.sin_addr.s_addr = inet_addr(pServerAddr);
	addr.sin_port = htons(nPort);

	if (connect(m_hSocket, (SOCKADDR*)&addr, sizeof(addr)) ==  SOCKET_ERROR)
	{
		m_eState = State_Error;
		SafeCallbackOnSocketError();
		return false;
	}	

	m_bNeedAuth = (nOption != 3);

	// 创建接收线程，支持异步模式
	UINT nThreadID;
	m_hWorkThread = (HANDLE)_beginthreadex(NULL, 0, &RecvThreadFunc, this, 0, &nThreadID);
	if (m_hWorkThread == 0)
	{
		Close();
		SafeCallbackOnError(Error_Unexpect, 0, "_beginthreadex() error");
		m_eState = State_Error;
		return false;
	}

	bool bSendSuc = false;
	if (nOption == 3)
	{
		// 非鉴权握手
		bSendSuc = SendHelloReqPackage();
	}
	else if (nOption == 4)	// 使用共享会话密钥
	{
		bSendSuc = SendUseSharedSKPackage();
	}
	else if (nOption == 2 && s_dwSessionID !=0)	// 连接成功，如果要重用会话密钥，则发送握手请求	
	{
		bSendSuc = SendShakeReqPackage();
	}
	else
	{
		// 重新鉴权
		bSendSuc = SendAuthReqPackage();
	}
		
	// 同步等待返回
WaitRetPkg:
	if (!bSendSuc)
	{
		SafeCallbackOnError(Error_Unexpect, 0, "send auth pkg fail");
		Close();
		m_eState = State_Error;
		return false;
	}

	auto pkg = ReceivePackage(10);
	if (pkg == NULL)
	{		
		if (m_eState == State_Connecting)
			SafeCallbackOnError(Error_TimeOut, 0, "connect timeout");
	
		m_eState = State_Error;
		Close();
		return false;
	}

	bool bAuthSuc = false;
	if (strcmp(pkg->GetServiceCode().c_str(), "__AUTH__") ==0)
	{
		bAuthSuc = HandleAuthRetPackage(pkg);
	}
	else if (strcmp(pkg->GetServiceCode().c_str(), "_SHAKE_") == 0)
	{
		bAuthSuc = HandleShakeRetPackage(pkg);
		if (!bAuthSuc)
		{
			// 握手失败，重新鉴权
			bSendSuc = SendAuthReqPackage();
			goto WaitRetPkg;
		}
	}
	else if (strcmp(pkg->GetServiceCode().c_str(), "_USESSK_") == 0)
	{
		bAuthSuc = HandleUseSharedSKRetPackage(pkg);
		if (!bAuthSuc)
		{
			// 使用共享密钥请求失败，重新鉴权
			bSendSuc = SendAuthReqPackage();
			goto WaitRetPkg;
		}
	}
	else if (strcmp(pkg->GetServiceCode().c_str(), "_HELLO_") == 0)
	{
		bAuthSuc = HandleHelloRetPackage(pkg);
	}
	else 
	{
		// 继续等待
		goto WaitRetPkg;
	}

	if (bAuthSuc)
	{
		m_eState = State_OK;
		// 回调OnAuthPass事件
		try
		{
			m_pCallback->OnAuthPass();
		}
		catch(...)
		{
			SafeCallbackOnError(Error_Exception, 0, "OnAuthPass exception");
		}
	}
	else
	{
		m_eState = State_Error;
	}

	return bAuthSuc;
}

CSmartPointer<IPackage> CClientComm::CreateNewPackage(const char *pServiceCode)
{
	CSmartPointer<IPackage> ptr;
	ptr.Attach(new CCommPackage(m_bNeedAuth  && m_eState == State_OK ? m_arrSessionKey : NULL, pServiceCode));
	return ptr;
}

CSmartPointer<IPackage> CClientComm::CreateReplyPackage(const CSmartPointer<IPackage>& pRecvPkg)
{
	CCommPackage* p = dynamic_cast<CCommPackage*>(pRecvPkg.GetRawPointer());
	
	CSmartPointer<IPackage> ptr;
	ptr.Attach(new CCommPackage(p));
	return ptr;
}

// socket接收线程
UINT CClientComm::RecvThreadFunc(void *pArg)
{
	CClientComm *pThis = (CClientComm*) pArg;
	assert(pThis != NULL);

	while(pThis->m_eState == State_OK || pThis->m_eState == State_Connecting)
	{
		pThis->RecvThreadProcess();
	}
	
	// 清空接收缓存
	pThis->m_nHasRecvLen = 0;
	
	// 退出线程
	_endthreadex(0);
	return 0;
}

bool CClientComm::IsConnectionOK()
{
	return m_eState == State_OK ;
}

bool CClientComm::IsSecureConnection()
{
	return m_bNeedAuth;
}

void CClientComm::RecvThreadProcess()
{
	// 查询状态
	fd_set fdsRead;
	FD_ZERO(&fdsRead);
	FD_SET(m_hSocket, &fdsRead);

	timeval waitTime;
	waitTime.tv_sec = 0;
	waitTime.tv_usec = 100 * 1000;
	
	int rc = select(0, &fdsRead, NULL, NULL, &waitTime);
	if (rc >0)	// 有数据到达
	{
		int nLen = recv(m_hSocket, (char *)m_pRecvBuf + m_nHasRecvLen, m_nRecvBufLen - m_nHasRecvLen, 0);
		if (nLen ==0)
		{
			// 连接已中断
			if (m_eState == State_OK)
			{
				m_eState = State_Disconnected;
				SafeCallbackOnError(Error_NetBroken, 0, "socket disconnected");
				SafeCallbackOnClose();
			}

			return;
		}
		else if (nLen <0)
		{
			bool bCallback = m_eState == State_OK;
			m_eState = State_Error;

			if (bCallback)
			{
				SafeCallbackOnSocketError();
				SafeCallbackOnClose();
			}
			
			m_eState = State_Error;
			return;
		}
		else
		{
			m_nHasRecvLen += nLen;

			// 可能一次收到多个包
			while(true)
			{
				// 检查包头是否完整
				if (m_nHasRecvLen < sizeof(CPackageHeader))
					return;

				// 提取整包长度
				int nPackageLen = GetPackageLenFromRecvBuf();
				
				// 重新分配接收缓冲区
				if (nPackageLen > m_nRecvBufLen)
				{
					BYTE *pNewBuf = new BYTE[nPackageLen];
					memset(pNewBuf, 0, nPackageLen); 
					m_nRecvBufLen = nPackageLen;

					memcpy(pNewBuf, m_pRecvBuf, m_nHasRecvLen);
					delete[] m_pRecvBuf;
					m_pRecvBuf = pNewBuf;
					return;
				}
				
				if (m_nHasRecvLen < nPackageLen)
					return;

				// 解析整包长度
				CCommPackage *p = new CCommPackage(m_bNeedAuth && m_eState==State_OK ? m_arrSessionKey : NULL, NULL);
				CSmartPointer<IPackage> ppkg;
				ppkg.Attach(p);

				string strErrMsg; 
				if (p->ParseRecvData(m_pRecvBuf, &nPackageLen, strErrMsg))
				{
					// 从缓冲区删除完整包数据
					memmove(m_pRecvBuf, m_pRecvBuf + nPackageLen, m_nHasRecvLen - nPackageLen);
					m_nHasRecvLen -= nPackageLen;

					// 检查是否底层错误包
					if (strcmp(ppkg->GetServiceCode().c_str(), "__ERROR_") ==0)
					{
						DWORD dwSysCode(0), dwUserCode(0);
						string strErrMsg;
						if (!ppkg->GetErrMsg(dwSysCode, dwUserCode, strErrMsg))
						{
							dwSysCode = Error_Bug;
							strErrMsg = "get comm package error msg fail";
						}
							
						SafeCallbackOnError(dwSysCode, dwUserCode, strErrMsg.c_str());
						m_eState = State_Error;		

						if (::InterlockedCompareExchange(&m_nSyncWaitResult, 0, 1) == 1)
							::SetEvent(m_hRecvEvent);
						return;
					}

					// 回调应用处理
					if (::InterlockedCompareExchange(&m_nSyncWaitResult, 0, 1) == 1)
					{
						CAutoLock Lock(&m_LockObject,1000);
						m_RecvPackages.push_back(ppkg.Detach());
						::SetEvent(m_hRecvEvent);
					}
					else
					{
						// 回调
						if (m_pCallback != NULL)
						{
							try
							{
								m_pCallback->OnReceivePackage(ppkg);
							}
							catch (...) {}
						}
					}					
				}
				else
				{
					// 丢弃全部接收数据
					m_nHasRecvLen = 0;

					SafeCallbackOnError(Error_Bug, 0, strErrMsg.c_str());
					SafeCallbackOnClose();
					m_eState = State_Error;
				}
			}
		}
	}
	else if (rc ==0) // 超时
	{
	}
	else if (rc == SOCKET_ERROR) // 错误
	{
		if (m_eState == State_OK)
		{
			SafeCallbackOnSocketError();
			SafeCallbackOnClose();
			m_eState = State_Error;
		}
	}
}

bool CClientComm::SendUseSharedSKPackage()
{
	CUseSharedSKReq req;
	memset(&req, 0, sizeof(req));

	int nSKLen = 16;
	int nTNLen = 16;
	if (!m_pCallback->OnGetSharedSK(req.m_arrTerminalNo, &nTNLen, m_arrSessionKey, &nSKLen))
	{
		SafeCallbackOnError(Error_Unexpect, 0, "get shared seesion key fail");
		return false;
	}

	// 生成Hash
	BYTE md5[16];
	memset(md5, 0, 16);
	MD5Hash(m_arrSessionKey, 16, md5);
	DWORD dwHash = ((DWORD)md5[0]) << 24 | ((DWORD)md5[1]) << 16 | ((DWORD)md5[2]) << 8 | ((DWORD)md5[3]);

	req.m_dwSharedSKHash = dwHash;
	
	CSmartPointer<IPackage> pkg = CreateNewPackage("_USESSK_");
	pkg->AddStruct("USESKREQ", false, false, (BYTE*)&req, sizeof(req));
	return SendPackage(pkg) != "";	
}

bool CClientComm::SendShakeReqPackage()
{
	CSessionShakeReq req;
	memset(&req, 0, sizeof(req));

	//ASSERT(pTerminalNo != NULL);
	req.m_dwTokenHash = s_dwTokenHash;
	req.m_dwSessionID = s_dwSessionID;

	CSmartPointer<IPackage> pkg = CreateNewPackage("_SHAKE_");
	pkg->AddStruct("SHAKEREQ", false, false, (BYTE*)&req, sizeof(req));
	return SendPackage(pkg) != "";
}

bool CClientComm::SendHelloReqPackage()
{
	CSmartPointer<IPackage> pkg = CreateNewPackage("_HELLO_");
	return SendPackage(pkg) != "";
}

bool CClientComm::SendAuthReqPackage()
{
	CConnAuthReq authReq;
	memset(&authReq, 0, sizeof(authReq));

	try
	{
		if (!m_pCallback->OnAuthRequest(&authReq))
		{
			SafeCallbackOnError(Error_Unexpect, 0, "callback OnAuthRequest() error");
			return false;
		}
	}
	catch(...)
	{
		SafeCallbackOnError(Error_Exception, 0, "callback OnAuthRequest() exception");
		return false;
	}

	// 保存TokenHash
	BYTE md5[16];
	memset(md5, 0, 16);
	MD5Hash(authReq.m_arrVerifyToken, sizeof(authReq.m_arrVerifyToken), md5);
	m_dwTokenHash = ((DWORD)md5[0]) << 24 | ((DWORD)md5[1]) << 16 | ((DWORD)md5[2]) << 8 | ((DWORD)md5[3]);

	{
		CAutoLock lock(&m_LockObject);
		s_dwTokenHash = m_dwTokenHash;
	}

	CSmartPointer<IPackage> ppkg = CreateNewPackage("__AUTH__");
	ppkg->AddStruct("AUTH_REQ", false, false, (BYTE*)&authReq, sizeof(authReq));
	return SendPackage(ppkg) != "";
}

bool CClientComm::HandleHelloRetPackage(const CSmartPointer<IPackage> &pRetPkg)
{
	DWORD dwSysCode(0), dwUserCode(0);
	string errMsg;
	if (pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
	{
		SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
		return false;
	}

	return true;
}

bool CClientComm::HandleShakeRetPackage(const CSmartPointer<IPackage>  &pRetPkg)
{
	// 握手通过
	DWORD dwSysCode(0), dwUserCode(0);
	string  errMsg;
	if (!pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
	{
		// 将缓存的Session信息恢复
		{
			CAutoLock lock(&m_LockObject);
			m_dwSessionID = s_dwSessionID;
			m_dwTokenHash = s_dwTokenHash;
			memcpy(m_arrSessionKey, s_arrSessionKey, 16);
		}

		return true;
	}

	// 握手失败，重新鉴权	
	SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
	return false;
}

bool CClientComm::HandleUseSharedSKRetPackage(const CSmartPointer<IPackage> &pRetPkg)
{
	// 握手通过
	DWORD dwSysCode(0), dwUserCode(0);
	string errMsg;
	if (pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
	{
		SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
		return false;
	}

	return true;
}

bool CClientComm::HandleAuthRetPackage(const CSmartPointer<IPackage> &pRetPkg)
{
	DWORD dwSysCode(0), dwUserCode(0);
	string errMsg;
	if (pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
	{
		SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
		return false;
	}
	
	int nBufLen = pRetPkg->GetStructLen("AUTH_RET");
	if (nBufLen != sizeof(CConnAuthRet))
	{
		char buf[64];
		sprintf(buf, "struct length of [AUTH_RET] is invalid: %d", nBufLen);

		SafeCallbackOnError(Error_Bug, 0, buf);
		return false;
	}

	CConnAuthRet authRet;
	memset(&authRet, 0, sizeof(authRet));

	int nArrayNum(0);
	bool bRet = pRetPkg->GetStructData("AUTH_RET", (BYTE*)&authRet, &nBufLen, &nArrayNum);
	if (!bRet)
	{
		SafeCallbackOnError(Error_Bug, 0, "get struct data of [AUTH_RET] fail");
		return false;
	}
	
	// 解密并保存会话密钥
	int nLen = 256;
	BYTE buf[256];
	memset(buf, 0, 256);

	try
	{
		if (!m_pCallback->OnSessionKeyRet(&authRet, buf, &nLen))
		{
			SafeCallbackOnError(Error_Unexpect, 0, "callback OnSessionKeyRet() error");
			return false;
		}
	}
	catch(...)
	{
		SafeCallbackOnError(Error_Exception, 0, "decrypt return session key exception");
		return false;
	}

	if (nLen == 16)
	{
		memcpy(m_arrSessionKey, buf, 16);

		// 缓存SessionKey
		{
			CAutoLock lock(&m_LockObject);
			memcpy(s_arrSessionKey, m_arrSessionKey, 16);
		}

		// 生成SessionID并保存
		BYTE md5[16];
		memset(md5, 0, 16);
		MD5Hash(m_arrSessionKey, 16, md5);

		m_dwSessionID = ((DWORD)md5[0]) << 24 | ((DWORD)md5[1]) << 16 | ((DWORD)md5[2]) << 8 | ((DWORD)md5[3]);

		{
			CAutoLock lock(&m_LockObject);
			s_dwSessionID = m_dwSessionID;
		}

		return true;
	}
	else
	{
		SafeCallbackOnError(Error_Bug, 0, "decrypted sesseion key length invalid");
		return false;
	}
}

int CClientComm::GetPackageLenFromRecvBuf()
{
	if (m_pRecvBuf == NULL || m_nRecvBufLen < sizeof(CPackageHeader) || m_nHasRecvLen < sizeof(CPackageHeader))
		return 0;

	CPackageHeader *pHead = (CPackageHeader*) m_pRecvBuf;
	return pHead->m_nPackageLen;	
}

void CClientComm::SetSocketOption()
{
	/*
	//Send time out
	int timeout = SOCKET_SENDTIMEOUT * 1000;
	int rc = setsockopt(nSocket, SOL_SOCKET, SO_SNDTIMEO, (char *)&timeout, sizeof(timeout));

	//recieve time out
	setsockopt(nSocket, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(timeout));

	// Send keepalives
	nSockOptVal = 1;
	setsockopt(nSocket, SOL_SOCKET, SO_KEEPALIVE, (const char FAR *)&nSockOptVal, sizeof(nSockOptVal));
	
	//Reuse Port
	BOOL bOptVal = TRUE;
	setsockopt(nSocket, SOL_SOCKET, SO_REUSEADDR, (const char FAR *)&bOptVal, sizeof(bOptVal));
	*/

	// Enable "hard" close
	LINGER LingerOpt;
	LingerOpt.l_onoff = 0;
	LingerOpt.l_linger = 0;
	setsockopt(m_hSocket, SOL_SOCKET, SO_LINGER, (const char FAR *)&LingerOpt, sizeof(LingerOpt));

	// set recv buf lenght
	setsockopt(m_hSocket, SOL_SOCKET, SO_RCVBUF, (const char *) SOCKET_RECV_BUF_LEN, sizeof(int));
}

// 返回序列号
string CClientComm::SendPackage(const CSmartPointer<IPackage>& pSendPkg)
{	
	CCommPackage *pkg = dynamic_cast<CCommPackage*>(pSendPkg.GetRawPointer());
	assert(pkg != NULL);
	
	if (m_bNeedAuth && m_eState != State_OK 
		&& pkg->GetServiceCode().compare("__AUTH__") != 0 
		&& pkg->GetServiceCode().compare("_HELLO_") != 0
		&& pkg->GetServiceCode().compare("_USESSK_") != 0 
		&& pkg->GetServiceCode().compare("_SHAKE_") != 0)
	{
		SafeCallbackOnError(Error_Bug, 0, "conn auth first");
		return "";
	}

	//Dbg("begin send package: [%s]", pSendPkg->GetServiceCode().c_str());

	int nBufLen = pkg->GetPackageLen();
	BYTE *pBuf = new BYTE[nBufLen];
	memset(pBuf, 0, nBufLen);
	pkg->GenerateSendData(pBuf, &nBufLen);
	
	
	int nTotalSend =0;
	while (nTotalSend < nBufLen)
	{
		int nSendLen = send(m_hSocket, (const char *)pBuf + nTotalSend, nBufLen - nTotalSend, 0);

		if (nSendLen == SOCKET_ERROR)
			break;

		nTotalSend += nSendLen;
	}	

	delete[] pBuf;
	
#ifdef _INCLUDE_SPBASE_
	if (nTotalSend < nBufLen)
		LogError(Severity_Low, Error_NetBroken, 0,  (const char*)CSimpleStringA::Format("send package [%s] fail, socket error: [%d]", 
		pSendPkg->GetServiceCode().c_str(),  WSAGetLastError()));
	//else
		//Dbg("send package [%s] succeed, total [%d] bytes", pSendPkg->GetServiceCode().c_str(), nTotalSend);
#endif

	return nTotalSend == nBufLen ? pkg->GetPackageReqID() : "";
}

CSmartPointer<IPackage> CClientComm::ReceivePackage(int nWaitSecond)
{
	if (m_eState != State_OK && m_eState != State_Connecting)
	{
		SafeCallbackOnError(Error_Bug, 0, "connection not ok!");
		return NULL;
	}

	DWORD dwWaitResult = WAIT_OBJECT_0;
	ResetEvent(m_hRecvEvent);
	::InterlockedExchange(&m_nSyncWaitResult, 1);

	if (m_RecvPackages.size() == 0)
	{
		dwWaitResult = WaitForSingleObject(m_hRecvEvent, nWaitSecond * 1000);
	}
	
	if (dwWaitResult == WAIT_OBJECT_0 && (m_eState == State_OK || m_eState == State_Connecting))
	{
		assert(m_RecvPackages.size() >0);
		CAutoLock lock(&m_LockObject, 1000);

		CSmartPointer<IPackage> ptr;
		ptr.Attach(m_RecvPackages.back());
		m_RecvPackages.pop_back();
		return ptr;
	}

	return NULL;
}

void CClientComm::Close()
{
	if (m_eState == State_OK)
		m_eState = State_Closed;

	if (m_hWorkThread != 0)
	{
		CloseHandle(m_hWorkThread);		
		m_hWorkThread = 0;
	}

	if (m_hSocket != INVALID_SOCKET)
	{
		shutdown(m_hSocket, SD_BOTH);
		closesocket(m_hSocket);
		m_hSocket = INVALID_SOCKET;
	}
}