SimpleRemoter/client/IOCPClient.h

// IOCPClient.h: interface for the IOCPClient class.
//
//////////////////////////////////////////////////////////////////////

#pragma once

#ifdef _WIN32
#include "stdafx.h"
#include <WinSock2.h>
#include <MSTcpIP.h>
#pragma comment(lib,"ws2_32.lib")
#endif

#include "Buffer.h"
#include "zstd/zstd.h"
#include "domain_pool.h"
#include "common/mask.h"
#include "common/header.h"
#define NO_AES
#include "common/encrypt.h"
#include "SafeThread.h"
#include "IOCPBase.h"

#define MAX_RECV_BUFFER  1024*32
#define MAX_SEND_BUFFER  1024*32

enum { S_STOP = 0, S_RUN, S_END };

typedef int (*DataProcessCB)(void* userData, PBYTE szBuffer, ULONG ulLength);


class ProtocolEncoder {
public:
	virtual ~ProtocolEncoder(){}
	virtual HeaderFlag GetHead() const {
		return "Shine";
	}
	virtual int GetHeadLen() const {
		return 13;
	}
	virtual int GetFlagLen() const {
		return 5;
	}
	virtual void Encode(unsigned char* data, int len, unsigned char* param = 0) {}
	virtual void Decode(unsigned char* data, int len, unsigned char* param = 0) {}
	virtual EncFun GetHeaderEncoder() const { return nullptr; }
};

class HellEncoder : public ProtocolEncoder {
private:
	EncFun m_HeaderEnc;
	Encoder *m_BodyEnc;
public:
	HellEncoder(EncFun head, Encoder *body) {
		m_HeaderEnc = head;
		m_BodyEnc = body;
	}
	~HellEncoder() {
		SAFE_DELETE(m_BodyEnc);
	}
	virtual HeaderFlag GetHead() const override {
		return ::GetHead(m_HeaderEnc);
	}
	virtual int GetHeadLen() const override {
		return 16;
	}
	virtual int GetFlagLen() const override {
		return 8;
	}
	virtual void Encode(unsigned char* data, int len, unsigned char* param = 0) override {
		return m_BodyEnc->Encode(data, len, param);
	}
	virtual void Decode(unsigned char* data, int len, unsigned char* param = 0) override {
		return m_BodyEnc->Decode(data, len, param);
	}
	virtual EncFun GetHeaderEncoder() const override {
		return m_HeaderEnc;
	}
};

class IOCPManager {
public:
	virtual ~IOCPManager() {}
	virtual BOOL IsAlive() const { return TRUE; }
	virtual BOOL IsReady() const { return TRUE; }
	virtual VOID OnReceive(PBYTE szBuffer, ULONG ulLength) { }

	static int DataProcess(void* user, PBYTE szBuffer, ULONG ulLength) {
		IOCPManager* m_Manager = (IOCPManager*)user;
		if (nullptr == m_Manager) {
			Mprintf("IOCPManager DataProcess on NULL ptr: %d\n", unsigned(szBuffer[0]));
			return FALSE;
		}
		// <20>ȴ<EFBFBD><C8B4><EFBFBD><EFBFBD><EFBFBD>׼<EFBFBD><D7BC><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ܴ<EFBFBD><DCB4><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>, 1<><31><EFBFBD>㹻<EFBFBD><E3B9BB>
		int i = 0;
		for (; i < 1000 && !m_Manager->IsReady(); ++i)
			Sleep(1);
		if (!m_Manager->IsReady()) {
			Mprintf("IOCPManager DataProcess is NOT ready: %d\n", unsigned(szBuffer[0]));
			return FALSE;
		}
		if (i) {
			Mprintf("IOCPManager DataProcess wait for %dms: %d\n", i, unsigned(szBuffer[0]));
		}
		m_Manager->OnReceive(szBuffer, ulLength);
		return TRUE;
	}
};

typedef BOOL(*TrailCheck)(void);

class IOCPClient : public IOCPBase
{
public:
	IOCPClient(const State& bExit, bool exit_while_disconnect = false, int mask=0, int encoder=0, 
		const std::string&pubIP="");
	virtual ~IOCPClient();

