[Redirector] Optimize code

This commit is contained in:
Connection Refused
2021-10-15 01:54:39 +08:00
parent 91b6692cfb
commit 19ff231f95
4 changed files with 47 additions and 57 deletions

View File

@@ -7,8 +7,6 @@ SOCKADDR_IN6 dnsAddr;
void HandleDNS(ENDPOINT_ID id, PSOCKADDR_IN6 target, char* packet, int length, PNF_UDP_OPTIONS option)
{
char buffer[1024];
auto remote = new SocksHelper::UDP();
if (remote->Associate())
{
@@ -16,14 +14,14 @@ void HandleDNS(ENDPOINT_ID id, PSOCKADDR_IN6 target, char* packet, int length, P
{
if (remote->Send(&dnsAddr, packet, length) == length)
{
char buffer[1024];
timeval timeout{};
timeout.tv_sec = 4;
int size = remote->Read(NULL, buffer, sizeof(buffer), &timeout);
if (size != 0 && size != SOCKET_ERROR)
{
nf_udpPostReceive(id, (unsigned char*)target, buffer, size, option);
}
}
}
}
@@ -54,7 +52,6 @@ bool DNSHandler::INIT()
return true;
}
cout << "[Redirector][DNSHandler::INIT] Convert address failed: " << WSAGetLastError() << " [" << dnsHost << ":" << dnsPort << "]" << endl;
return false;
}
@@ -64,8 +61,10 @@ bool DNSHandler::IsDNS(PSOCKADDR_IN6 target)
{
return ((PSOCKADDR_IN)target)->sin_port == htons(53);
}
return target->sin6_port == htons(53);
else
{
return target->sin6_port == htons(53);
}
}
void DNSHandler::CreateHandler(ENDPOINT_ID id, PSOCKADDR_IN6 target, const char* packet, int length, PNF_UDP_OPTIONS options)

View File

@@ -348,10 +348,10 @@ void udpSend(ENDPOINT_ID id, const unsigned char* target, const char* buffer, in
if (!remote->CreateUDP())
return;
auto data = (PNF_UDP_OPTIONS)new char[sizeof(NF_UDP_OPTIONS) + options->optionsLength]();
memcpy(data, options, sizeof(NF_UDP_OPTIONS) + options->optionsLength - 1);
auto option = (PNF_UDP_OPTIONS)new char[sizeof(NF_UDP_OPTIONS) + options->optionsLength]();
memcpy(option, options, sizeof(NF_UDP_OPTIONS) + options->optionsLength - 1);
thread(udpReceiveHandler, id, remote, data).detach();
thread(udpReceiveHandler, id, remote, option).detach();
}
if (remote->Send((PSOCKADDR_IN6)target, buffer, length) != length)

View File

