Ensure TCP TIME-WAIT is not terminated prematurely.

Netstack's TIME-WAIT state for a TCP socket could be terminated prematurely if
the socket entered TIME-WAIT using shutdown(..., SHUT_RDWR) and then was closed
using close(). This fixes that bug and updates the tests to verify that Netstack
correctly honors TIME-WAIT under such conditions.

Fixes #3106

PiperOrigin-RevId: 326456443

3 files changed, 180 insertions, 5 deletions

diff --git a/test/syscalls/linux/socket_generic_stress.cc b/test/syscalls/linux/socket_generic_stress.cc
index 6a232238d..19239e9e9 100644
--- a/test/syscalls/linux/socket_generic_stress.cc
+++ b/test/syscalls/linux/socket_generic_stress.cc
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <poll.h>
 #include <stdio.h>
 #include <sys/ioctl.h>
 #include <sys/socket.h>
@@ -37,6 +38,14 @@ TEST_P(ConnectStressTest, Reset65kTimes) {
     char sent_data[100] = {};
     ASSERT_THAT(write(sockets->first_fd(), sent_data, sizeof(sent_data)),
                 SyscallSucceedsWithValue(sizeof(sent_data)));
+    // Poll the other FD to make sure that the data is in the receive buffer
+    // before closing it to ensure a RST is triggered.
+    const int kTimeout = 10000;
+    struct pollfd pfd = {
+        .fd = sockets->second_fd(),
+        .events = POLL_IN,
+    };
+    ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
   }
 }
 
