From 3b917921d7feab08a94d54f29af5a57f35c9fc5e Mon Sep 17 00:00:00 2001
From: Bhasker Hariharan <bhaskerh@google.com>
Date: Thu, 21 Apr 2022 15:01:19 -0700
Subject: [PATCH] Fix race in listen.

Netstack has an internal StateBound state that a listening socket
transitions to before it actually transitions to StateListen. This
state is not handled by the dispatcher because normally a socket in
StateBound is never registered with the demuxer. But in case of listen
there is a small race window where the socket is in StateBound after
registering with dispatcher but before it transitions to Listen.

An incoming segment that hits the window causes the dispatcher to
panic. This change makes it so that we transition to Listen before
registering and if registration fails we now transition to a CLOSED
state same as what linux does.

See: https://github.com/google/gvisor/issues/7428

Fixes #7428

PiperOrigin-RevId: 443489134
---
 pkg/tcpip/transport/tcp/endpoint.go |  7 ++++-
 test/syscalls/linux/tcp_socket.cc   | 45 +++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go
index 79ec2b651..bea8049e6 100644
--- a/pkg/tcpip/transport/tcp/endpoint.go
+++ b/pkg/tcpip/transport/tcp/endpoint.go
@@ -2618,13 +2618,18 @@ func (e *endpoint) listen(backlog int) tcpip.Error {
 		return &tcpip.ErrInvalidEndpointState{}
 	}
 
+	// Setting this state after RegisterTransportEndpoint will result in a
+	// race where the endpoint is in Bound but reachable via the demuxer. Instead
+	// we set it to listen so that incoming packets will just be queued to the
+	// inbound segment queue by the TCP processor.
+	e.setEndpointState(StateListen)
 	// Register the endpoint.
 	if err := e.stack.RegisterTransportEndpoint(e.effectiveNetProtos, ProtocolNumber, e.TransportEndpointInfo.ID, e, e.boundPortFlags, e.boundBindToDevice); err != nil {
+		e.transitionToStateCloseLocked()
 		return err
 	}
 
 	e.isRegistered = true
-	e.setEndpointState(StateListen)
 
 	// The queue may be non-zero when we're restoring the endpoint, and it
 	// may be pre-populated with some previously accepted (but not Accepted)
diff --git a/test/syscalls/linux/tcp_socket.cc b/test/syscalls/linux/tcp_socket.cc
index 7891abea1..602845ae5 100644
--- a/test/syscalls/linux/tcp_socket.cc
+++ b/test/syscalls/linux/tcp_socket.cc
@@ -1064,6 +1064,51 @@ TEST_P(SimpleTcpSocketTest, NonBlockingConnectNoListener) {
               SyscallFailsWithErrno(ECONNABORTED));
 }
 
+TEST_P(SimpleTcpSocketTest, ListenConnectParallel) {
+  int family = GetParam();
+  sockaddr_storage addr =
+      ASSERT_NO_ERRNO_AND_VALUE(InetLoopbackAddr(GetParam()));
+  socklen_t addrlen = sizeof(addr);
+  constexpr int sock_type = SOCK_STREAM;
+
+  FileDescriptor l =
+      ASSERT_NO_ERRNO_AND_VALUE(Socket(family, sock_type, IPPROTO_TCP));
+  EXPECT_THAT(bind(l.get(), AsSockAddr(&addr), addrlen), SyscallSucceeds());
+
+  // Get the address bound by the listening socket.
+  EXPECT_THAT(getsockname(l.get(), AsSockAddr(&addr), &addrlen),
+              SyscallSucceeds());
+
+  constexpr int num_threads = 100;
+  ScopedThread t([&l]() {
+    absl::SleepFor(absl::Microseconds(1000));
+    EXPECT_THAT(listen(l.get(), num_threads), SyscallSucceeds());
+  });
+
+  // Initiate connects in a separate thread.
+  std::vector<ScopedThread*> threads;
+  for (int i = 0; i < num_threads; i++) {
+    ScopedThread t([&addr, &addrlen, family]() {
+      const FileDescriptor c = ASSERT_NO_ERRNO_AND_VALUE(
+          Socket(family, SOCK_STREAM | SOCK_NONBLOCK, IPPROTO_TCP));
+
+      // Now connect to the bound address and this should fail as nothing
+      // is listening on the bound address.
+      EXPECT_THAT(RetryEINTR(connect)(c.get(), AsSockAddr(&addr), addrlen),
+                  SyscallFailsWithErrno(EINPROGRESS));
+      // Wait for the connect to fail or succeed as it can race with the socket
+      // listening.
+      struct pollfd poll_fd = {c.get(), POLLERR | POLLOUT, 0};
+      EXPECT_THAT(RetryEINTR(poll)(&poll_fd, 1, 1000),
+                  SyscallSucceedsWithValue(1));
+    });
+    threads.push_back(&t);
+  }
+  for (auto t : threads) {
+    t->Join();
+  }
+}
+
 TEST_P(SimpleTcpSocketTest, NonBlockingConnectNoListenerRead) {
   sockaddr_storage addr =
       ASSERT_NO_ERRNO_AND_VALUE(InetLoopbackAddr(GetParam()));