Set TOS for AsyncServer listener socket

Summary: Added set/get APIs to configure TOS for listener sockets. The setListenerTos() sets the TOS for the server sockets and all accepted connections are expected to inherit the same. These APIs would be used by higher layers (like thrift server) to set the TOS on the server socket.

Reviewed By: jmswen

Differential Revision: D28651968

fbshipit-source-id: 30f251970269155adbf5e88e1079096dbeceb216

folly/io/async/AsyncServerSocket.cpp

@@ -791,6 +791,34 @@ void AsyncServerSocket::setTosReflect(bool enable) {
   tosReflect_ = true;
 }
 
+void AsyncServerSocket::setListenerTos(uint32_t tos) {
+  if (!kIsLinux || tos == 0) {
+    listenerTos_ = 0;
+    return;
+  }
+
+  for (auto& handler : sockets_) {
+    if (handler.socket_ == NetworkSocket()) {
+      continue;
+    }
+
+    const auto proto =
+        (handler.addressFamily_ == AF_INET) ? IPPROTO_IP : IPPROTO_IPV6;
+    const auto optName =
+        (handler.addressFamily_ == AF_INET) ? IP_TOS : IPV6_TCLASS;
+
+    int ret =
+        netops::setsockopt(handler.socket_, proto, optName, &tos, sizeof(tos));
+
+    if (ret == 0) {
+      VLOG(10) << "Set TOS " << tos << " for for socket " << handler.socket_;
+    } else {
+      folly::throwSystemError(errno, "failed to set TOS for socket");
+    }
+  }
+  listenerTos_ = tos;
+}
+
 void AsyncServerSocket::setupSocket(NetworkSocket fd, int family) {
   // Put the socket in non-blocking mode
   if (netops::set_socket_non_blocking(fd) != 0) {

folly/io/async/AsyncServerSocket.h

@@ -633,6 +633,13 @@ class AsyncServerSocket : public DelayedDestruction, public AsyncSocketBase {
 
   bool getTosReflect() { return tosReflect_; }
 
+  /**
+   * Set/Get default TOS for listener socket
+   */
+  void setListenerTos(uint32_t tos);
+
+  uint32_t getListenerTos() const { return listenerTos_; }
+
   /**
    * Get the number of connections dropped by the AsyncServerSocket
    */
@@ -956,6 +963,7 @@ class AsyncServerSocket : public DelayedDestruction, public AsyncSocketBase {
   std::weak_ptr<ShutdownSocketSet> wShutdownSocketSet_;
   ConnectionEventCallback* connectionEventCallback_{nullptr};
   bool tosReflect_{false};
+  uint32_t listenerTos_{0};
   bool zeroCopyVal_{false};
   folly::observer::AtomicObserver<std::chrono::nanoseconds> queueTimeout_{
       folly::observer::makeStaticObserver(std::chrono::nanoseconds::zero())};

folly/io/async/test/AsyncSocketTest.cpp

@@ -121,4 +121,58 @@ TEST(AsyncSocketTest, tosReflect) {
   ASSERT_EQ(value, 1);
 }
 
+TEST(AsyncSocketTest, listenerTosV6) {
+  EventBase base;
+  auto server1 = AsyncServerSocket::newSocket(&base);
+  server1->bind(0);
+  server1->listen(10);
+  auto fd = server1->getNetworkSocket();
+
+  // Verify if Listener TOS is disabled by default
+  EXPECT_FALSE(server1->getListenerTos());
+  int value;
+  socklen_t valueLength = sizeof(value);
+  int rc =
+      netops::getsockopt(fd, IPPROTO_IPV6, IPV6_TCLASS, &value, &valueLength);
+  ASSERT_EQ(rc, 0);
+  ASSERT_EQ(value, 0);
+
+  // Set listener Tos to 116 (0x74, represents dscp 29)
+  server1->setListenerTos(116);
+
+  // Verify if listener DSCP is set now
+  EXPECT_TRUE(server1->getListenerTos());
+  rc = netops::getsockopt(fd, IPPROTO_IPV6, IPV6_TCLASS, &value, &valueLength);
+  ASSERT_EQ(rc, 0);
+  ASSERT_EQ(value, 116);
+}
+
+TEST(AsyncSocketTest, listenerTosV4) {
+  EventBase base;
+  auto server1 = AsyncServerSocket::newSocket(&base);
+  folly::IPAddress ip("127.0.0.1");
+  std::vector<folly::IPAddress> serverIp;
+  serverIp.push_back(ip);
+  server1->bind(serverIp, 0);
+  server1->listen(10);
+  auto fd = server1->getNetworkSocket();
+
+  // Verify if Listener TOS is disabled by default
+  EXPECT_FALSE(server1->getListenerTos());
+  int value;
+  socklen_t valueLength = sizeof(value);
+  int rc = netops::getsockopt(fd, IPPROTO_IP, IP_TOS, &value, &valueLength);
+  ASSERT_EQ(rc, 0);
+  ASSERT_EQ(value, 0);
+
+  // Set listener Tos to 140 (0x8c, represents dscp 35)
+  server1->setListenerTos(140);
+
+  // Verify if listener DSCP is set now
+  EXPECT_TRUE(server1->getListenerTos());
+  rc = netops::getsockopt(fd, IPPROTO_IP, IP_TOS, &value, &valueLength);
+  ASSERT_EQ(rc, 0);
+  ASSERT_EQ(value, 140);
+}
+
 } // namespace folly

folly/io/async/test/AsyncSocketTest2.cpp

@@ -6209,6 +6209,131 @@ TEST(AsyncSocketTest, V4TosReflectTest) {
   ASSERT_EQ(rc, 0);
   ASSERT_EQ(value, 0x2c);
 }
+
+TEST(AsyncSocketTest, V6AcceptedTosTest) {
+  EventBase eventBase;
+
+  // This test verifies if the ListenerTos set on a socket is
+  // propagated properly to accepted socket connections
+
+  // Create a server socket
+  std::shared_ptr<AsyncServerSocket> serverSocket(
+      AsyncServerSocket::newSocket(&eventBase));
+  folly::IPAddress ip("::1");
+  std::vector<folly::IPAddress> serverIp;
+  serverIp.push_back(ip);
+  serverSocket->bind(serverIp, 0);
+  serverSocket->listen(16);
+  folly::SocketAddress serverAddress;
+  serverSocket->getAddress(&serverAddress);
+
+  // Set listener TOS to 0x74 i.e. dscp 29
+  serverSocket->setListenerTos(0x74);
+
+  // Add a callback to accept one connection then stop the loop
+  TestAcceptCallback acceptCallback;
+  acceptCallback.setConnectionAcceptedFn(
+      [&](NetworkSocket /* fd */, const folly::SocketAddress& /* addr */) {
+        serverSocket->removeAcceptCallback(&acceptCallback, &eventBase);
+      });
+  acceptCallback.setAcceptErrorFn([&](const std::exception& /* ex */) {
+    serverSocket->removeAcceptCallback(&acceptCallback, &eventBase);
+  });
+  serverSocket->addAcceptCallback(&acceptCallback, &eventBase);
+  serverSocket->startAccepting();
+
+  // Create a client socket, setsockopt() the TOS before connecting
+  auto clientThread = [](std::shared_ptr<AsyncSocket>& clientSock,
+                         ConnCallback* ccb,
+                         EventBase* evb,
+                         folly::SocketAddress sAddr) {
+    clientSock = AsyncSocket::newSocket(evb);
+    SocketOptionKey v6Opts = {IPPROTO_IPV6, IPV6_TCLASS};
+    SocketOptionMap optionMap;
+    optionMap.insert({v6Opts, 0x2c});
+    SocketAddress bindAddr("0.0.0.0", 0);
+    clientSock->connect(ccb, sAddr, 30, optionMap, bindAddr);
+  };
+
+  std::shared_ptr<AsyncSocket> socket(nullptr);
+  ConnCallback cb;
+  clientThread(socket, &cb, &eventBase, serverAddress);
+
+  eventBase.loop();
+
+  // Verify if the connection is accepted and if the accepted socket has
+  // setsockopt on the TOS for the same value that the listener was set to
+  auto fd = acceptCallback.getEvents()->at(1).fd;
+  ASSERT_NE(fd, NetworkSocket());
+  int value;
+  socklen_t valueLength = sizeof(value);
+  int rc =
+      netops::getsockopt(fd, IPPROTO_IPV6, IPV6_TCLASS, &value, &valueLength);
+  ASSERT_EQ(rc, 0);
+  ASSERT_EQ(value, 0x74);
+}
+
+TEST(AsyncSocketTest, V4AcceptedTosTest) {
+  EventBase eventBase;
+
+  // This test verifies if the ListenerTos set on a socket is
+  // propagated properly to accepted socket connections
+
+  // Create a server socket
+  std::shared_ptr<AsyncServerSocket> serverSocket(
+      AsyncServerSocket::newSocket(&eventBase));
+  folly::IPAddress ip("127.0.0.1");
+  std::vector<folly::IPAddress> serverIp;
+  serverIp.push_back(ip);
+  serverSocket->bind(serverIp, 0);
+  serverSocket->listen(16);
+  folly::SocketAddress serverAddress;
+  serverSocket->getAddress(&serverAddress);
+
+  // Set listener TOS to 0x74 i.e. dscp 29
+  serverSocket->setListenerTos(0x74);
+
+  // Add a callback to accept one connection then stop the loop
+  TestAcceptCallback acceptCallback;
+  acceptCallback.setConnectionAcceptedFn(
+      [&](NetworkSocket /* fd */, const folly::SocketAddress& /* addr */) {
+        serverSocket->removeAcceptCallback(&acceptCallback, &eventBase);
+      });
+  acceptCallback.setAcceptErrorFn([&](const std::exception& /* ex */) {
+    serverSocket->removeAcceptCallback(&acceptCallback, &eventBase);
+  });
+  serverSocket->addAcceptCallback(&acceptCallback, &eventBase);
+  serverSocket->startAccepting();
+
+  // Create a client socket, setsockopt() the TOS before connecting
+  auto clientThread = [](std::shared_ptr<AsyncSocket>& clientSock,
+                         ConnCallback* ccb,
+                         EventBase* evb,
+                         folly::SocketAddress sAddr) {
+    clientSock = AsyncSocket::newSocket(evb);
+    SocketOptionKey v4Opts = {IPPROTO_IP, IP_TOS};
+    SocketOptionMap optionMap;
+    optionMap.insert({v4Opts, 0x2c});
+    SocketAddress bindAddr("0.0.0.0", 0);
+    clientSock->connect(ccb, sAddr, 30, optionMap, bindAddr);
+  };
+
+  std::shared_ptr<AsyncSocket> socket(nullptr);
+  ConnCallback cb;
+  clientThread(socket, &cb, &eventBase, serverAddress);
+
+  eventBase.loop();
+
+  // Verify if the connection is accepted and if the accepted socket has
+  // setsockopt on the TOS for the same value that the listener was set to
+  auto fd = acceptCallback.getEvents()->at(1).fd;
+  ASSERT_NE(fd, NetworkSocket());
+  int value;
+  socklen_t valueLength = sizeof(value);
+  int rc = netops::getsockopt(fd, IPPROTO_IP, IP_TOS, &value, &valueLength);
+  ASSERT_EQ(rc, 0);
+  ASSERT_EQ(value, 0x74);
+}
 #endif
 
 #if defined(__linux__)