	int SendLoginInfo(const LOGIN_INFOR& logInfo) {
		LOGIN_INFOR tmp = logInfo;
		int iRet = Send2Server((char*)&tmp, sizeof(LOGIN_INFOR));

		return iRet;
	}
	virtual BOOL ConnectServer(const char* szServerIP, unsigned short uPort);

	std::string GetClientIP() const {
		return m_sLocPublicIP;
	}

	std::map<std::string, std::string> GetClientIPHeader() const {
		return m_sLocPublicIP.empty() ? std::map<std::string, std::string>{} : 
			std::map<std::string, std::string>{ {"X-Forwarded-For", m_sLocPublicIP} };
	}

	BOOL Send2Server(const char* szBuffer, ULONG ulOriginalLength, PkgMask* mask = NULL) {
		return OnServerSending(szBuffer, ulOriginalLength, mask);
	}

	void SetServerAddress(const char* szServerIP, unsigned short uPort) {
		m_Domain = szServerIP ? szServerIP : "127.0.0.1";
		m_nHostPort = uPort;
	}

	BOOL IsRunning() const
	{
		return m_bIsRunning;
	}

	void SetExit() {
		m_bIsRunning = FALSE;
	}

	VOID setManagerCallBack(void* Manager, DataProcessCB dataProcess);
	VOID RunEventLoop(TrailCheck checker);
	VOID RunEventLoop(const BOOL &bCondition);
	bool IsConnected() const { return m_bConnected == TRUE; }
	BOOL Reconnect(void* manager) {
		Disconnect();
		if (manager) m_Manager = manager;
		return ConnectServer(NULL, 0);
	}
	const State& GetState() const {
		return g_bExit;
	}
protected:
	virtual int ReceiveData(char* buffer, int bufSize, int flags) {
		// TCP<43>汾<EFBFBD><E6B1BE><EFBFBD><EFBFBD> recv
		return recv(m_sClientSocket, buffer, bufSize - 1, 0);
	}
	virtual bool ProcessRecvData(CBuffer* m_CompressedBuffer, char* szBuffer, int len, int flag);
	virtual VOID Disconnect(); // <20><><EFBFBD><EFBFBD>֧<EFBFBD><D6A7> TCP/UDP
	virtual int SendTo(const char* buf, int len, int flags) {
		return ::send(m_sClientSocket, buf, len, flags);
	}
	BOOL OnServerSending(const char* szBuffer, ULONG ulOriginalLength, PkgMask* mask);
	static DWORD WINAPI WorkThreadProc(LPVOID lParam);
	VOID OnServerReceiving(CBuffer *m_CompressedBuffer, char* szBuffer, ULONG ulReceivedLength);
	BOOL SendWithSplit(const char* src, ULONG srcSize, ULONG ulSplitLength, int cmd, PkgMask* mask);

protected:
	sockaddr_in			m_ServerAddr;
	SOCKET				m_sClientSocket;
	BOOL				m_bWorkThread;
	HANDLE				m_hWorkThread;
	BOOL				m_bIsRunning;
	BOOL				m_bConnected;

#if USING_CTX
	ZSTD_CCtx*			m_Cctx;						// ѹ<><D1B9><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
	ZSTD_DCtx*			m_Dctx;						// <20><>ѹ<EFBFBD><D1B9><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
#endif

	const State&		g_bExit;					// ȫ<><C8AB>״̬<D7B4><CCAC>
	void*				m_Manager;					// <20>û<EFBFBD><C3BB><EFBFBD><EFBFBD><EFBFBD>
	DataProcessCB		m_DataProcess;				// <20><><EFBFBD><EFBFBD><EFBFBD>û<EFBFBD><C3BB><EFBFBD><EFBFBD><EFBFBD>
	ProtocolEncoder*	m_Encoder;					// <20><><EFBFBD><EFBFBD>
	DomainPool			m_Domain;
	std::string			m_sCurIP;
	int					m_nHostPort;
	bool				m_exit_while_disconnect;
	PkgMask*			m_masker;
	BOOL				m_EncoderType;
	std::string			m_sLocPublicIP;
};