263 lines
8.5 KiB
C++
263 lines
8.5 KiB
C++
#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();
|
||
|
||
// RTT 反代理（试用版执法）相关。详见 IOCPServer.cpp 中的实现注释。
|
||
HANDLE m_hRttThread = NULL;
|
||
bool m_bTrialMode = false; // StartServer 时根据 IsTrail(passcode) 缓存
|
||
std::atomic<bool> m_bSioTcpInfoProbed{false}; // 是否已完成 OS 兼容性探测
|
||
std::atomic<bool> m_bSioTcpInfoSupported{false}; // 探测结果（不支持则 RTT 线程会自行退出）
|
||
// 全 server 进程一次性 latch，IP 段触发与 RTT 触发共用。任一先触发后另一者只记日志不再弹框，
|
||
// 避免对运营商反复打扰；不影响每条 abusive 连接的独立日志。
|
||
static std::atomic<bool> s_TrialAbuseWarned;
|
||
|
||
private:
|
||
static DWORD WINAPI ListenThreadProc(LPVOID lParam);
|
||
static DWORD WINAPI WorkThreadProc(LPVOID lParam);
|
||
static DWORD WINAPI RttPollThreadProc(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);
|