From 1eedab470ee60c0510411f1ca1de62908ae18f5e Mon Sep 17 00:00:00 2001
From: Liu Woon Yung <pirorin187@gmail.com>
Date: Wed, 3 Jul 2019 19:24:29 +0800
Subject: [PATCH] Fixed various TCP-related deadlocks and improved
 disconnection mechanism. 1. TCPServer::listener_loop() not disconnecting
 failed/disconnected sockets, leaving the socket unusable. 2. Fixed
 TCPServer::recv_locking() and TCPServer::send_locking() not indicating the
 non-active connection state as an error, which previously resulted in the
 caller entering an infinite loop. 3. Synchronize the TCP stream with the
 client before closing the socket, to avoid possibly losing bytes in flight.

---
 src/cpp/transport/tcp/TCPServerLinux.cpp   | 28 ++++++++++++++++++++++
 src/cpp/transport/tcp/TCPServerWindows.cpp | 28 ++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/src/cpp/transport/tcp/TCPServerLinux.cpp b/src/cpp/transport/tcp/TCPServerLinux.cpp
index 277cf388f..0ac621bbf 100644
--- a/src/cpp/transport/tcp/TCPServerLinux.cpp
+++ b/src/cpp/transport/tcp/TCPServerLinux.cpp
@@ -391,6 +391,21 @@ bool TCPv4Agent::close_connection(TCPConnection& connection)
         lock.unlock();
         /* Add lock for close. */
         std::unique_lock<std::mutex> conn_lock(connection.mtx);
+
+        /* Synchronize the stream with the client, to avoid losing bytes currently in flight. */
+        shutdown(connection_platform.poll_fd->fd, SHUT_WR);
+        fd_set read_fdset;
+        struct timeval timeout = {
+            120,
+            0
+        };
+        FD_ZERO(&read_fdset);
+        FD_SET(connection_platform.poll_fd->fd, &read_fdset);
+        if (select(connection_platform.poll_fd->fd + 1, &read_fdset, NULL, NULL, &timeout) > 0)
+        {
+            char dummy;
+            while (recv(connection_platform.poll_fd->fd, &dummy, sizeof(dummy), 0) > 0) {};
+        }
         if (0 == ::close(connection_platform.poll_fd->fd))
         {
             connection_platform.poll_fd->fd = -1;
@@ -479,6 +494,11 @@ void TCPv4Agent::listener_loop()
                     }
                 }
             }
+            else if (0 < ((POLLERR | POLLHUP) & conn.poll_fd->revents))
+			{
+                /* Disconnect failed connections. */
+                close_connection(conn);
+            }
         }
     }
 }
@@ -519,6 +539,10 @@ size_t TCPv4Agent::recv_locking(
             errcode = (0 == poll_rv) ? 0 : 1;
         }
     }
+    else
+    {
+        errcode = 1;
+    }
     return rv;
 }
 
@@ -544,6 +568,10 @@ size_t TCPv4Agent::send_locking(
             errcode = 1;
         }
     }
+    else
+    {
+        errcode = 1;
+    }
     return rv;
 }
 
diff --git a/src/cpp/transport/tcp/TCPServerWindows.cpp b/src/cpp/transport/tcp/TCPServerWindows.cpp
index 42b928d9c..cb96dfdda 100644
--- a/src/cpp/transport/tcp/TCPServerWindows.cpp
+++ b/src/cpp/transport/tcp/TCPServerWindows.cpp
@@ -363,6 +363,21 @@ bool TCPv4Agent::close_connection(TCPConnection& connection)
         lock.unlock();
         /* Add lock for close. */
         std::unique_lock<std::mutex> conn_lock(connection.mtx);
+
+        /* Synchronize the stream with the client, to avoid losing bytes currently in flight. */
+        shutdown(connection_platform.poll_fd->fd, SD_SEND);
+        fd_set read_fdset;
+        struct timeval timeout = {
+            120,
+            0
+        };
+        FD_ZERO(&read_fdset);
+        FD_SET(connection_platform.poll_fd->fd, &read_fdset);
+        if (select(connection_platform.poll_fd->fd + 1, &read_fdset, NULL, NULL, &timeout) > 0)
+        {
+            char dummy;
+            while (recv(connection_platform.poll_fd->fd, &dummy, sizeof(dummy), 0) > 0) {};
+        }
         if (0 == closesocket(connection_platform.poll_fd->fd))
         {
             connection_platform.poll_fd->fd = INVALID_SOCKET;
@@ -416,6 +431,11 @@ bool TCPv4Agent::read_message(int timeout)
                     rv = true;
                 }
             }
+            else if (0 < ((POLLERR | POLLHUP) & conn.poll_fd->revents))
+            {
+                /* Disconnect failed connections. */
+                close_connection(conn);
+            }
         }
     }
     else
@@ -490,6 +510,10 @@ size_t TCPv4Agent::recv_locking(
             errcode = (0 == poll_rv) ? 0 : 1;
         }
     }
+    else
+    {
+        errcode = 1;
+    }
     return rv;
 }
 
@@ -515,6 +539,10 @@ size_t TCPv4Agent::send_locking(
             errcode = 1;
         }
     }
+    else
+    {
+        errcode = 1;
+    }
     return rv;
 }