mirror of
https://github.com/netchx/netch.git
synced 2026-03-14 17:43:18 +08:00
[Redirector] Optimize code
This commit is contained in:
@@ -5,38 +5,31 @@ extern USHORT dnsPort;
|
||||
|
||||
SOCKADDR_IN6 dnsAddr;
|
||||
|
||||
void ProcessPacket(ENDPOINT_ID id, SOCKADDR_IN6 target, char* packet, int length, PNF_UDP_OPTIONS option)
|
||||
void HandleDNS(ENDPOINT_ID id, PSOCKADDR_IN6 target, char* packet, int length, PNF_UDP_OPTIONS option)
|
||||
{
|
||||
char buffer[1024];
|
||||
|
||||
auto tcpSocket = SocksHelper::Utils::Connect();
|
||||
if (tcpSocket != INVALID_SOCKET)
|
||||
auto remote = new SocksHelper::UDP();
|
||||
if (remote->Associate())
|
||||
{
|
||||
if (SocksHelper::Utils::Handshake(tcpSocket))
|
||||
if (remote->CreateUDP())
|
||||
{
|
||||
SocksHelper::UDP udpConn;
|
||||
udpConn.tcpSocket = tcpSocket;
|
||||
|
||||
if (udpConn.Associate())
|
||||
if (remote->Send(&dnsAddr, packet, length) == length)
|
||||
{
|
||||
if (udpConn.CreateUDP())
|
||||
{
|
||||
if (udpConn.Send(&dnsAddr, packet, length) == length)
|
||||
{
|
||||
timeval timeout{};
|
||||
timeout.tv_sec = 4;
|
||||
timeval timeout{};
|
||||
timeout.tv_sec = 4;
|
||||
|
||||
int size = udpConn.Read(NULL, buffer, sizeof(buffer), &timeout);
|
||||
if (size != 0 && size != SOCKET_ERROR)
|
||||
{
|
||||
nf_udpPostReceive(id, (unsigned char*)&target, buffer, size, option);
|
||||
}
|
||||
}
|
||||
int size = remote->Read(NULL, buffer, sizeof(buffer), &timeout);
|
||||
if (size != 0 && size != SOCKET_ERROR)
|
||||
{
|
||||
nf_udpPostReceive(id, (unsigned char*)target, buffer, size, option);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
delete remote;
|
||||
delete target;
|
||||
delete[] packet;
|
||||
delete[] option;
|
||||
}
|
||||
@@ -77,13 +70,13 @@ bool DNSHandler::IsDNS(PSOCKADDR_IN6 target)
|
||||
|
||||
void DNSHandler::CreateHandler(ENDPOINT_ID id, PSOCKADDR_IN6 target, const char* packet, int length, PNF_UDP_OPTIONS options)
|
||||
{
|
||||
SOCKADDR_IN6 remote;
|
||||
auto remote = new SOCKADDR_IN6();
|
||||
auto buffer = new char[length]();
|
||||
auto option = (PNF_UDP_OPTIONS)new char[sizeof(NF_UDP_OPTIONS) + options->optionsLength];
|
||||
|
||||
memcpy(&remote, target, sizeof(SOCKADDR_IN6));
|
||||
memcpy(remote, target, sizeof(SOCKADDR_IN6));
|
||||
memcpy(buffer, packet, length);
|
||||
memcpy(option, options, sizeof(NF_UDP_OPTIONS) + options->optionsLength - 1);
|
||||
|
||||
thread(ProcessPacket, id, remote, buffer, length, option).detach();
|
||||
thread(HandleDNS, id, remote, buffer, length, option).detach();
|
||||
}
|
||||
|
||||
@@ -329,66 +329,38 @@ void udpSend(ENDPOINT_ID id, const unsigned char* target, const char* buffer, in
|
||||
nf_udpPostSend(id, target, buffer, length, options);
|
||||
return;
|
||||
}
|
||||
auto udpConn = udpContext[id];
|
||||
auto remote = udpContext[id];
|
||||
udpContextLock.unlock();
|
||||
|
||||
UP += length;
|
||||
|
||||
if (udpConn->tcpSocket == INVALID_SOCKET)
|
||||
if (remote->tcpSocket == INVALID_SOCKET || remote->udpSocket == INVALID_SOCKET)
|
||||
{
|
||||
auto tcpSocket = SocksHelper::Utils::Connect();
|
||||
if (tcpSocket == INVALID_SOCKET)
|
||||
{
|
||||
printf("[Redirector][EventHandler][udpSend][%llu] Connect to remote server failed\n", id);
|
||||
if (remote->tcpSocket) closesocket(remote->tcpSocket);
|
||||
if (remote->udpSocket) closesocket(remote->udpSocket);
|
||||
|
||||
remote->tcpSocket = INVALID_SOCKET;
|
||||
remote->udpSocket = INVALID_SOCKET;
|
||||
|
||||
if (!remote->Associate())
|
||||
return;
|
||||
}
|
||||
|
||||
if (!SocksHelper::Utils::Handshake(tcpSocket))
|
||||
{
|
||||
closesocket(tcpSocket);
|
||||
|
||||
printf("[Redirector][EventHandler][udpSend][%llu] Handshake failed\n", id);
|
||||
if (!remote->CreateUDP())
|
||||
return;
|
||||
}
|
||||
|
||||
udpConn->tcpSocket = tcpSocket;
|
||||
}
|
||||
|
||||
if (udpConn->udpSocket == INVALID_SOCKET)
|
||||
{
|
||||
if (!udpConn->Associate())
|
||||
{
|
||||
closesocket(udpConn->tcpSocket);
|
||||
udpConn->tcpSocket = INVALID_SOCKET;
|
||||
|
||||
printf("[Redirector][EventHandler][udpSend][%llu] UDP Associate failed\n", id);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!udpConn->CreateUDP())
|
||||
{
|
||||
closesocket(udpConn->tcpSocket);
|
||||
udpConn->tcpSocket = INVALID_SOCKET;
|
||||
|
||||
printf("[Redirector][EventHandler][udpSend][%llu] Create UDP socket failed\n", id);
|
||||
return;
|
||||
}
|
||||
|
||||
auto data = (PNF_UDP_OPTIONS)new char[sizeof(NF_UDP_OPTIONS) + options->optionsLength]();
|
||||
memcpy(data, options, sizeof(NF_UDP_OPTIONS) + options->optionsLength - 1);
|
||||
|
||||
thread(udpBeginReceive, id, udpConn, data).detach();
|
||||
thread(udpReceiveHandler, id, remote, data).detach();
|
||||
}
|
||||
|
||||
if (udpConn->Send((PSOCKADDR_IN6)target, buffer, length) != length)
|
||||
if (remote->Send((PSOCKADDR_IN6)target, buffer, length) != length)
|
||||
{
|
||||
closesocket(udpConn->tcpSocket);
|
||||
closesocket(udpConn->udpSocket);
|
||||
if (remote->tcpSocket) closesocket(remote->tcpSocket);
|
||||
if (remote->udpSocket) closesocket(remote->udpSocket);
|
||||
|
||||
udpConn->tcpSocket = INVALID_SOCKET;
|
||||
udpConn->udpSocket = INVALID_SOCKET;
|
||||
|
||||
printf("[Redirector][EventHandler][udpSend][%llu] Send data failed\n", id);
|
||||
remote->tcpSocket = INVALID_SOCKET;
|
||||
remote->udpSocket = INVALID_SOCKET;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -417,19 +389,17 @@ void udpClosed(ENDPOINT_ID id, PNF_UDP_CONN_INFO info)
|
||||
}
|
||||
}
|
||||
|
||||
void udpBeginReceive(ENDPOINT_ID id, SocksHelper::PUDP udpConn, PNF_UDP_OPTIONS options)
|
||||
void udpReceiveHandler(ENDPOINT_ID id, SocksHelper::PUDP remote, PNF_UDP_OPTIONS options)
|
||||
{
|
||||
char buffer[1458];
|
||||
|
||||
while (udpConn->udpSocket != INVALID_SOCKET)
|
||||
while (remote->tcpSocket != INVALID_SOCKET && remote->udpSocket != INVALID_SOCKET)
|
||||
{
|
||||
SOCKADDR_IN6 target;
|
||||
|
||||
int length = udpConn->Read(&target, buffer, sizeof(buffer), NULL);
|
||||
int length = remote->Read(&target, buffer, sizeof(buffer), NULL);
|
||||
if (length == 0 || length == SOCKET_ERROR)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
DL += length;
|
||||
|
||||
|
||||
@@ -24,6 +24,6 @@ void udpCanReceive(ENDPOINT_ID id);
|
||||
void udpReceive(ENDPOINT_ID id, const unsigned char* target, const char* buffer, int length, PNF_UDP_OPTIONS options);
|
||||
void udpClosed(ENDPOINT_ID id, PNF_UDP_CONN_INFO info);
|
||||
|
||||
void udpBeginReceive(ENDPOINT_ID id, SocksHelper::PUDP conn, PNF_UDP_OPTIONS options);
|
||||
void udpReceiveHandler(ENDPOINT_ID id, SocksHelper::PUDP conn, PNF_UDP_OPTIONS options);
|
||||
|
||||
#endif
|
||||
|
||||
@@ -7,7 +7,7 @@ extern wstring tgtPort;
|
||||
extern string tgtUsername;
|
||||
extern string tgtPassword;
|
||||
|
||||
SOCKET SocksHelper::Utils::Connect()
|
||||
SOCKET SocksHelper::Connect()
|
||||
{
|
||||
auto client = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
|
||||
if (client == INVALID_SOCKET)
|
||||
@@ -27,7 +27,7 @@ SOCKET SocksHelper::Utils::Connect()
|
||||
}
|
||||
}
|
||||
|
||||
timeval timeout;
|
||||
timeval timeout{};
|
||||
timeout.tv_sec = 4;
|
||||
|
||||
if (!WSAConnectByNameW(client, (LPWSTR)tgtHost.c_str(), (LPWSTR)tgtPort.c_str(), NULL, NULL, NULL, NULL, &timeout, NULL))
|
||||
@@ -41,7 +41,7 @@ SOCKET SocksHelper::Utils::Connect()
|
||||
return client;
|
||||
}
|
||||
|
||||
bool SocksHelper::Utils::Handshake(SOCKET client)
|
||||
bool SocksHelper::Handshake(SOCKET client)
|
||||
{
|
||||
char buffer[1024];
|
||||
memset(buffer, 0, sizeof(buffer));
|
||||
@@ -117,7 +117,7 @@ bool SocksHelper::Utils::Handshake(SOCKET client)
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SocksHelper::Utils::SplitAddr(SOCKET client, PSOCKADDR_IN6 addr)
|
||||
bool SocksHelper::SplitAddr(SOCKET client, PSOCKADDR_IN6 addr)
|
||||
{
|
||||
char addressType;
|
||||
if (recv(client, (char*)&addressType, 1, 0) != 1)
|
||||
@@ -180,6 +180,13 @@ SocksHelper::TCP::~TCP()
|
||||
|
||||
bool SocksHelper::TCP::Connect(PSOCKADDR_IN6 target)
|
||||
{
|
||||
this->tcpSocket = SocksHelper::Connect();
|
||||
if (this->tcpSocket == INVALID_SOCKET)
|
||||
return false;
|
||||
|
||||
if (!SocksHelper::Handshake(this->tcpSocket))
|
||||
return false;
|
||||
|
||||
/* Connect Request */
|
||||
if (target->sin6_family == AF_INET)
|
||||
{
|
||||
@@ -230,7 +237,7 @@ bool SocksHelper::TCP::Connect(PSOCKADDR_IN6 target)
|
||||
return false;
|
||||
|
||||
SOCKADDR_IN6 addr;
|
||||
return SocksHelper::Utils::SplitAddr(this->tcpSocket, &addr);
|
||||
return SocksHelper::SplitAddr(this->tcpSocket, &addr);
|
||||
}
|
||||
|
||||
int SocksHelper::TCP::Send(const char* buffer, int length)
|
||||
@@ -264,18 +271,45 @@ SocksHelper::UDP::~UDP()
|
||||
|
||||
this->udpSocket = INVALID_SOCKET;
|
||||
}
|
||||
}
|
||||
|
||||
if (this->tcpThread.joinable())
|
||||
void SocksHelper::UDP::Run(SOCKET tcpSocket, SOCKET udpSocket)
|
||||
{
|
||||
char buffer[1];
|
||||
|
||||
while (tcpSocket != INVALID_SOCKET)
|
||||
{
|
||||
this->tcpThread.join();
|
||||
if (recv(tcpSocket, buffer, sizeof(buffer), 0) != sizeof(buffer))
|
||||
break;
|
||||
|
||||
if (send(tcpSocket, buffer, sizeof(buffer), 0) != sizeof(buffer))
|
||||
break;
|
||||
}
|
||||
|
||||
if (tcpSocket != INVALID_SOCKET)
|
||||
{
|
||||
closesocket(tcpSocket);
|
||||
|
||||
tcpSocket = INVALID_SOCKET;
|
||||
}
|
||||
|
||||
if (udpSocket != INVALID_SOCKET)
|
||||
{
|
||||
closesocket(udpSocket);
|
||||
|
||||
udpSocket = INVALID_SOCKET;
|
||||
}
|
||||
}
|
||||
|
||||
bool SocksHelper::UDP::Associate()
|
||||
{
|
||||
this->tcpSocket = SocksHelper::Connect();
|
||||
if (this->tcpSocket == INVALID_SOCKET)
|
||||
return false;
|
||||
|
||||
if (!SocksHelper::Handshake(this->tcpSocket))
|
||||
return false;
|
||||
|
||||
char buffer[10];
|
||||
buffer[0] = 0x05;
|
||||
buffer[1] = 0x03;
|
||||
@@ -306,7 +340,7 @@ bool SocksHelper::UDP::Associate()
|
||||
return false;
|
||||
}
|
||||
|
||||
return SocksHelper::Utils::SplitAddr(this->tcpSocket, &this->address);
|
||||
return SocksHelper::SplitAddr(this->tcpSocket, &this->address);
|
||||
}
|
||||
|
||||
bool SocksHelper::UDP::CreateUDP()
|
||||
@@ -350,8 +384,7 @@ bool SocksHelper::UDP::CreateUDP()
|
||||
}
|
||||
}
|
||||
|
||||
this->tcpThread = thread(SocksHelper::UDP::Run, this->tcpSocket, this->udpSocket);
|
||||
this->tcpThread.detach();
|
||||
thread(SocksHelper::UDP::Run, this->tcpSocket, this->udpSocket).detach();
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -444,31 +477,3 @@ int SocksHelper::UDP::Read(PSOCKADDR_IN6 target, char* buffer, int length, PTIME
|
||||
|
||||
return bufferLength - (addr.sin6_family == AF_INET ? 10 : 22);
|
||||
}
|
||||
|
||||
void SocksHelper::UDP::Run(SOCKET tcpSocket, SOCKET udpSocket)
|
||||
{
|
||||
char buffer[1];
|
||||
|
||||
while (tcpSocket != INVALID_SOCKET)
|
||||
{
|
||||
if (recv(tcpSocket, buffer, sizeof(buffer), 0) != sizeof(buffer))
|
||||
break;
|
||||
|
||||
if (send(tcpSocket, buffer, sizeof(buffer), 0) != sizeof(buffer))
|
||||
break;
|
||||
}
|
||||
|
||||
if (tcpSocket != INVALID_SOCKET)
|
||||
{
|
||||
closesocket(tcpSocket);
|
||||
|
||||
tcpSocket = INVALID_SOCKET;
|
||||
}
|
||||
|
||||
if (udpSocket != INVALID_SOCKET)
|
||||
{
|
||||
closesocket(udpSocket);
|
||||
|
||||
udpSocket = INVALID_SOCKET;
|
||||
}
|
||||
}
|
||||
@@ -5,12 +5,9 @@
|
||||
|
||||
namespace SocksHelper
|
||||
{
|
||||
namespace Utils
|
||||
{
|
||||
SOCKET Connect();
|
||||
bool Handshake(SOCKET client);
|
||||
bool SplitAddr(SOCKET client, PSOCKADDR_IN6 addr);
|
||||
}
|
||||
SOCKET Connect();
|
||||
bool Handshake(SOCKET client);
|
||||
bool SplitAddr(SOCKET client, PSOCKADDR_IN6 addr);
|
||||
|
||||
typedef class TCP
|
||||
{
|
||||
@@ -30,6 +27,8 @@ namespace SocksHelper
|
||||
public:
|
||||
~UDP();
|
||||
|
||||
static void Run(SOCKET tcpSocket, SOCKET udpSocket);
|
||||
|
||||
bool Associate();
|
||||
bool CreateUDP();
|
||||
|
||||
@@ -38,11 +37,8 @@ namespace SocksHelper
|
||||
|
||||
SOCKET tcpSocket = INVALID_SOCKET;
|
||||
SOCKET udpSocket = INVALID_SOCKET;
|
||||
private:
|
||||
static void Run(SOCKET tcpSocket, SOCKET udpSocket);
|
||||
|
||||
SOCKADDR_IN6 address = { 0 };
|
||||
thread tcpThread;
|
||||
} *PUDP;
|
||||
};
|
||||
|
||||
|
||||
@@ -144,20 +144,6 @@ void TCPHandler::Handle(SOCKET client)
|
||||
id = (addr.sin6_family == AF_INET6) ? addr.sin6_port : ((PSOCKADDR_IN)&addr)->sin_port;
|
||||
}
|
||||
|
||||
auto remote = SocksHelper::Utils::Connect();
|
||||
if (remote == INVALID_SOCKET)
|
||||
{
|
||||
closesocket(client);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!SocksHelper::Utils::Handshake(remote))
|
||||
{
|
||||
closesocket(client);
|
||||
closesocket(remote);
|
||||
return;
|
||||
}
|
||||
|
||||
tcpLock.lock();
|
||||
if (tcpContext.find(id) == tcpContext.end())
|
||||
{
|
||||
@@ -170,23 +156,20 @@ void TCPHandler::Handle(SOCKET client)
|
||||
auto target = tcpContext[id];
|
||||
tcpLock.unlock();
|
||||
|
||||
auto conn = new SocksHelper::TCP();
|
||||
conn->tcpSocket = remote;
|
||||
|
||||
if (!conn->Connect(&target))
|
||||
auto remote = new SocksHelper::TCP();
|
||||
if (!remote->Connect(&target))
|
||||
{
|
||||
delete conn;
|
||||
|
||||
closesocket(client);
|
||||
|
||||
delete remote;
|
||||
return;
|
||||
}
|
||||
|
||||
thread(TCPHandler::Send, client, conn).detach();
|
||||
TCPHandler::Read(client, conn);
|
||||
thread(TCPHandler::Send, client, remote).detach();
|
||||
TCPHandler::Read(client, remote);
|
||||
|
||||
closesocket(client);
|
||||
closesocket(remote);
|
||||
delete conn;
|
||||
delete remote;
|
||||
}
|
||||
|
||||
void TCPHandler::Read(SOCKET client, SocksHelper::PTCP remote)
|
||||
|
||||
Reference in New Issue
Block a user