diff --git a/Redirector/EventHandler.cpp b/Redirector/EventHandler.cpp index a017a934..b9e58260 100644 --- a/Redirector/EventHandler.cpp +++ b/Redirector/EventHandler.cpp @@ -9,6 +9,9 @@ extern vector handleList; extern USHORT tcpListen; +mutex udpContextLock; +map 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 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 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; } diff --git a/Redirector/SocksHelper.cpp b/Redirector/SocksHelper.cpp index 1ea3f2ab..14cd140b 100644 --- a/Redirector/SocksHelper.cpp +++ b/Redirector/SocksHelper.cpp @@ -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) { diff --git a/Redirector/SocksHelper.h b/Redirector/SocksHelper.h index 75b8f60c..a5b39c36 100644 --- a/Redirector/SocksHelper.h +++ b/Redirector/SocksHelper.h @@ -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(); diff --git a/Redirector/TCPHandler.cpp b/Redirector/TCPHandler.cpp index 7a6a8071..20692d68 100644 --- a/Redirector/TCPHandler.cpp +++ b/Redirector/TCPHandler.cpp @@ -10,10 +10,10 @@ bool TCPHandler::Init() { auto lg = lock_guard(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; }