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>
#include "LangManager.h"
#include <map>
#include <set>
#include <string>

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

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;
    HWND                m_hMainWnd = nullptr;

    // IP 连接限流和封禁
    struct ConnectionInfo {
        int count;          // 连接次数
        time_t windowStart; // 统计窗口起始时间
    };
    std::map<std::string, ConnectionInfo> m_ConnectionCount;  // IP -> 连接统计
    std::map<std::string, time_t> m_BannedIPs;                // IP -> 封禁到期时间
    CRITICAL_SECTION m_BanLock;
    // 白名单已移至 IPWhitelist 单例 (common/IPWhitelist.h)

    bool IsIPBanned(const std::string& ip);
    bool IsIPBlacklisted(const std::string& ip);
    void RecordConnection(const std::string& ip);
    void BanIP(const std::string& ip, int seconds);
    void LoadIPWhitelist();
    void LoadIPBlacklist();

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(HWND hWnd = nullptr);
    ~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 CDialogLang
{
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_ClientID = 0;
    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),
        CDialogLang(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, uint64_t clientID)
    {
        if (m_bIsClosed) {
            Mprintf("%s SayByeBye: %llu [Already Closed]\n", ToPekingTimeAsString(0).c_str(), clientID);
            BYTE bToken = COMMAND_BYE;
            return m_ContextObject->Send2Client(&bToken, 1) ? 0 : 0x20260223;
        }
        m_ClientID = clientID;
        m_bConnected = TRUE;
        m_nDisconnectTime = 0;
        m_ContextObject = pContext;
        m_iocpServer = pContext->GetServer();
        m_ContextObject->hDlg = this;
        m_ContextObject->hWnd = GetSafeHwnd();
        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(_L(msg->text)), (LPARAM)new CString(_L(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;
        __super::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;
    }
    virtual uint64_t GetClientID() const {
        // 优先用 UpdateContext 设过的 m_ClientID（重连场景），否则取子连接 ctx 自身的 ID。
        // 子连接通过 TOKEN_CONN_AUTH 通过校验后，ctx->GetClientID() 已被钉成主连接的 clientID，
        // 这样 dialog 拿到的 ID 既准确又免去 IP 反查兜底（NAT/127.0.0.1 场景靠谱）。
        if (m_ClientID != 0) return m_ClientID;
        return m_ContextObject ? m_ContextObject->GetClientID() : 0;
    }
    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);