@@ -12,7 +12,7 @@ SOCKET SocksHelper::Connect()
auto client = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
if (client == INVALID_SOCKET)
{
printf("[Redirector][SocksHelper::Utils::Connect] Create socket failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::Connect] Create socket failed: %d\n", WSAGetLastError());
return INVALID_SOCKET;
}
@@ -20,7 +20,7 @@ SOCKET SocksHelper::Connect()
int v6only = 0;
if (setsockopt(client, IPPROTO_IPV6, IPV6_V6ONLY, (char*)&v6only, sizeof(v6only)) == SOCKET_ERROR)
{
printf("[Redirector][SocksHelper::Utils::Connect] Set socket option failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::Connect] Set socket option failed: %d\n", WSAGetLastError());
closesocket(client);
return INVALID_SOCKET;
@@ -32,7 +32,7 @@ SOCKET SocksHelper::Connect()
if (!WSAConnectByNameW(client, (LPWSTR)tgtHost.c_str(), (LPWSTR)tgtPort.c_str(), NULL, NULL, NULL, NULL, &timeout, NULL))
{
printf("[Redirector][SocksHelper::Utils::Connect] Connect to remote server failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::Connect] Connect to remote server failed: %d\n", WSAGetLastError());
closesocket(client);
return INVALID_SOCKET;
@@ -53,14 +53,14 @@ bool SocksHelper::Handshake(SOCKET client)
buffer[3] = 0x02;
if (send(client, buffer, 4, 0) != 4)
{
printf("[Redirector][SocksHelper::Utils::Handshake] Send client hello failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::Handshake] Send client hello failed: %d\n", WSAGetLastError());
return false;
}
/* Server Choice */
if (recv(client, buffer, 2, 0) != 2)
{
printf("[Redirector][SocksHelper::Utils::Handshake] Receive server choice failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::Handshake] Receive server choice failed: %d\n", WSAGetLastError());
return false;
}
@@ -92,20 +92,20 @@ bool SocksHelper::Handshake(SOCKET client)
auto length = 1 + 1 + ulength + 1 + plength;
if (send(client, buffer, length, 0) != length)
{
printf("[Redirector][SocksHelper::Utils::Handshake] Send authentication request failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::Handshake] Send authentication request failed: %d\n", WSAGetLastError());
return false;
}
/* Server Response */
if (recv(client, buffer, 2, 0) != 2)
{
printf("[Redirector][SocksHelper::Utils::Handshake] Receive server response failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::Handshake] Receive server response failed: %d\n", WSAGetLastError());
return false;
}
if (buffer[1] != 0x00)
{
puts("[Redirector][SocksHelper::Utils::Handshake] Authentication failed");
puts("[Redirector][SocksHelper::Handshake] Authentication failed");
return false;
}
}
@@ -122,7 +122,7 @@ bool SocksHelper::SplitAddr(SOCKET client, PSOCKADDR_IN6 addr)
char addressType;
if (recv(client, (char*)&addressType, 1, 0) != 1)
{
printf("[Redirector][SocksHelper::Utils::SplitAddr] Read address type failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::SplitAddr] Read address type failed: %d\n", WSAGetLastError());
return false;
}
@@ -133,13 +133,13 @@ bool SocksHelper::SplitAddr(SOCKET client, PSOCKADDR_IN6 addr)
if (recv(client, (char*)&address->sin_addr, 4, 0) != 4)
{
printf("[Redirector][SocksHelper::Utils::SplitAddr] Read IPv4 address failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::SplitAddr] Read IPv4 address failed: %d\n", WSAGetLastError());
return false;
}
if (recv(client, (char*)&address->sin_port, 2, 0) != 2)
{
printf("[Redirector][SocksHelper::Utils::SplitAddr] Read IPv4 port failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::SplitAddr] Read IPv4 port failed: %d\n", WSAGetLastError());
return false;
}
}
@@ -149,19 +149,19 @@ bool SocksHelper::SplitAddr(SOCKET client, PSOCKADDR_IN6 addr)
if (recv(client, (char*)&addr->sin6_addr, 16, 0) != 16)
{
printf("[Redirector][SocksHelper::Utils::SplitAddr] Read IPv6 address failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::SplitAddr] Read IPv6 address failed: %d\n", WSAGetLastError());
return false;
}
if (recv(client, (char*)&addr->sin6_port, 2, 0) != 2)
{
printf("[Redirector][SocksHelper::Utils::SplitAddr] Read IPv6 port failed: %d\n", WSAGetLastError());
printf("[Redirector][SocksHelper::SplitAddr] Read IPv6 port failed: %d\n", WSAGetLastError());
return false;
}
}
else
{
printf("[Redirector][SocksHelper::Utils::SplitAddr] Unsupported address family: %d\n", addressType);
printf("[Redirector][SocksHelper::SplitAddr] Unsupported address family: %d\n", addressType);
return false;
}
@@ -393,6 +393,9 @@ int SocksHelper::UDP::Send(PSOCKADDR_IN6 target, const char* buffer, int length)
if (this->udpSocket == INVALID_SOCKET)
return SOCKET_ERROR;
if (target->sin6_family != AF_INET && target->sin6_family != AF_INET6)
return SOCKET_ERROR;
auto data = new char[3 + 1 + 16 + 2 + (ULONG64)length]();
data[3] = (target->sin6_family == AF_INET) ? 0x01 : 0x04;
@@ -403,18 +406,11 @@ int SocksHelper::UDP::Send(PSOCKADDR_IN6 target, const char* buffer, int length)
memcpy(data + 4, &ipv4->sin_addr, 4);
memcpy(data + 8, &ipv4->sin_port, 2);
}
else if (target->sin6_family == AF_INET6)
else
{
memcpy(data + 4, &target->sin6_addr, 16);
memcpy(data + 20, &target->sin6_port, 2);
}
else
{
delete[] data;
puts("[Redirector][SocksHelper::UDP::Send] Unsupported address family");
return length;
}
memcpy(data + 3 + 1 + (target->sin6_family == AF_INET ? 4 : 16) + 2, buffer, length);
auto dataLength = 3 + 1 + (target->sin6_family == AF_INET ? 4 : 16) + 2 + length;
@@ -442,14 +438,14 @@ int SocksHelper::UDP::Read(PSOCKADDR_IN6 target, char* buffer, int length, PTIME
FD_ZERO(&fds);
FD_SET(this->udpSocket, &fds);
int code = select(NULL, &fds, NULL, NULL, timeout);
if (code == 0 || code == SOCKET_ERROR)
return code;
int size = select(NULL, &fds, NULL, NULL, timeout);
if (size == 0 || size == SOCKET_ERROR)
return size;
}
int bufferLength = recvfrom(this->udpSocket, buffer, length, 0, NULL, NULL);
if (bufferLength == 0 || bufferLength == SOCKET_ERROR)
return bufferLength;
int size = recvfrom(this->udpSocket, buffer, length, 0, NULL, NULL);
if (size == 0 || size == SOCKET_ERROR)
return size;
SOCKADDR_IN6 addr;
if (buffer[3] == 0x01)
@@ -460,20 +456,24 @@ int SocksHelper::UDP::Read(PSOCKADDR_IN6 target, char* buffer, int length, PTIME
memcpy(&ipv4->sin_addr, buffer + 4, 4);
memcpy(&ipv4->sin_port, buffer + 8, 2);
memcpy(buffer, buffer + 10, (ULONG64)bufferLength - 10);
memcpy(buffer, buffer + 10, (ULONG64)size - 10);
}
else
else if (buffer[3] == 0x04)
{
addr.sin6_family = AF_INET6;
memcpy(&addr.sin6_addr, buffer + 4, 16);
memcpy(&addr.sin6_port, buffer + 20, 2);
memcpy(buffer, buffer + 22, (ULONG64)bufferLength - 22);
memcpy(buffer, buffer + 22, (ULONG64)size - 22);
}
else
{
return SOCKET_ERROR;
}
if (target != NULL)
memcpy(target, &addr, sizeof(SOCKADDR_IN6));
return bufferLength - (addr.sin6_family == AF_INET ? 10 : 22);
return size - (addr.sin6_family == AF_INET ? 10 : 22);
}

