Proper cleanup of sockets that used REUSEPORT

Fixed a small logic error that broke proper accounting of MultiPortEndpoints.

PiperOrigin-RevId: 246502126
Change-Id: I1a7d6ea134f811612e545676212899a3707bc2c2

2 files changed, 35 insertions, 1 deletions

diff --git a/pkg/tcpip/stack/transport_demuxer.go b/pkg/tcpip/stack/transport_demuxer.go
index 66c564613..807c3ba5e 100644
--- a/pkg/tcpip/stack/transport_demuxer.go
+++ b/pkg/tcpip/stack/transport_demuxer.go
@@ -171,7 +171,7 @@ func (ep *multiPortEndpoint) singleRegisterEndpoint(t TransportEndpoint) {
 	// A new endpoint is added into endpointsArr and its index there is
 	// saved in endpointsMap. This will allows to remove endpoint from
 	// the array fast.
-	ep.endpointsMap[ep] = len(ep.endpointsArr)
+	ep.endpointsMap[t] = len(ep.endpointsArr)
 	ep.endpointsArr = append(ep.endpointsArr, t)
 }
 
diff --git a/test/syscalls/linux/socket_inet_loopback.cc b/test/syscalls/linux/socket_inet_loopback.cc
index f86a0f30c..b216d14cb 100644
--- a/test/syscalls/linux/socket_inet_loopback.cc
+++ b/test/syscalls/linux/socket_inet_loopback.cc
@@ -1040,6 +1040,40 @@ TEST_P(SocketMultiProtocolInetLoopbackTest, PortReuseTwoSockets) {
   }
 }
 
+// Check that when a socket was bound to an address with REUSEPORT and then
+// closed, we can bind a different socket to the same address without needing
+// REUSEPORT.
+TEST_P(SocketMultiProtocolInetLoopbackTest, NoReusePortFollowingReusePort) {
+  auto const& param = GetParam();
+  TestAddress const& test_addr = V4Loopback();
+  sockaddr_storage addr = test_addr.addr;
+
+  auto s = ASSERT_NO_ERRNO_AND_VALUE(Socket(test_addr.family(), param.type, 0));
+  int fd = s.get();
+  socklen_t addrlen = test_addr.addr_len;
+  int portreuse = 1;
+  ASSERT_THAT(
+      setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &portreuse, sizeof(portreuse)),
+      SyscallSucceeds());
+  ASSERT_THAT(bind(fd, reinterpret_cast<sockaddr*>(&addr), addrlen),
+              SyscallSucceeds());
+  ASSERT_THAT(getsockname(fd, reinterpret_cast<sockaddr*>(&addr), &addrlen),
+              SyscallSucceeds());
+  ASSERT_EQ(addrlen, test_addr.addr_len);
+
+  s.reset();
+
+  // Open a new socket and bind to the same address, but w/o REUSEPORT.
+  s = ASSERT_NO_ERRNO_AND_VALUE(Socket(test_addr.family(), param.type, 0));
+  fd = s.get();
+  portreuse = 0;
+  ASSERT_THAT(
+      setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &portreuse, sizeof(portreuse)),
+      SyscallSucceeds());
+  ASSERT_THAT(bind(fd, reinterpret_cast<sockaddr*>(&addr), addrlen),
+              SyscallSucceeds());
+}
+
 INSTANTIATE_TEST_SUITE_P(
     AllFamlies, SocketMultiProtocolInetLoopbackTest,
     ::testing::Values(ProtocolTestParam{"TCP", SOCK_STREAM},