repos_assist/DLLInjection/RegistryInjection/InjectionDLL/InjectionDLL/PipeClient.cpp

#include "StdAfx.h"
#include "PipeClient.h"
#include "Utility.h"

#define BUFSIZE 512

CPipeClient::CPipeClient(LPCTSTR lpPipeName, DWORD dwMode)
{
    m_hPipeInst = INVALID_HANDLE_VALUE;
    m_bClientStop = FALSE;
    m_dwMode = dwMode;
    memset(m_szPipeName, 0, MAX_PATH*sizeof(TCHAR));
    memset(m_szWriteBuff, 0, 1024*sizeof(TCHAR));
    memset(m_szReceiveBuff, 0, 1024*sizeof(TCHAR));
    if ( lpPipeName )
        _stprintf_s(m_szPipeName, _T("%s"), lpPipeName);
}

CPipeClient::~CPipeClient(void)
{
    StopWork();
    if ( m_hPipeInst != INVALID_HANDLE_VALUE )
        CloseHandle(m_hPipeInst);
}

BOOL CPipeClient::StartWork()
{
    if ( !m_bClientStop )
    {
        BOOL bRet = TRUE;
        HANDLE hConnect = CreateThread(NULL, 0, ConnectThread, this, 0, NULL);
        HANDLE hReadMsg = CreateThread(NULL, 0, ReadMsgThread, this, 0, NULL);

        if ( hConnect == NULL || hReadMsg == NULL )
            bRet = FALSE;

        if ( hConnect )
            CloseHandle(hConnect);

        if ( hReadMsg )
            CloseHandle(hReadMsg);

        return bRet;
    }

    return TRUE;
}

DWORD CPipeClient::ConnectThread(LPVOID lpParam)
{
    CPipeClient *pInstance = (CPipeClient*)lpParam;
    if ( !pInstance ) 
        return 0L;

    while(!pInstance->m_bClientStop)
    {
        if ( pInstance->m_hPipeInst != INVALID_HANDLE_VALUE ) {
            // 1分钟检测;
            Sleep(60000);
            Utility::dprintf(_T("m_hPipeInst 已存在\n"));
            continue;
        }

        // 等待10秒;
        if ( !WaitNamedPipe(pInstance->m_szPipeName, 10000) )
        {// 如果管道不存在,会立即返回而不考虑超时值,所以此处仍要Sleep;
            Utility::dprintf(_T("<%ld> WaitNamedPipe 失败\n"), Utility::g_WndInfo.dwProcessId);
            Sleep(10000);
            continue;
        }

        pInstance->m_hPipeInst = CreateFile(
            pInstance->m_szPipeName,        // pipe name 
            GENERIC_READ | GENERIC_WRITE,   // read and write access 
            0,                              // no sharing 
            NULL,                           // default security attributes
            OPEN_EXISTING,                  // opens existing pipe 
            0,                              // default attributes 
            NULL);                          // no template file 

        // 创建成功，退出;
        if ( pInstance->m_hPipeInst != INVALID_HANDLE_VALUE ) 
        {
            // 管道连接成功，修改管道通信模式：message-read mode. 
            BOOL fSuccess = SetNamedPipeHandleState( 
                pInstance->m_hPipeInst,     // pipe handle 
                &pInstance->m_dwMode,       // new pipe mode 
                NULL,                       // don't set maximum bytes 
                NULL);                      // don't set maximum time 

            if (!fSuccess) {
                Utility::dprintf(_T("SetNamedPipeHandleState failed. GLE=%d\n"), GetLastError() ); 
                CloseHandle(pInstance->m_hPipeInst);
            }
        }
        else
        {
            // Exit if an error other than ERROR_PIPE_BUSY occurs. 
            if ( GetLastError() != ERROR_PIPE_BUSY ) 
            {
                Utility::dprintf(_T("Could not open pipe. GLE=%d\n"), GetLastError() ); 
            }
        }
    }

    Utility::dprintf(_T("<%ld> ConnectThread 退出\n"), Utility::g_WndInfo.dwProcessId);

    return 0;
}

DWORD CPipeClient::ReadMsgThread(LPVOID lpParam)
{
    DWORD cbRead = 0;
    BOOL bSuccess = FALSE;
    TCHAR chBuf[BUFSIZE]; 
    DWORD dwDataIndex = 0;
    DWORD dwError = 0;
    
    CPipeClient *pInstance = (CPipeClient*)lpParam;
    if ( !pInstance ) 
        return 0L;

    while(!pInstance->m_bClientStop)
    {
        if ( pInstance->m_hPipeInst == INVALID_HANDLE_VALUE ) {
            Sleep(2000);
            continue;
        }

        do 
        { 
            bSuccess = ReadFile( 
                pInstance->m_hPipeInst,     // pipe handle 
                chBuf,                      // buffer to receive reply 
                BUFSIZE*sizeof(TCHAR),      // size of buffer 
                &cbRead,                    // number of bytes read 
                NULL);                      // not overlapped 

            if ( !bSuccess && (dwError = GetLastError()) != ERROR_MORE_DATA )
                break; 

            Utility::dprintf(_T("读取数据：%ld, %ld, %s"), dwError, cbRead, chBuf);
            //Utility::dprintf(_T("读取数据：%ld, %ld"), dwError, cbRead);
#if 0
            TCHAR szMsg[8912] = {0};
            _stprintf_s(szMsg, _T("读取数据：%d, %ld, %ld, %s\n"), (int)bSuccess, dwError, cbRead, chBuf);
            OutputDebugString(szMsg);
#endif
            // 追回数据;
            memcpy(pInstance->m_szReceiveBuff + dwDataIndex, chBuf, cbRead);
            dwDataIndex += cbRead;
            Sleep(50);
        } while ( !bSuccess );  // repeat loop if ERROR_MORE_DATA 

        // 清空缓存数据;
        dwDataIndex = 0;
        memset(chBuf, 0, BUFSIZE*sizeof(TCHAR));
       
        if ( bSuccess )
        {
            Utility::dprintf(_T("读取到的消息=%d"), sizeof(pInstance->m_szReceiveBuff));
            Sleep(3000);
            Utility::dprintf(_T("读取到的消息=%s"), pInstance->m_szReceiveBuff);

            // 消息处理;
            // ...
        }
        else
        {
            Utility::dprintf(_T("ReadFile from pipe failed. GLE=%d\n"), dwError );
            if ( dwError == ERROR_PIPE_NOT_CONNECTED || dwError == ERROR_BROKEN_PIPE)
            {
                CloseHandle(pInstance->m_hPipeInst);
                pInstance->m_hPipeInst = INVALID_HANDLE_VALUE;
            }

#ifdef _DEBUG
            Sleep(10000);
#endif
        }

         memset(pInstance->m_szReceiveBuff, 0, BUFSIZE*sizeof(TCHAR));
    }

    Utility::dprintf(_T("<%ld> ReadMsgThread 退出\n"),Utility::g_WndInfo.dwProcessId);

    return 0;
}