[Redirector] Optimize code

This commit is contained in:
Connection Refused
2021-10-15 01:35:52 +08:00
parent 62154cf07b
commit 91b6692cfb
6 changed files with 91 additions and 144 deletions

View File

@@ -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();
}

View File

@@ -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;

View File

@@ -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

View File

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

View File

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

View File

@@ -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)