[Redirector] Update UDPHandler

This commit is contained in:
Connection Refused
2021-09-29 15:52:12 +08:00
parent 5726c6ee2e
commit 3c84491e24
4 changed files with 124 additions and 27 deletions

View File

@@ -9,6 +9,9 @@ extern vector<wstring> handleList;
extern USHORT tcpListen;
mutex udpContextLock;
map<ENDPOINT_ID, SocksHelper::PUDP> udpContext;
wstring ConvertIP(PSOCKADDR addr)
{
WCHAR buffer[MAX_PATH] = L"";
@@ -227,7 +230,8 @@ void udpCreated(ENDPOINT_ID id, PNF_UDP_CONN_INFO info)
return;
}
nf_udpDisableFiltering(id);
lock_guard<mutex> lg(udpContextLock);
udpContext[id] = new SocksHelper::UDP();
}
void udpConnectRequest(ENDPOINT_ID id, PNF_UDP_CONN_REQUEST info)
@@ -243,7 +247,75 @@ void udpCanSend(ENDPOINT_ID id)
void udpSend(ENDPOINT_ID id, const unsigned char* target, const char* buffer, int length, PNF_UDP_OPTIONS options)
{
nf_udpPostSend(id, target, buffer, length, options);
udpContextLock.lock();
if (udpContext.find(id) == udpContext.end())
{
udpContextLock.unlock();
nf_udpPostSend(id, target, buffer, length, options);
return;
}
auto conn = udpContext[id];
udpContextLock.unlock();
if (conn->tcpSocket == INVALID_SOCKET)
{
auto tcpSocket = SocksHelper::Utils::Connect();
if (tcpSocket == INVALID_SOCKET)
{
printf("[Redirector][EventHandler][udpSend][%llu] Connect to remote server failed\n", id);
return;
}
if (!SocksHelper::Utils::Handshake(tcpSocket))
{
closesocket(tcpSocket);
printf("[Redirector][EventHandler][udpSend][%llu] Handshake failed\n", id);
return;
}
conn->tcpSocket = tcpSocket;
}
if (conn->udpSocket == INVALID_SOCKET)
{
if (!conn->Associate())
{
closesocket(conn->tcpSocket);
conn->tcpSocket = INVALID_SOCKET;
printf("[Redirector][EventHandler][udpSend][%llu] UDP Associate failed\n", id);
return;
}
if (!conn->CreateUDP())
{
closesocket(conn->tcpSocket);
conn->tcpSocket = INVALID_SOCKET;
printf("[Redirector][EventHandler][udpSend][%llu] Create UDP socket failed\n", id);
return;
}
PNF_UDP_OPTIONS data = new NF_UDP_OPTIONS();
memcpy(data, options, sizeof(NF_UDP_OPTIONS));
thread(udpBeginReceive, id, conn, data).detach();
}
if (conn->Send((PSOCKADDR)target, buffer, length) == SOCKET_ERROR)
{
closesocket(conn->tcpSocket);
closesocket(conn->udpSocket);
conn->tcpSocket = INVALID_SOCKET;
conn->udpSocket = INVALID_SOCKET;
printf("[Redirector][EventHandler][udpSend][%llu] Send data failed\n", id);
return;
}
}
void udpCanReceive(ENDPOINT_ID id)
@@ -258,6 +330,36 @@ void udpReceive(ENDPOINT_ID id, const unsigned char* target, const char* buffer,
void udpClosed(ENDPOINT_ID id, PNF_UDP_CONN_INFO info)
{
UNREFERENCED_PARAMETER(id);
UNREFERENCED_PARAMETER(info);
printf("[Redirector][EventHandler][udpClosed][%llu]\n", id);
lock_guard<mutex> lg(udpContextLock);
if (udpContext.find(id) != udpContext.end())
{
delete udpContext[id];
udpContext.erase(id);
}
}
void udpBeginReceive(ENDPOINT_ID id, SocksHelper::PUDP conn, PNF_UDP_OPTIONS data)
{
char buffer[1458];
while (conn->udpSocket != INVALID_SOCKET)
{
SOCKADDR_IN6 target;
int targetLength = sizeof(SOCKADDR_IN6);
int length = conn->Read((PSOCKADDR)&target, buffer, 1458);
if (length == 0 || length == SOCKET_ERROR)
{
break;
}
nf_udpPostReceive(id, (unsigned char*)&target, buffer, length, data);
}
delete data;
}

View File

@@ -238,7 +238,7 @@ bool SocksHelper::TCP::Connect(PSOCKADDR target)
return Utils::ReadAddr(this->tcpSocket, buffer[3], (PSOCKADDR)&addr);
}
int SocksHelper::TCP::Send(char* buffer, int length)
int SocksHelper::TCP::Send(const char* buffer, int length)
{
if (this->tcpSocket != INVALID_SOCKET)
{
@@ -371,7 +371,7 @@ bool SocksHelper::UDP::CreateUDP()
return true;
}
int SocksHelper::UDP::Send(PSOCKADDR target, char* buffer, int length)
int SocksHelper::UDP::Send(PSOCKADDR target, const char* buffer, int length)
{
if (this->udpSocket == INVALID_SOCKET)
{
@@ -425,17 +425,11 @@ int SocksHelper::UDP::Read(PSOCKADDR target, char* buffer, int length)
return SOCKET_ERROR;
}
auto targetLength = 0;
auto bufferLength = recvfrom(this->udpSocket, buffer, length, 0, target, &targetLength);
if (bufferLength <= 0)
int targetLength = 0;
int bufferLength = recvfrom(this->udpSocket, buffer, length, 0, target, &targetLength);
if (bufferLength == 0 || bufferLength == SOCKET_ERROR)
{
if (bufferLength == SOCKET_ERROR)
{
printf("[Redirector][SocksHelper::UDP::Read] Receive packet failed: %d\n", WSAGetLastError());
return SOCKET_ERROR;
}
return 0;
return bufferLength;
}
memset(target, 0, sizeof(SOCKADDR_IN6));
@@ -467,7 +461,7 @@ void SocksHelper::UDP::run()
{
char buffer[1];
while (true)
while (this->tcpSocket != INVALID_SOCKET)
{
if (recv(this->tcpSocket, buffer, 1, 0) != 1)
{

View File

@@ -21,10 +21,10 @@ namespace SocksHelper
bool Connect(PSOCKADDR target);
int Send(char* buffer, int length);
int Send(const char* buffer, int length);
int Read(char* buffer, int length);
SOCKET tcpSocket = NULL;
SOCKET tcpSocket = INVALID_SOCKET;
} *PTCP;
typedef class UDP
@@ -37,11 +37,11 @@ namespace SocksHelper
bool Associate();
bool CreateUDP();
int Send(PSOCKADDR target, char* buffer, int length);
int Send(PSOCKADDR target, const char* buffer, int length);
int Read(PSOCKADDR target, char* buffer, int length);
SOCKET tcpSocket = NULL;
SOCKET udpSocket = NULL;
SOCKET tcpSocket = INVALID_SOCKET;
SOCKET udpSocket = INVALID_SOCKET;
private:
void run();

View File

@@ -10,10 +10,10 @@ bool TCPHandler::Init()
{
auto lg = lock_guard<mutex>(tcpLock);
if (tcpSocket)
if (tcpSocket != INVALID_SOCKET)
{
closesocket(tcpSocket);
tcpSocket = NULL;
tcpSocket = INVALID_SOCKET;
}
auto client = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
@@ -120,6 +120,7 @@ void TCPHandler::Accept()
auto client = accept(tcpSocket, NULL, NULL);
if (client == INVALID_SOCKET)
{
printf("[Redirector][TCPHandler::Accept] Accept client failed: %d\n", WSAGetLastError());
return;
}
@@ -191,8 +192,8 @@ void TCPHandler::Read(SOCKET client, SocksHelper::PTCP remote)
while (tcpSocket != INVALID_SOCKET)
{
auto length = remote->Read(buffer, 1446);
if (!length || length == SOCKET_ERROR)
int length = remote->Read(buffer, 1446);
if (length == 0 || length == SOCKET_ERROR)
{
return;
}
@@ -210,8 +211,8 @@ void TCPHandler::Send(SOCKET client, SocksHelper::PTCP remote)
while (tcpSocket != INVALID_SOCKET)
{
auto length = recv(client, buffer, 1446, 0);
if (!length || length == SOCKET_ERROR)
int length = recv(client, buffer, 1446, 0);
if (length == 0 || length == SOCKET_ERROR)
{
return;
}