ClientComm.cpp 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792
  1. // ClientComm.cpp: implementation of the CClientComm class.
  2. //
  3. //////////////////////////////////////////////////////////////////////
  4. #include "ErrorCode.h"
  5. #include <winpr/wtypes.h>
  6. #ifdef _INCLUDE_SPBASE_
  7. #include "SpBase.h"
  8. #else
  9. #define Dbg
  10. #endif
  11. #include "ClientComm.h"
  12. #include "Package.h"
  13. #include <assert.h>
  14. #include <string.h>
  15. #include <winpr/crt.h>
  16. #include <winpr/thread.h>
  17. #include <chrono>
  18. #include <iomanip>
  19. #include <ctime>
  20. #include <iostream>
  21. DWORD CClientComm::s_dwSessionID =0;
  22. DWORD CClientComm::s_dwTokenHash =0;
  23. BYTE CClientComm::s_arrSessionKey[16] = {0};
  24. void outputCurrentTime(std::string msg)
  25. {
  26. using namespace chrono;
  27. system_clock::time_point now = system_clock::now();
  28. std::time_t last = system_clock::to_time_t(now);
  29. std::cout << msg << std::put_time(std::localtime(&last), "%F %T") << endl;
  30. }
  31. //////////////////////////////////////////////////////////////////////
  32. // Construction/Destruction
  33. //////////////////////////////////////////////////////////////////////
  34. #define SafeCallbackOnSocketError() \
  35. try \
  36. { \
  37. int nErrCode = WSAGetLastError(); \
  38. string errMsg = GetSysErrorMsg(nErrCode); \
  39. if (m_pCallback != NULL)\
  40. m_pCallback->OnError(Error_IO, 0, errMsg.c_str()); \
  41. } \
  42. catch(...) {}
  43. #define SafeCallbackOnError(dwSysCode, dwUserCode, pErrMsg) \
  44. try \
  45. { \
  46. if (m_pCallback != NULL)\
  47. m_pCallback->OnError(dwSysCode, dwUserCode, pErrMsg); \
  48. } \
  49. catch(...) {}
  50. #define SafeCallbackOnClose() \
  51. try \
  52. { \
  53. if (m_pCallback != NULL)\
  54. m_pCallback->OnClose();\
  55. } \
  56. catch (...) {}
  57. CClientComm::CClientComm(CSecureClientBase *pCallback)
  58. :SOCKET_RECV_BUF_LEN(8096)
  59. {
  60. assert(pCallback != NULL);
  61. m_pCallback = pCallback;
  62. m_hSocket = INVALID_SOCKET;
  63. m_hWorkThread = 0;
  64. m_eState = State_None;
  65. m_bNeedAuth = true;
  66. m_hRecvEvent = ::CreateEventA(NULL, FALSE, FALSE, NULL);
  67. m_nSyncWaitResult = 0;
  68. m_nRecvBufLen = SOCKET_RECV_BUF_LEN * 100; //800K
  69. m_pRecvBuf = new BYTE[m_nRecvBufLen];
  70. memset(m_pRecvBuf, 0, m_nRecvBufLen);
  71. m_nHasRecvLen =0;
  72. m_dwSessionID =0;
  73. m_dwTokenHash =0;
  74. memset(m_arrSessionKey, 0, 16);
  75. }
  76. CClientComm::~CClientComm()
  77. {
  78. Close();
  79. if (m_hRecvEvent != 0)
  80. {
  81. CloseHandle(m_hRecvEvent);
  82. m_hRecvEvent = 0;
  83. }
  84. delete[] m_pRecvBuf;
  85. m_pRecvBuf = NULL;
  86. m_nRecvBufLen = 0;
  87. }
  88. // 创建连接,@option:1、重新鉴权新建会话密钥;2、通过握手使用缓存会话密钥;
  89. // 3、不使用会话密钥,即非安全通道 4、不协商,直接使用共享会话密钥
  90. bool CClientComm::Connect(const char *pServerAddr, int nPort, int nOption)
  91. {
  92. m_eState = State_Connecting;
  93. m_hSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
  94. if (m_hSocket == INVALID_SOCKET)
  95. {
  96. m_eState = State_Error;
  97. SafeCallbackOnSocketError();
  98. return false;
  99. }
  100. //SetSocketOption();
  101. sockaddr_in addr;
  102. ZeroMemory(&addr, sizeof(sockaddr_in));
  103. addr.sin_family = AF_INET;
  104. addr.sin_addr.s_addr = inet_addr(pServerAddr);
  105. addr.sin_port = htons(nPort);
  106. if (connect(m_hSocket, (SOCKADDR*)&addr, sizeof(addr)) == SOCKET_ERROR)
  107. {
  108. m_eState = State_Error;
  109. SafeCallbackOnSocketError();
  110. return false;
  111. }
  112. m_bNeedAuth = (nOption != 3);
  113. // 创建接收线程,支持异步模式
  114. DWORD nThreadID;
  115. m_hWorkThread = (HANDLE)CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)&RecvThreadFunc, this, 0, &nThreadID);
  116. if (m_hWorkThread == 0)
  117. {
  118. Close();
  119. SafeCallbackOnError(Error_Unexpect, 0, "_beginthreadex() error");
  120. m_eState = State_Error;
  121. return false;
  122. }
  123. bool bSendSuc = false;
  124. if (nOption == 3)
  125. {
  126. // 非鉴权握手
  127. bSendSuc = SendHelloReqPackage();
  128. }
  129. else if (nOption == 4) // 使用共享会话密钥
  130. {
  131. bSendSuc = SendUseSharedSKPackage();
  132. }
  133. else if (nOption == 2 && s_dwSessionID !=0) // 连接成功,如果要重用会话密钥,则发送握手请求
  134. {
  135. bSendSuc = SendShakeReqPackage();
  136. }
  137. else
  138. {
  139. // 重新鉴权
  140. bSendSuc = SendAuthReqPackage();
  141. }
  142. // 同步等待返回
  143. WaitRetPkg:
  144. if (!bSendSuc)
  145. {
  146. SafeCallbackOnError(Error_Unexpect, 0, "send auth pkg fail");
  147. Close();
  148. m_eState = State_Error;
  149. return false;
  150. }
  151. auto pkg = ReceivePackage(10);
  152. if (pkg == NULL)
  153. {
  154. if (m_eState == State_Connecting)
  155. SafeCallbackOnError(Error_TimeOut, 0, "connect timeout");
  156. m_eState = State_Error;
  157. Close();
  158. return false;
  159. }
  160. bool bAuthSuc = false;
  161. if (strcmp(pkg->GetServiceCode().c_str(), "__AUTH__") ==0)
  162. {
  163. bAuthSuc = HandleAuthRetPackage(pkg);
  164. }
  165. else if (strcmp(pkg->GetServiceCode().c_str(), "_SHAKE_") == 0)
  166. {
  167. bAuthSuc = HandleShakeRetPackage(pkg);
  168. if (!bAuthSuc)
  169. {
  170. // 握手失败,重新鉴权
  171. bSendSuc = SendAuthReqPackage();
  172. goto WaitRetPkg;
  173. }
  174. }
  175. else if (strcmp(pkg->GetServiceCode().c_str(), "_USESSK_") == 0)
  176. {
  177. bAuthSuc = HandleUseSharedSKRetPackage(pkg);
  178. if (!bAuthSuc)
  179. {
  180. // 使用共享密钥请求失败,重新鉴权
  181. bSendSuc = SendAuthReqPackage();
  182. goto WaitRetPkg;
  183. }
  184. }
  185. else if (strcmp(pkg->GetServiceCode().c_str(), "_HELLO_") == 0)
  186. {
  187. bAuthSuc = HandleHelloRetPackage(pkg);
  188. }
  189. else
  190. {
  191. // 继续等待
  192. goto WaitRetPkg;
  193. }
  194. if (bAuthSuc)
  195. {
  196. m_eState = State_OK;
  197. // 回调OnAuthPass事件
  198. try
  199. {
  200. m_pCallback->OnAuthPass();
  201. }
  202. catch(...)
  203. {
  204. SafeCallbackOnError(Error_Exception, 0, "OnAuthPass exception");
  205. }
  206. }
  207. else
  208. {
  209. m_eState = State_Error;
  210. }
  211. return bAuthSuc;
  212. }
  213. CSmartPointer<IPackage> CClientComm::CreateNewPackage(const char *pServiceCode)
  214. {
  215. CSmartPointer<IPackage> ptr;
  216. ptr.Attach(new CCommPackage(m_bNeedAuth && m_eState == State_OK ? m_arrSessionKey : NULL, pServiceCode));
  217. return ptr;
  218. }
  219. CSmartPointer<IPackage> CClientComm::CreateReplyPackage(const CSmartPointer<IPackage>& pRecvPkg)
  220. {
  221. CCommPackage* p = dynamic_cast<CCommPackage*>(pRecvPkg.GetRawPointer());
  222. CSmartPointer<IPackage> ptr;
  223. ptr.Attach(new CCommPackage(p));
  224. return ptr;
  225. }
  226. // socket接收线程
  227. DWORD CClientComm::RecvThreadFunc(void *pArg)
  228. {
  229. CClientComm *pThis = (CClientComm*) pArg;
  230. assert(pThis != NULL);
  231. while(pThis->m_eState == State_OK || pThis->m_eState == State_Connecting)
  232. {
  233. pThis->RecvThreadProcess();
  234. }
  235. // 清空接收缓存
  236. pThis->m_nHasRecvLen = 0;
  237. // 退出线程
  238. //_endthreadex(0);
  239. return 0;
  240. }
  241. bool CClientComm::IsConnectionOK()
  242. {
  243. return m_eState == State_OK ;
  244. }
  245. bool CClientComm::IsSecureConnection()
  246. {
  247. return m_bNeedAuth;
  248. }
  249. void CClientComm::RecvThreadProcess()
  250. {
  251. // 查询状态
  252. fd_set fdsRead;
  253. FD_ZERO(&fdsRead);
  254. FD_SET(m_hSocket, &fdsRead);
  255. timeval waitTime;
  256. waitTime.tv_sec = 0;
  257. waitTime.tv_usec = 100 * 1000;
  258. int rc = select(m_hSocket + 1, &fdsRead, NULL, NULL, &waitTime);
  259. if (rc >0) // 有数据到达
  260. {
  261. int nLen = recv(m_hSocket, (char *)m_pRecvBuf + m_nHasRecvLen, m_nRecvBufLen - m_nHasRecvLen, 0);
  262. if (nLen ==0)
  263. {
  264. // 连接已中断
  265. if (m_eState == State_OK)
  266. {
  267. m_eState = State_Disconnected;
  268. SafeCallbackOnError(Error_NetBroken, 0, "socket disconnected");
  269. SafeCallbackOnClose();
  270. }
  271. return;
  272. }
  273. else if (nLen <0)
  274. {
  275. bool bCallback = m_eState == State_OK;
  276. m_eState = State_Error;
  277. if (bCallback)
  278. {
  279. SafeCallbackOnSocketError();
  280. SafeCallbackOnClose();
  281. }
  282. m_eState = State_Error;
  283. return;
  284. }
  285. else
  286. {
  287. m_nHasRecvLen += nLen;
  288. // 可能一次收到多个包
  289. while(true)
  290. {
  291. // 检查包头是否完整
  292. if (m_nHasRecvLen < sizeof(CPackageHeader))
  293. return;
  294. // 提取整包长度
  295. int nPackageLen = GetPackageLenFromRecvBuf();
  296. // 重新分配接收缓冲区
  297. if (nPackageLen > m_nRecvBufLen)
  298. {
  299. BYTE *pNewBuf = new BYTE[nPackageLen];
  300. memset(pNewBuf, 0, nPackageLen);
  301. m_nRecvBufLen = nPackageLen;
  302. memcpy(pNewBuf, m_pRecvBuf, m_nHasRecvLen);
  303. delete[] m_pRecvBuf;
  304. m_pRecvBuf = pNewBuf;
  305. return;
  306. }
  307. if (m_nHasRecvLen < nPackageLen)
  308. return;
  309. // 解析整包长度
  310. CCommPackage *p = new CCommPackage(m_bNeedAuth && m_eState==State_OK ? m_arrSessionKey : NULL, NULL);
  311. CSmartPointer<IPackage> ppkg;
  312. ppkg.Attach(p);
  313. string strErrMsg;
  314. if (p->ParseRecvData(m_pRecvBuf, &nPackageLen, strErrMsg))
  315. {
  316. // 从缓冲区删除完整包数据
  317. memmove(m_pRecvBuf, m_pRecvBuf + nPackageLen, m_nHasRecvLen - nPackageLen);
  318. m_nHasRecvLen -= nPackageLen;
  319. // 检查是否底层错误包
  320. if (strcmp(ppkg->GetServiceCode().c_str(), "__ERROR_") ==0)
  321. {
  322. DWORD dwSysCode(0), dwUserCode(0);
  323. string strErrMsg;
  324. if (!ppkg->GetErrMsg(dwSysCode, dwUserCode, strErrMsg))
  325. {
  326. dwSysCode = Error_Bug;
  327. strErrMsg = "get comm package error msg fail";
  328. }
  329. SafeCallbackOnError(dwSysCode, dwUserCode, strErrMsg.c_str());
  330. m_eState = State_Error;
  331. if (::InterlockedCompareExchange(&m_nSyncWaitResult, 0, 1) == 1)
  332. {
  333. ::SetEvent(m_hRecvEvent);
  334. //outputCurrentTime("setEvent 1");
  335. }
  336. return;
  337. }
  338. // 回调应用处理
  339. if (::InterlockedCompareExchange(&m_nSyncWaitResult, 0, 1) == 1)
  340. {
  341. CAutoLock Lock(&m_LockObject,1000);
  342. m_RecvPackages.push_back(ppkg.Detach());
  343. ::SetEvent(m_hRecvEvent);
  344. //outputCurrentTime("setEvent 2");
  345. }
  346. else
  347. {
  348. // 回调
  349. if (m_pCallback != NULL)
  350. {
  351. try
  352. {
  353. m_pCallback->OnReceivePackage(ppkg);
  354. }
  355. catch (...) {}
  356. }
  357. }
  358. }
  359. else
  360. {
  361. // 丢弃全部接收数据
  362. m_nHasRecvLen = 0;
  363. SafeCallbackOnError(Error_Bug, 0, strErrMsg.c_str());
  364. SafeCallbackOnClose();
  365. m_eState = State_Error;
  366. }
  367. }
  368. }
  369. }
  370. else if (rc ==0) // 超时
  371. {
  372. }
  373. else if (rc == SOCKET_ERROR) // 错误
  374. {
  375. if (m_eState == State_OK)
  376. {
  377. SafeCallbackOnSocketError();
  378. SafeCallbackOnClose();
  379. m_eState = State_Error;
  380. }
  381. }
  382. }
  383. bool CClientComm::SendUseSharedSKPackage()
  384. {
  385. CUseSharedSKReq req;
  386. memset(&req, 0, sizeof(req));
  387. int nSKLen = 16;
  388. int nTNLen = 16;
  389. if (!m_pCallback->OnGetSharedSK(req.m_arrTerminalNo, &nTNLen, m_arrSessionKey, &nSKLen))
  390. {
  391. SafeCallbackOnError(Error_Unexpect, 0, "get shared seesion key fail");
  392. return false;
  393. }
  394. // 生成Hash
  395. BYTE md5[16];
  396. memset(md5, 0, 16);
  397. MD5Hash(m_arrSessionKey, 16, md5);
  398. DWORD dwHash = ((DWORD)md5[0]) << 24 | ((DWORD)md5[1]) << 16 | ((DWORD)md5[2]) << 8 | ((DWORD)md5[3]);
  399. req.m_dwSharedSKHash = dwHash;
  400. CSmartPointer<IPackage> pkg = CreateNewPackage("_USESSK_");
  401. pkg->AddStruct("USESKREQ", false, false, (BYTE*)&req, sizeof(req));
  402. return SendPackage(pkg) != "";
  403. }
  404. bool CClientComm::SendShakeReqPackage()
  405. {
  406. CSessionShakeReq req;
  407. memset(&req, 0, sizeof(req));
  408. //ASSERT(pTerminalNo != NULL);
  409. req.m_dwTokenHash = s_dwTokenHash;
  410. req.m_dwSessionID = s_dwSessionID;
  411. CSmartPointer<IPackage> pkg = CreateNewPackage("_SHAKE_");
  412. pkg->AddStruct("SHAKEREQ", false, false, (BYTE*)&req, sizeof(req));
  413. return SendPackage(pkg) != "";
  414. }
  415. bool CClientComm::SendHelloReqPackage()
  416. {
  417. CSmartPointer<IPackage> pkg = CreateNewPackage("_HELLO_");
  418. return SendPackage(pkg) != "";
  419. }
  420. bool CClientComm::SendAuthReqPackage()
  421. {
  422. CConnAuthReq authReq;
  423. memset(&authReq, 0, sizeof(authReq));
  424. try
  425. {
  426. if (!m_pCallback->OnAuthRequest(&authReq))
  427. {
  428. SafeCallbackOnError(Error_Unexpect, 0, "callback OnAuthRequest() error");
  429. return false;
  430. }
  431. }
  432. catch(...)
  433. {
  434. SafeCallbackOnError(Error_Exception, 0, "callback OnAuthRequest() exception");
  435. return false;
  436. }
  437. // 保存TokenHash
  438. BYTE md5[16];
  439. memset(md5, 0, 16);
  440. MD5Hash(authReq.m_arrVerifyToken, sizeof(authReq.m_arrVerifyToken), md5);
  441. m_dwTokenHash = ((DWORD)md5[0]) << 24 | ((DWORD)md5[1]) << 16 | ((DWORD)md5[2]) << 8 | ((DWORD)md5[3]);
  442. {
  443. CAutoLock lock(&m_LockObject);
  444. s_dwTokenHash = m_dwTokenHash;
  445. }
  446. CSmartPointer<IPackage> ppkg = CreateNewPackage("__AUTH__");
  447. ppkg->AddStruct("AUTH_REQ", false, false, (BYTE*)&authReq, sizeof(authReq));
  448. return SendPackage(ppkg) != "";
  449. }
  450. bool CClientComm::HandleHelloRetPackage(const CSmartPointer<IPackage> &pRetPkg)
  451. {
  452. DWORD dwSysCode(0), dwUserCode(0);
  453. string errMsg;
  454. if (pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
  455. {
  456. SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
  457. return false;
  458. }
  459. return true;
  460. }
  461. bool CClientComm::HandleShakeRetPackage(const CSmartPointer<IPackage> &pRetPkg)
  462. {
  463. // 握手通过
  464. DWORD dwSysCode(0), dwUserCode(0);
  465. string errMsg;
  466. if (!pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
  467. {
  468. // 将缓存的Session信息恢复
  469. {
  470. CAutoLock lock(&m_LockObject);
  471. m_dwSessionID = s_dwSessionID;
  472. m_dwTokenHash = s_dwTokenHash;
  473. memcpy(m_arrSessionKey, s_arrSessionKey, 16);
  474. }
  475. return true;
  476. }
  477. // 握手失败,重新鉴权
  478. SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
  479. return false;
  480. }
  481. bool CClientComm::HandleUseSharedSKRetPackage(const CSmartPointer<IPackage> &pRetPkg)
  482. {
  483. // 握手通过
  484. DWORD dwSysCode(0), dwUserCode(0);
  485. string errMsg;
  486. if (pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
  487. {
  488. SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
  489. return false;
  490. }
  491. return true;
  492. }
  493. bool CClientComm::HandleAuthRetPackage(const CSmartPointer<IPackage> &pRetPkg)
  494. {
  495. DWORD dwSysCode(0), dwUserCode(0);
  496. string errMsg;
  497. if (pRetPkg->GetErrMsg(dwSysCode, dwUserCode, errMsg))
  498. {
  499. SafeCallbackOnError(dwSysCode, dwUserCode, errMsg.c_str());
  500. return false;
  501. }
  502. int nBufLen = pRetPkg->GetStructLen("AUTH_RET");
  503. if (nBufLen != sizeof(CConnAuthRet))
  504. {
  505. char buf[64];
  506. sprintf(buf, "struct length of [AUTH_RET] is invalid: %d", nBufLen);
  507. SafeCallbackOnError(Error_Bug, 0, buf);
  508. return false;
  509. }
  510. CConnAuthRet authRet;
  511. memset(&authRet, 0, sizeof(authRet));
  512. int nArrayNum(0);
  513. bool bRet = pRetPkg->GetStructData("AUTH_RET", (BYTE*)&authRet, &nBufLen, &nArrayNum);
  514. if (!bRet)
  515. {
  516. SafeCallbackOnError(Error_Bug, 0, "get struct data of [AUTH_RET] fail");
  517. return false;
  518. }
  519. // 解密并保存会话密钥
  520. int nLen = 256;
  521. BYTE buf[256];
  522. memset(buf, 0, 256);
  523. try
  524. {
  525. if (!m_pCallback->OnSessionKeyRet(&authRet, buf, &nLen))
  526. {
  527. SafeCallbackOnError(Error_Unexpect, 0, "callback OnSessionKeyRet() error");
  528. return false;
  529. }
  530. }
  531. catch(...)
  532. {
  533. SafeCallbackOnError(Error_Exception, 0, "decrypt return session key exception");
  534. return false;
  535. }
  536. if (nLen == 16)
  537. {
  538. memcpy(m_arrSessionKey, buf, 16);
  539. // 缓存SessionKey
  540. {
  541. CAutoLock lock(&m_LockObject);
  542. memcpy(s_arrSessionKey, m_arrSessionKey, 16);
  543. }
  544. // 生成SessionID并保存
  545. BYTE md5[16];
  546. memset(md5, 0, 16);
  547. MD5Hash(m_arrSessionKey, 16, md5);
  548. m_dwSessionID = ((DWORD)md5[0]) << 24 | ((DWORD)md5[1]) << 16 | ((DWORD)md5[2]) << 8 | ((DWORD)md5[3]);
  549. {
  550. CAutoLock lock(&m_LockObject);
  551. s_dwSessionID = m_dwSessionID;
  552. }
  553. return true;
  554. }
  555. else
  556. {
  557. SafeCallbackOnError(Error_Bug, 0, "decrypted sesseion key length invalid");
  558. return false;
  559. }
  560. }
  561. int CClientComm::GetPackageLenFromRecvBuf()
  562. {
  563. if (m_pRecvBuf == NULL || m_nRecvBufLen < sizeof(CPackageHeader) || m_nHasRecvLen < sizeof(CPackageHeader))
  564. return 0;
  565. CPackageHeader *pHead = (CPackageHeader*) m_pRecvBuf;
  566. return pHead->m_nPackageLen;
  567. }
  568. void CClientComm::SetSocketOption()
  569. {
  570. /*
  571. //Send time out
  572. int timeout = SOCKET_SENDTIMEOUT * 1000;
  573. int rc = setsockopt(nSocket, SOL_SOCKET, SO_SNDTIMEO, (char *)&timeout, sizeof(timeout));
  574. //recieve time out
  575. setsockopt(nSocket, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(timeout));
  576. // Send keepalives
  577. nSockOptVal = 1;
  578. setsockopt(nSocket, SOL_SOCKET, SO_KEEPALIVE, (const char FAR *)&nSockOptVal, sizeof(nSockOptVal));
  579. //Reuse Port
  580. BOOL bOptVal = TRUE;
  581. setsockopt(nSocket, SOL_SOCKET, SO_REUSEADDR, (const char FAR *)&bOptVal, sizeof(bOptVal));
  582. */
  583. // Enable "hard" close
  584. LINGER LingerOpt;
  585. LingerOpt.l_onoff = 0;
  586. LingerOpt.l_linger = 0;
  587. setsockopt(m_hSocket, SOL_SOCKET, SO_LINGER, (const char FAR *)&LingerOpt, sizeof(LingerOpt));
  588. // set recv buf lenght
  589. setsockopt(m_hSocket, SOL_SOCKET, SO_RCVBUF, (const char *) SOCKET_RECV_BUF_LEN, sizeof(int));
  590. }
  591. // 返回序列号
  592. string CClientComm::SendPackage(const CSmartPointer<IPackage>& pSendPkg)
  593. {
  594. CCommPackage *pkg = dynamic_cast<CCommPackage*>(pSendPkg.GetRawPointer());
  595. assert(pkg != NULL);
  596. if (m_bNeedAuth && m_eState != State_OK
  597. && pkg->GetServiceCode().compare("__AUTH__") != 0
  598. && pkg->GetServiceCode().compare("_HELLO_") != 0
  599. && pkg->GetServiceCode().compare("_USESSK_") != 0
  600. && pkg->GetServiceCode().compare("_SHAKE_") != 0)
  601. {
  602. SafeCallbackOnError(Error_Bug, 0, "conn auth first");
  603. return "";
  604. }
  605. //Dbg("begin send package: [%s]", pSendPkg->GetServiceCode().c_str());
  606. int nBufLen = pkg->GetPackageLen();
  607. BYTE *pBuf = new BYTE[nBufLen];
  608. memset(pBuf, 0, nBufLen);
  609. pkg->GenerateSendData(pBuf, &nBufLen);
  610. int nTotalSend =0;
  611. while (nTotalSend < nBufLen)
  612. {
  613. int nSendLen = send(m_hSocket, (const char *)pBuf + nTotalSend, nBufLen - nTotalSend, 0);
  614. if (nSendLen == SOCKET_ERROR)
  615. break;
  616. nTotalSend += nSendLen;
  617. }
  618. delete[] pBuf;
  619. #ifdef _INCLUDE_SPBASE_
  620. if (nTotalSend < nBufLen)
  621. LogError(Severity_Low, Error_NetBroken, 0, (const char*)CSimpleStringA::Format("send package [%s] fail, socket error: [%d]",
  622. pSendPkg->GetServiceCode().c_str(), WSAGetLastError()));
  623. //else
  624. //Dbg("send package [%s] succeed, total [%d] bytes", pSendPkg->GetServiceCode().c_str(), nTotalSend);
  625. #endif
  626. return nTotalSend == nBufLen ? pkg->GetPackageReqID() : "";
  627. }
  628. CSmartPointer<IPackage> CClientComm::ReceivePackage(int nWaitSecond)
  629. {
  630. if (m_eState != State_OK && m_eState != State_Connecting)
  631. {
  632. SafeCallbackOnError(Error_Bug, 0, "connection not ok!");
  633. return NULL;
  634. }
  635. DWORD dwWaitResult = WAIT_OBJECT_0;
  636. ResetEvent(m_hRecvEvent);
  637. //outputCurrentTime("ResetEvent 1");
  638. ::InterlockedExchange(&m_nSyncWaitResult, 1);
  639. if (m_RecvPackages.size() == 0)
  640. dwWaitResult = WaitForSingleObject(m_hRecvEvent, nWaitSecond * 1000);
  641. if (dwWaitResult == WAIT_OBJECT_0 && (m_eState == State_OK || m_eState == State_Connecting))
  642. {
  643. assert(m_RecvPackages.size() >0);
  644. CAutoLock lock(&m_LockObject, 1000);
  645. CSmartPointer<IPackage> ptr;
  646. ptr.Attach(m_RecvPackages.back());
  647. m_RecvPackages.pop_back();
  648. return ptr;
  649. }
  650. return NULL;
  651. }
  652. void CClientComm::Close()
  653. {
  654. if (m_eState == State_OK)
  655. m_eState = State_Closed;
  656. if (m_hWorkThread != 0)
  657. {
  658. CloseHandle(m_hWorkThread);
  659. m_hWorkThread = 0;
  660. }
  661. if (m_hSocket != INVALID_SOCKET)
  662. {
  663. shutdown(m_hSocket, SD_BOTH);
  664. closesocket(m_hSocket);
  665. m_hSocket = INVALID_SOCKET;
  666. }
  667. }