handle incomplete MSG_WAITALL during tcp receive

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);

src/Network/Core.cpp

@@ -143,29 +143,28 @@ void GetServerInfo(std::string Data) {
     const std::string buffer = ([&]() -> std::string {
         int32_t Header;
         std::vector<char> 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);

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);

src/Network/Resources.cpp

@@ -164,7 +164,7 @@ std::vector<char> 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) {

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<char> 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 "";
     }