From 91b6692cfbebc38e6cfc6393a92fe2e3bdd561c0 Mon Sep 17 00:00:00 2001 From: Connection Refused Date: Fri, 15 Oct 2021 01:35:52 +0800 Subject: [PATCH] [Redirector] Optimize code --- Redirector/DNSHandler.cpp | 39 ++++++++---------- Redirector/EventHandler.cpp | 68 +++++++++---------------------- Redirector/EventHandler.h | 2 +- Redirector/SocksHelper.cpp | 81 ++++++++++++++++++++----------------- Redirector/SocksHelper.h | 14 +++---- Redirector/TCPHandler.cpp | 31 ++++---------- 6 files changed, 91 insertions(+), 144 deletions(-) diff --git a/Redirector/DNSHandler.cpp b/Redirector/DNSHandler.cpp index edd5c25d..17f5ee6b 100644 --- a/Redirector/DNSHandler.cpp +++ b/Redirector/DNSHandler.cpp @@ -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(); } diff --git a/Redirector/EventHandler.cpp b/Redirector/EventHandler.cpp index ab4605e8..d0a0925e 100644 --- a/Redirector/EventHandler.cpp +++ b/Redirector/EventHandler.cpp @@ -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; diff --git a/Redirector/EventHandler.h b/Redirector/EventHandler.h index dfc70e98..f84c70a2 100644 --- a/Redirector/EventHandler.h +++ b/Redirector/EventHandler.h @@ -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 diff --git a/Redirector/SocksHelper.cpp b/Redirector/SocksHelper.cpp index de82c8c0..119b6271 100644 --- a/Redirector/SocksHelper.cpp +++ b/Redirector/SocksHelper.cpp @@ -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; - } -} \ No newline at end of file diff --git a/Redirector/SocksHelper.h b/Redirector/SocksHelper.h index 5f151718..8e2dfab4 100644 --- a/Redirector/SocksHelper.h +++ b/Redirector/SocksHelper.h @@ -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; }; diff --git a/Redirector/TCPHandler.cpp b/Redirector/TCPHandler.cpp index 4584e1e1..e4db5605 100644 --- a/Redirector/TCPHandler.cpp +++ b/Redirector/TCPHandler.cpp @@ -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)