SimpleRemoter/server/2015Remote/IOCPServer.h

#pragma once

#include "StdAfx.h"
#include <WinSock2.h>
#pragma comment(lib,"ws2_32.lib")
#include "Server.h"
#include <Mstcpip.h>

#define	NC_CLIENT_CONNECT		0x0001
#define	NC_RECEIVE				0x0004
#define NC_RECEIVE_COMPLETE		0x0005 // 完整接收

// ZLIB 压缩库
#include "zlib/zlib.h"

inline int z_uncompress(z_stream* strm, Bytef* dest, uLongf* destLen, const Bytef* src, uLong srcLen)
{
    inflateReset(strm);

    strm->next_in = (Bytef*)src;
    strm->avail_in = srcLen;
    strm->next_out = dest;
    strm->avail_out = *destLen;

    int ret = inflate(strm, Z_FINISH);

    *destLen = strm->total_out;

    if (ret == Z_STREAM_END) return Z_OK;
    return ret;
}

// ZSTD
#include "zstd/zstd.h"
#ifdef _WIN64
#pragma comment(lib, "zstd/zstd_x64.lib")
#else
#pragma comment(lib, "zstd/zstd.lib")
#endif
#define C_FAILED(p) ZSTD_isError(p)
#define C_SUCCESS(p) (!C_FAILED(p))
#define ZSTD_CLEVEL 5

#define Mcompress(dest, destLen, source, sourceLen) m_Cctx ? ZSTD_compress2(m_Cctx, dest, *(destLen), source, sourceLen):\
	ZSTD_compress(dest, *(destLen), source, sourceLen, ZSTD_CLEVEL_DEFAULT)

#define Muncompress(dest, destLen, source, sourceLen) m_Dctx ? ZSTD_decompressDCtx(m_Dctx, dest, *(destLen), source, sourceLen):\
	ZSTD_decompress(dest, *(destLen), source, sourceLen)

class IOCPServer : public Server
{
protected:
    int					m_nPort;
    SOCKET				m_sListenSocket;
    HANDLE				m_hCompletionPort;
    UINT				m_ulMaxConnections;
    HANDLE				m_hListenEvent;
    HANDLE				m_hListenThread;
    BOOL				m_bTimeToKill;
    HANDLE				m_hKillEvent;
    ULONG				m_ulThreadPoolMin;
    ULONG				m_ulThreadPoolMax;
    ULONG				m_ulCPULowThreadsHold;
    ULONG				m_ulCPUHighThreadsHold;
    ULONG				m_ulCurrentThread;
    ULONG				m_ulBusyThread;

    ULONG				m_ulKeepLiveTime;
    pfnNotifyProc		m_NotifyProc;
    pfnOfflineProc		m_OfflineProc;
    ULONG				m_ulWorkThreadCount;
    CRITICAL_SECTION	m_cs;
    ContextObjectList	m_ContextConnectionList;
    ContextObjectList	m_ContextFreePoolList;

private:
    static DWORD WINAPI ListenThreadProc(LPVOID lParam);
    static DWORD WINAPI WorkThreadProc(LPVOID lParam);

    BOOL InitializeIOCP(VOID);
    VOID OnAccept();
    PCONTEXT_OBJECT AllocateContext(SOCKET s);
    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, z_stream *z);
    BOOL OnClientInitializing(PCONTEXT_OBJECT  ContextObject, DWORD dwTrans);
    BOOL OnClientReceiving(PCONTEXT_OBJECT  ContextObject, DWORD dwTrans, ZSTD_DCtx* ctx, z_stream* z);
    BOOL OnClientPreSending(CONTEXT_OBJECT* ContextObject, PBYTE szBuffer, size_t ulOriginalLength);
    BOOL OnClientPostSending(CONTEXT_OBJECT* ContextObject, ULONG ulCompressedLength);
    int AddWorkThread(int n)
    {
        EnterCriticalSection(&m_cs);
        m_ulWorkThreadCount += n;
        int ret = m_ulWorkThreadCount;
        LeaveCriticalSection(&m_cs);
        return ret;
    }

