From 9e168d6924e7f11f9e4c56629f87998b1b7cf304 Mon Sep 17 00:00:00 2001
From: Wenzhe Lu <wenzhelu@meta.com>
Date: Tue, 26 Sep 2023 23:58:01 -0700
Subject: [PATCH] make it possible to disable SO_REUSEADDR in
 folly::AsyncServerSocket

Summary:
Discussion thread: https://fb.workplace.com/groups/560979627394613/permalink/2764923823666838/

It's unfortunately possible that TW Agent assign same port to 2 stacked TW task, due to some race condition when those tasks restarted about the same time.

We would like to disable SO_REUSEADDR so that when above scenario happen `AsyncServerSocket::bind` throws exception.

Reviewed By: Gownta

Differential Revision: D49623771

fbshipit-source-id: b40dfbe7c7ec62d18ad1f0d292fd31099c8705e1
---
 folly/io/async/AsyncServerSocket.cpp    | 22 ++++++++++++-
 folly/io/async/AsyncServerSocket.h      |  9 ++++++
 folly/io/async/test/AsyncSocketTest.cpp | 43 +++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/folly/io/async/AsyncServerSocket.cpp b/folly/io/async/AsyncServerSocket.cpp
index 67c425e2117..c06b0e73b21 100644
--- a/folly/io/async/AsyncServerSocket.cpp
+++ b/folly/io/async/AsyncServerSocket.cpp
@@ -568,6 +568,26 @@ void AsyncServerSocket::bind(uint16_t port) {
   }
 }
 
+void AsyncServerSocket::setEnableReuseAddr(bool enable) {
+  enableReuseAddr_ = enable;
+  for (auto& handler : sockets_) {
+    if (handler.socket_ == NetworkSocket()) {
+      continue;
+    }
+
+    int val = (enable) ? 1 : 0;
+    if (netops::setsockopt(
+            handler.socket_, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val)) !=
+        0) {
+      auto errnoCopy = errno;
+      LOG(ERROR) << "failed to set SO_REUSEADDR on async server socket "
+                 << errnoCopy;
+      folly::throwSystemErrorExplicit(
+          errnoCopy, "failed to set SO_REUSEADDR on async server socket");
+    }
+  }
+}
+
 void AsyncServerSocket::listen(int backlog) {
   if (eventBase_) {
     eventBase_->dcheckIsInEventBaseThread();
@@ -830,7 +850,7 @@ void AsyncServerSocket::setupSocket(NetworkSocket fd, int family) {
   // Set reuseaddr to avoid 2MSL delay on server restart
   int one = 1;
   // AF_UNIX does not support SO_REUSEADDR, setting this would confuse Windows
-  if (family != AF_UNIX &&
+  if (family != AF_UNIX && enableReuseAddr_ &&
       netops::setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) !=
           0) {
     auto errnoCopy = errno;
diff --git a/folly/io/async/AsyncServerSocket.h b/folly/io/async/AsyncServerSocket.h
index 4b845cd31ab..b8d22a236e5 100644
--- a/folly/io/async/AsyncServerSocket.h
+++ b/folly/io/async/AsyncServerSocket.h
@@ -765,6 +765,13 @@ class AsyncServerSocket : public DelayedDestruction, public AsyncSocketBase {
     }
   }
 
+  /**
+   * Set whether or not SO_REUSEADDR should be enabled on the server socket,
+   * allowing multiple sockets binds to the same <address>:<port>
+   * It's enabled by default.
+   */
+  void setEnableReuseAddr(bool enable);
+
   /**
    * Get whether or not SO_REUSEPORT is enabled on the server socket.
    */
@@ -1004,6 +1011,8 @@ class AsyncServerSocket : public DelayedDestruction, public AsyncSocketBase {
   CallbackAssignFunction callbackAssignFunc_;
   bool keepAliveEnabled_;
   bool reusePortEnabled_{false};
+  // SO_REUSEADDR is enabled by default
+  bool enableReuseAddr_{true};
   bool closeOnExec_;
   bool tfo_{false};
   bool noTransparentTls_{false};
diff --git a/folly/io/async/test/AsyncSocketTest.cpp b/folly/io/async/test/AsyncSocketTest.cpp
index f038eefc05d..79d2455a16a 100644
--- a/folly/io/async/test/AsyncSocketTest.cpp
+++ b/folly/io/async/test/AsyncSocketTest.cpp
@@ -69,6 +69,49 @@ TEST(AsyncSocketTest, REUSEPORT) {
   serverSocket2->startAccepting();
 }
 
+TEST(AsyncSocketTest, DisableReuseAddr) {
+  EventBase base;
+  auto serverSocket = AsyncServerSocket::newSocket(&base);
+  serverSocket->setEnableReuseAddr(false /* enable */);
+  // idempotent
+  serverSocket->setEnableReuseAddr(false /* enable */);
+  serverSocket->setEnableReuseAddr(false /* enable */);
+  serverSocket->bind(0);
+
+  SocketAddress address;
+  serverSocket->getAddress(&address);
+  int port = address.getPort();
+
+  auto serverSocket2 = AsyncServerSocket::newSocket(&base);
+  serverSocket2->setEnableReuseAddr(false /* enable */);
+  // idempotent
+  serverSocket2->setEnableReuseAddr(false /* enable */);
+  serverSocket2->setEnableReuseAddr(false /* enable */);
+  EXPECT_THROW(serverSocket2->bind(port), std::system_error);
+  // it's ok to bind to a different port
+  serverSocket2->bind(0);
+}
+
+TEST(AsyncSocketTest, EnableThenDisableReuseAddr) {
+  EventBase base;
+  auto serverSocket = AsyncServerSocket::newSocket(&base);
+  serverSocket->bind(0);
+
+  SocketAddress address;
+  serverSocket->getAddress(&address);
+  int port = address.getPort();
+
+  auto serverSocket2 = AsyncServerSocket::newSocket(&base);
+  // defaulty SO_REUSEADDR enabled so can bind to same port
+  serverSocket2->bind(port);
+  serverSocket2->setEnableReuseAddr(false /* enable */);
+  serverSocket->setEnableReuseAddr(false /* enable */);
+
+  EXPECT_THROW(serverSocket2->bind(port), std::system_error);
+  // it's ok to bind to a different port
+  serverSocket2->bind(0);
+}
+
 TEST(AsyncSocketTest, v4v6samePort) {
   EventBase base;
   auto serverSocket = AsyncServerSocket::newSocket(&base);