@@ -58,7 +67,45 @@ INSTANTIATE_TEST_SUITE_P(
 // a persistent listener (if applicable).
 using PersistentListenerConnectStressTest = SocketPairTest;
 
-TEST_P(PersistentListenerConnectStressTest, 65kTimes) {
+TEST_P(PersistentListenerConnectStressTest, 65kTimesShutdownCloseFirst) {
+  for (int i = 0; i < 1 << 16; ++i) {
+    auto sockets = ASSERT_NO_ERRNO_AND_VALUE(NewSocketPair());
+    ASSERT_THAT(shutdown(sockets->first_fd(), SHUT_RDWR), SyscallSucceeds());
+    if (GetParam().type == SOCK_STREAM) {
+      // Poll the other FD to make sure that we see the FIN from the other
+      // side before closing the second_fd. This ensures that the first_fd
+      // enters TIME-WAIT and not second_fd.
+      const int kTimeout = 10000;
+      struct pollfd pfd = {
+          .fd = sockets->second_fd(),
+          .events = POLL_IN,
+      };
+      ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
+    }
+    ASSERT_THAT(shutdown(sockets->second_fd(), SHUT_RDWR), SyscallSucceeds());
+  }
+}
+
+TEST_P(PersistentListenerConnectStressTest, 65kTimesShutdownCloseSecond) {
+  for (int i = 0; i < 1 << 16; ++i) {
+    auto sockets = ASSERT_NO_ERRNO_AND_VALUE(NewSocketPair());
+    ASSERT_THAT(shutdown(sockets->second_fd(), SHUT_RDWR), SyscallSucceeds());
+    if (GetParam().type == SOCK_STREAM) {
+      // Poll the other FD to make sure that we see the FIN from the other
+      // side before closing the first_fd. This ensures that the second_fd
+      // enters TIME-WAIT and not first_fd.
+      const int kTimeout = 10000;
+      struct pollfd pfd = {
+          .fd = sockets->first_fd(),
+          .events = POLL_IN,
+      };
+      ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
+    }
+    ASSERT_THAT(shutdown(sockets->first_fd(), SHUT_RDWR), SyscallSucceeds());
+  }
+}
+
+TEST_P(PersistentListenerConnectStressTest, 65kTimesClose) {
   for (int i = 0; i < 1 << 16; ++i) {
     auto sockets = ASSERT_NO_ERRNO_AND_VALUE(NewSocketPair());
   }
diff --git a/test/syscalls/linux/socket_inet_loopback.cc b/test/syscalls/linux/socket_inet_loopback.cc
index 18b9e4b70..c3b42682f 100644
--- a/test/syscalls/linux/socket_inet_loopback.cc
+++ b/test/syscalls/linux/socket_inet_loopback.cc
@@ -865,7 +865,7 @@ TEST_P(SocketInetLoopbackTest, TCPResetAfterClose) {
 // results in the stack.Seed() being different which can cause
 // sequence number of final connect to be one that is considered
 // old and can cause the test to be flaky.
-TEST_P(SocketInetLoopbackTest, TCPTimeWaitTest_NoRandomSave) {
+TEST_P(SocketInetLoopbackTest, TCPPassiveCloseNoTimeWaitTest_NoRandomSave) {
   auto const& param = GetParam();
   TestAddress const& listener = param.listener;
   TestAddress const& connector = param.connector;
@@ -920,14 +920,27 @@ TEST_P(SocketInetLoopbackTest, TCPTimeWaitTest_NoRandomSave) {
                   &conn_addrlen),
       SyscallSucceeds());
 
-  // close the accept FD to trigger TIME_WAIT on the accepted socket which
+  // shutdown the accept FD to trigger TIME_WAIT on the accepted socket which
   // should cause the conn_fd to follow CLOSE_WAIT->LAST_ACK->CLOSED instead of
   // TIME_WAIT.
-  accepted.reset();
-  absl::SleepFor(absl::Seconds(1));
+  ASSERT_THAT(shutdown(accepted.get(), SHUT_RDWR), SyscallSucceeds());
+  {
+    const int kTimeout = 10000;
+    struct pollfd pfd = {
+        .fd = conn_fd.get(),
+        .events = POLLIN,
+    };
+    ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
+    ASSERT_EQ(pfd.revents, POLLIN);
+  }
+
   conn_fd.reset();
+  // This sleep is required to give conn_fd time to transition to TIME-WAIT.
   absl::SleepFor(absl::Seconds(1));
 
+  // At this point conn_fd should be the one that moved to CLOSE_WAIT and
+  // eventually to CLOSED.
+
   // Now bind and connect a new socket and verify that we can immediately
   // rebind the address bound by the conn_fd as it never entered TIME_WAIT.
   const FileDescriptor conn_fd2 = ASSERT_NO_ERRNO_AND_VALUE(
@@ -942,6 +955,97 @@ TEST_P(SocketInetLoopbackTest, TCPTimeWaitTest_NoRandomSave) {
               SyscallSucceeds());
 }
 
+TEST_P(SocketInetLoopbackTest, TCPActiveCloseTimeWaitTest_NoRandomSave) {
+  auto const& param = GetParam();
+  TestAddress const& listener = param.listener;
+  TestAddress const& connector = param.connector;
+
+  // Create the listening socket.
+  const FileDescriptor listen_fd = ASSERT_NO_ERRNO_AND_VALUE(
+      Socket(listener.family(), SOCK_STREAM, IPPROTO_TCP));
+  sockaddr_storage listen_addr = listener.addr;
+  ASSERT_THAT(bind(listen_fd.get(), reinterpret_cast<sockaddr*>(&listen_addr),
+                   listener.addr_len),
+              SyscallSucceeds());
+  ASSERT_THAT(listen(listen_fd.get(), SOMAXCONN), SyscallSucceeds());
+
+  // Get the port bound by the listening socket.
+  socklen_t addrlen = listener.addr_len;
+  ASSERT_THAT(getsockname(listen_fd.get(),
+                          reinterpret_cast<sockaddr*>(&listen_addr), &addrlen),
+              SyscallSucceeds());
+
+  uint16_t const port =
+      ASSERT_NO_ERRNO_AND_VALUE(AddrPort(listener.family(), listen_addr));
+
+  // Connect to the listening socket.
+  FileDescriptor conn_fd = ASSERT_NO_ERRNO_AND_VALUE(
+      Socket(connector.family(), SOCK_STREAM, IPPROTO_TCP));
+
+  // We disable saves after this point as a S/R causes the netstack seed
+  // to be regenerated which changes what ports/ISN is picked for a given
+  // tuple (src ip,src port, dst ip, dst port). This can cause the final
+  // SYN to use a sequence number that looks like one from the current
+  // connection in TIME_WAIT and will not be accepted causing the test
+  // to timeout.
+  //
+  // TODO(gvisor.dev/issue/940): S/R portSeed/portHint
+  DisableSave ds;
+
+  sockaddr_storage conn_addr = connector.addr;
+  ASSERT_NO_ERRNO(SetAddrPort(connector.family(), &conn_addr, port));
+  ASSERT_THAT(RetryEINTR(connect)(conn_fd.get(),
+                                  reinterpret_cast<sockaddr*>(&conn_addr),
+                                  connector.addr_len),
+              SyscallSucceeds());
+
+  // Accept the connection.
+  auto accepted =
+      ASSERT_NO_ERRNO_AND_VALUE(Accept(listen_fd.get(), nullptr, nullptr));
+
+  // Get the address/port bound by the connecting socket.
+  sockaddr_storage conn_bound_addr;
+  socklen_t conn_addrlen = connector.addr_len;
+  ASSERT_THAT(
+      getsockname(conn_fd.get(), reinterpret_cast<sockaddr*>(&conn_bound_addr),
+                  &conn_addrlen),
+      SyscallSucceeds());
+
+  // shutdown the conn FD to trigger TIME_WAIT on the connect socket.
+  ASSERT_THAT(shutdown(conn_fd.get(), SHUT_RDWR), SyscallSucceeds());
+  {
+    const int kTimeout = 10000;
+    struct pollfd pfd = {
+        .fd = accepted.get(),
+        .events = POLLIN,
+    };
+    ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
+    ASSERT_EQ(pfd.revents, POLLIN);
+  }
+  ScopedThread t([&]() {
+    constexpr int kTimeout = 10000;
+    constexpr int16_t want_events = POLLHUP;
+    struct pollfd pfd = {
+        .fd = conn_fd.get(),
+        .events = want_events,
+    };
+    ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
+  });
+
+  accepted.reset();
+  t.Join();
+  conn_fd.reset();
+
+  // Now bind and connect a new socket and verify that we can't immediately
+  // rebind the address bound by the conn_fd as it is in TIME_WAIT.
+  conn_fd = ASSERT_NO_ERRNO_AND_VALUE(
+      Socket(connector.family(), SOCK_STREAM, IPPROTO_TCP));
+
+  ASSERT_THAT(bind(conn_fd.get(), reinterpret_cast<sockaddr*>(&conn_bound_addr),
+                   conn_addrlen),
+              SyscallFailsWithErrno(EADDRINUSE));
+}
+
 TEST_P(SocketInetLoopbackTest, AcceptedInheritsTCPUserTimeout) {
   auto const& param = GetParam();
   TestAddress const& listener = param.listener;
diff --git a/test/syscalls/linux/tcp_socket.cc b/test/syscalls/linux/tcp_socket.cc
index 0cea7d11f..a6325a761 100644
--- a/test/syscalls/linux/tcp_socket.cc
+++ b/test/syscalls/linux/tcp_socket.cc
@@ -720,6 +720,30 @@ TEST_P(TcpSocketTest, TcpSCMPriority) {
   ASSERT_EQ(cmsg, nullptr);
 }
 
+TEST_P(TcpSocketTest, TimeWaitPollHUP) {
+  shutdown(s_, SHUT_RDWR);
+  ScopedThread t([&]() {
+    constexpr int kTimeout = 10000;
+    constexpr int16_t want_events = POLLHUP;
+    struct pollfd pfd = {
+        .fd = s_,
+        .events = want_events,
+    };
+    ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
+  });
+  shutdown(t_, SHUT_RDWR);
+  t.Join();
+  // At this point s_ should be in TIME-WAIT and polling for POLLHUP should
+  // return with 1 FD.
+  constexpr int kTimeout = 10000;
+  constexpr int16_t want_events = POLLHUP;
+  struct pollfd pfd = {
+      .fd = s_,
+      .events = want_events,
+  };
+  ASSERT_THAT(poll(&pfd, 1, kTimeout), SyscallSucceedsWithValue(1));
+}
+
 INSTANTIATE_TEST_SUITE_P(AllInetTests, TcpSocketTest,
                          ::testing::Values(AF_INET, AF_INET6));