diff --git a/include/Network/network.hpp b/include/Network/network.hpp index f58bdad4..092853e8 100644 --- a/include/Network/network.hpp +++ b/include/Network/network.hpp @@ -42,7 +42,7 @@ extern std::string magic; int KillSocket(uint64_t Dead); void UUl(const std::string& R); void UDPSend(std::string Data); -bool CheckBytes(int32_t Bytes); +bool CheckBytes(int32_t Bytes, int32_t Expected = -1); void GameSend(std::string_view Data); void SendLarge(std::string Data); std::string TCPRcv(uint64_t Sock); @@ -56,3 +56,4 @@ void UDPClientMain(const std::string& IP, int Port); void TCPGameServer(const std::string& IP, int Port); bool SecurityWarning(); void CoreSend(std::string data); +int RecvWaitAll(int sockfd, char *buf, int len); diff --git a/src/Network/Core.cpp b/src/Network/Core.cpp index bcb80509..8dddfdd8 100644 --- a/src/Network/Core.cpp +++ b/src/Network/Core.cpp @@ -143,29 +143,28 @@ void GetServerInfo(std::string Data) { const std::string buffer = ([&]() -> std::string { int32_t Header; std::vector data(sizeof(Header)); - int Temp = recv(ISock, data.data(), sizeof(Header), MSG_WAITALL); + int Temp = RecvWaitAll(ISock, data.data(), sizeof(Header)); - auto checkBytes = ([&](const int32_t bytes) -> bool { + auto checkBytes = ([&](const int32_t bytes, const int32_t expected = -1) -> bool { if (bytes == 0) { return false; } else if (bytes < 0) { return false; } + if (expected != -1 && bytes != expected) { + return false; + } return true; }); - if (!checkBytes(Temp)) { + if (!checkBytes(Temp, sizeof(Header))) { return ""; } memcpy(&Header, data.data(), sizeof(Header)); - if (!checkBytes(Temp)) { - return ""; - } - data.resize(Header, 0); - Temp = recv(ISock, data.data(), Header, MSG_WAITALL); - if (!checkBytes(Temp)) { + Temp = RecvWaitAll(ISock, data.data(), Header); + if (!checkBytes(Temp, Header)) { return ""; } return std::string(data.data(), Header); diff --git a/src/Network/GlobalHandler.cpp b/src/Network/GlobalHandler.cpp index da4b6f28..c2290df6 100644 --- a/src/Network/GlobalHandler.cpp +++ b/src/Network/GlobalHandler.cpp @@ -50,17 +50,6 @@ int KillSocket(uint64_t Dead) { return a; } -bool CheckBytes(uint32_t Bytes) { - if (Bytes == 0) { - debug("(Proxy) Connection closing"); - return false; - } else if (Bytes < 0) { - debug("(Proxy) send failed with error: " + std::to_string(WSAGetLastError())); - return false; - } - return true; -} - void GameSend(std::string_view Data) { static std::mutex Lock; std::scoped_lock Guard(Lock); diff --git a/src/Network/Resources.cpp b/src/Network/Resources.cpp index 1696dde8..26bc62e3 100644 --- a/src/Network/Resources.cpp +++ b/src/Network/Resources.cpp @@ -164,7 +164,7 @@ std::vector TCPRcvRaw(SOCKET Sock, uint64_t& GRcv, uint64_t Size) { do { // receive at most some MB at a time int Len = std::min(int(Size - Rcv), 1 * 1024 * 1024); - int Temp = recv(Sock, &File[Rcv], Len, MSG_WAITALL); + int Temp = RecvWaitAll(Sock, &File[Rcv], Len); if (Temp == -1 || Temp == 0) { debug("Recv returned: " + std::to_string(Temp)); if (Temp == -1) { diff --git a/src/Network/VehicleEvent.cpp b/src/Network/VehicleEvent.cpp index dcbb41e5..cecacea3 100644 --- a/src/Network/VehicleEvent.cpp +++ b/src/Network/VehicleEvent.cpp @@ -28,9 +28,9 @@ int LastPort; std::string LastIP; SOCKET TCPSock = -1; -bool CheckBytes(int32_t Bytes) { +bool CheckBytes(int32_t Bytes, int32_t Expected) { if (Bytes == 0) { - debug("(TCP) Connection closing... CheckBytes(16)"); + debug("(TCP) Connection closing..."); Terminate = true; return false; } else if (Bytes < 0) { @@ -39,6 +39,11 @@ bool CheckBytes(int32_t Bytes) { Terminate = true; return false; } + if (Expected != -1 && Bytes != Expected) { + debug(std::format("(TCP) Short recv detected, expected {} bytes, got {} bytes", Expected, Bytes)); + Terminate = true; + return false; + } return true; } void UUl(const std::string& R) { @@ -75,6 +80,24 @@ void TCPSend(const std::string& Data, uint64_t Sock) { } while (Sent < Size); } +int RecvWaitAll(int sockfd, char *buf, int len) { + // handle MSG_WAITALL not actually filling the whole buffer + // happens frequently in wine, and can also happen natively when the OS pauses the execution for various reasons + int offset = 0; + while (offset < len) { + int recv_status = recv(sockfd, &buf[offset], len - offset, MSG_WAITALL); + if (recv_status == 0) { + // do not discard received data when the other side closes the socket cleanly + return offset; + } + if (recv_status == -1) { + return -1; + } + offset += recv_status; + } + return offset; +} + std::string TCPRcv(SOCKET Sock) { if (Sock == -1) { Terminate = true; @@ -84,8 +107,8 @@ std::string TCPRcv(SOCKET Sock) { int32_t Header; int Temp; std::vector Data(sizeof(Header)); - Temp = recv(Sock, Data.data(), sizeof(Header), MSG_WAITALL); - if (!CheckBytes(Temp)) { + Temp = RecvWaitAll(Sock, Data.data(), sizeof(Header)); + if (!CheckBytes(Temp, sizeof(Header))) { UUl("Socket Closed Code 3"); return ""; } @@ -97,8 +120,8 @@ std::string TCPRcv(SOCKET Sock) { } Data.resize(Header, 0); - Temp = recv(Sock, Data.data(), Header, MSG_WAITALL); - if (!CheckBytes(Temp)) { + Temp = RecvWaitAll(Sock, Data.data(), Header); + if (!CheckBytes(Temp, Header)) { UUl("Socket Closed Code 5"); return ""; }