View File

@@ -13,6 +13,7 @@ bool TCPHandler::INIT()
if (tcpSocket != INVALID_SOCKET)
{
closesocket(tcpSocket);
tcpSocket = INVALID_SOCKET;
}
@@ -94,11 +95,9 @@ void TCPHandler::CreateHandler(SOCKADDR_IN6 client, SOCKADDR_IN6 remote)
{
auto lg = lock_guard<mutex>(tcpLock);
auto id = (client.sin6_family == AF_INET6) ? client.sin6_port : ((PSOCKADDR_IN)&client)->sin_port;
auto id = (client.sin6_family == AF_INET) ? ((PSOCKADDR_IN)&client)->sin_port : client.sin6_port;
if (tcpContext.find(id) != tcpContext.end())
{
tcpContext.erase(id);
}
tcpContext[id] = remote;
}
@@ -107,11 +106,9 @@ void TCPHandler::DeleteHandler(SOCKADDR_IN6 client)
{
auto lg = lock_guard<mutex>(tcpLock);
auto id = (client.sin6_family == AF_INET6) ? client.sin6_port : ((PSOCKADDR_IN)&client)->sin_port;
auto id = (client.sin6_family == AF_INET) ? ((PSOCKADDR_IN)&client)->sin_port : client.sin6_port;
if (tcpContext.find(id) != tcpContext.end())
{
tcpContext.erase(id);
}
}
void TCPHandler::Accept()
@@ -132,16 +129,18 @@ void TCPHandler::Accept()
void TCPHandler::Handle(SOCKET client)
{
USHORT id = 0;
{
SOCKADDR_IN6 addr;
int addrLength = sizeof(SOCKADDR_IN6);
if (getpeername(client, (PSOCKADDR)&addr, &addrLength) == SOCKET_ERROR)
{
closesocket(client);
return;
}
id = (addr.sin6_family == AF_INET6) ? addr.sin6_port : ((PSOCKADDR_IN)&addr)->sin_port;
id = (addr.sin6_family == AF_INET) ? ((PSOCKADDR_IN)&addr)->sin_port : addr.sin6_port;
}
tcpLock.lock();
@@ -180,14 +179,10 @@ void TCPHandler::Read(SOCKET client, SocksHelper::PTCP remote)
{
int length = remote->Read(buffer, sizeof(buffer));
if (length == 0 || length == SOCKET_ERROR)
{
return;
}
if (send(client, buffer, length, 0) != length)
{
return;
}
}
}
@@ -199,13 +194,9 @@ void TCPHandler::Send(SOCKET client, SocksHelper::PTCP remote)
{
int length = recv(client, buffer, sizeof(buffer), 0);
if (length == 0 || length == SOCKET_ERROR)
{
return;
}
if (remote->Send(buffer, length) != length)
{
return;
}
}
}