diff --git a/Redirector/DNSHandler.cpp b/Redirector/DNSHandler.cpp index 14d0bf7b..0ce0a550 100644 --- a/Redirector/DNSHandler.cpp +++ b/Redirector/DNSHandler.cpp @@ -23,7 +23,10 @@ void ProcessPacket(ENDPOINT_ID id, SOCKADDR_IN6 target, char* packet, int length { if (udpConn.Send(&dnsAddr, packet, length) == length) { - int size = udpConn.Read(NULL, buffer, sizeof(buffer)); + 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); diff --git a/Redirector/SocksHelper.cpp b/Redirector/SocksHelper.cpp index 8dd7b5aa..2f517904 100644 --- a/Redirector/SocksHelper.cpp +++ b/Redirector/SocksHelper.cpp @@ -398,11 +398,22 @@ int SocksHelper::UDP::Send(PSOCKADDR_IN6 target, const char* buffer, int length) return length; } -int SocksHelper::UDP::Read(PSOCKADDR_IN6 target, char* buffer, int length) +int SocksHelper::UDP::Read(PSOCKADDR_IN6 target, char* buffer, int length, PTIMEVAL timeout = NULL) { if (!this->udpSocket) return SOCKET_ERROR; + if (timeout != NULL) + { + fd_set fds; + FD_ZERO(&fds); + FD_SET(this->udpSocket, &fds); + + int code = select(this->udpSocket, &fds, NULL, NULL, timeout); + if (code == 0 || code == SOCKET_ERROR) + return code; + } + int bufferLength = recvfrom(this->udpSocket, buffer, length, 0, NULL, NULL); if (bufferLength == 0 || bufferLength == SOCKET_ERROR) return bufferLength; diff --git a/Redirector/SocksHelper.h b/Redirector/SocksHelper.h index a69d19d5..c6068720 100644 --- a/Redirector/SocksHelper.h +++ b/Redirector/SocksHelper.h @@ -34,7 +34,7 @@ namespace SocksHelper bool CreateUDP(); int Send(PSOCKADDR_IN6 target, const char* buffer, int length); - int Read(PSOCKADDR_IN6 target, char* buffer, int length); + int Read(PSOCKADDR_IN6 target, char* buffer, int length, PTIMEVAL timeout = NULL); SOCKET tcpSocket = INVALID_SOCKET; SOCKET udpSocket = INVALID_SOCKET;