public:
    IOCPServer(void);
    ~IOCPServer(void);
    int GetPort() const override
    {
        return m_nPort;
    }

    UINT StartServer(pfnNotifyProc NotifyProc, pfnOfflineProc OffProc, USHORT uPort);

    BOOL Send2Client(CONTEXT_OBJECT* ContextObject, PBYTE szBuffer, ULONG ulOriginalLength) override
    {
        return OnClientPreSending(ContextObject, szBuffer, ulOriginalLength);
    }

    void UpdateMaxConnection(int maxConn);

    void Destroy();
    void Disconnect(CONTEXT_OBJECT *ctx) {}
};

typedef IOCPServer ISocketBase;

typedef IOCPServer CIOCPServer;

typedef CONTEXT_OBJECT ClientContext;

#define m_Socket sClientSocket
#define m_DeCompressionBuffer InDeCompressedBuffer

// 所有动态创建的对话框的基类
class CDialogBase : public CDialog
{
public:
    CONTEXT_OBJECT* m_ContextObject;
    Server* m_iocpServer;
    CString m_IPAddress;
    bool m_bIsClosed;
    bool m_bIsProcessing;
    HICON m_hIcon;
    BOOL m_bConnected;
    uint64_t m_nDisconnectTime = 0;
    CDialogBase(UINT nIDTemplate, CWnd* pParent, Server* pIOCPServer, CONTEXT_OBJECT* pContext, int nIcon) :
        m_bIsClosed(false), m_bIsProcessing(false),
        m_ContextObject(pContext),
        m_iocpServer(pIOCPServer),
        CDialog(nIDTemplate, pParent)
    {
        m_bConnected = TRUE;
		m_nDisconnectTime = 0;
        m_IPAddress = pContext->GetPeerName().c_str();
        m_hIcon = nIcon > 0 ? LoadIcon(AfxGetInstanceHandle(), MAKEINTRESOURCE(nIcon)) : NULL;
    }
    int UpdateContext(CONTEXT_OBJECT* pContext)
    {
        m_bConnected = TRUE;
        m_nDisconnectTime = 0;
        m_ContextObject = pContext;
        m_iocpServer = pContext->GetServer();
        m_ContextObject->hDlg = this;
        return 0;
    }
    virtual ~CDialogBase() {}

public:
    virtual BOOL ReceiveCommonMsg()
    {
        switch (m_ContextObject->InDeCompressedBuffer.GetBYTE(0)) {
        case TOKEN_CLIENT_MSG: {
            ClientMsg* msg = (ClientMsg*)m_ContextObject->InDeCompressedBuffer.GetBuffer(0);
            PostMessageA(WM_SHOWERRORMSG, (WPARAM)new CString(msg->text), (LPARAM)new CString(msg->title));
            return TRUE;
        }
        }
        return FALSE;
    }
    virtual void OnReceiveComplete(void) = 0;
    // 标记为是否正在接受数据
    void MarkReceiving(bool recv = true)
    {
        m_bIsProcessing = recv;
    }
    bool IsProcessing() const
    {
        return m_bIsProcessing;
    }
    void OnClose()
    {
        m_bIsClosed = true;
        m_bConnected = FALSE;
        while (m_bIsProcessing)
            Sleep(200);
        if(m_hIcon) DestroyIcon(m_hIcon);
        m_hIcon = NULL;
        CDialog::OnClose();

        if (GetSafeHwnd())
            DestroyWindow();
    }
    virtual void PostNcDestroy() override
    {
        delete this;
    }
    virtual BOOL ShouldReconnect()
    {
        return FALSE;
    }
    // 取消 SOCKET 读取，该函数可以被多次调用
    void CancelIO()
    {
        m_bIsClosed = TRUE;

        m_ContextObject->CancelIO();
    }
    BOOL IsClosed() const
    {
        return m_bIsClosed;
    }
    BOOL SayByeBye()
    {
        if (!m_bConnected) return FALSE;

        Mprintf("%s SayByeBye: %s\n", ToPekingTimeAsString(0).c_str(), m_ContextObject->GetPeerName().c_str());
        BYTE bToken = COMMAND_BYE;
        return m_ContextObject->Send2Client(&bToken, 1);
    }
};

typedef CDialogBase DialogBase;

BOOL ParseReceivedData(CONTEXT_OBJECT* ContextObject, DWORD dwTrans, pfnNotifyProc m_NotifyProc, ZSTD_DCtx *ctx=NULL, z_stream* z=NULL);

BOOL WriteContextData(CONTEXT_OBJECT* ContextObject, PBYTE szBuffer, size_t ulOriginalLength, ZSTD_CCtx *ctx=NULL, z_stream* z = NULL);