Add support for TFO connections

Summary:
This adds support to establish connections
over TFO.

The API introduced here retains the same
connect() + write() api that clients currently
use.

If enableTFO() is called then the connect will
be deferred to the first write. This only works
with request response protocols since a write
must trigger the connect. There is a tradeoff here
for the simpler API, and we can address this with
other signals such as a short timeout in the future.

Even though the client might enable TFO, the program
might run on machines without TFO support.
There were 2 choices for supporting machines where
TFO might not be enabled:
1. Fallback to normal connect if tfo sendmsg fails
2. Check tfo supported on the machine before using it

Both these have their tradeoffs, however option 1 does
not require us to read from procfs in the common code
path.

Reviewed By: Orvid

Differential Revision: D3327480

fbshipit-source-id: 9ac3a0c7ad2d206b158fdc305641fedbd93aa44d

folly/detail/SocketFastOpen.cpp

@@ -20,21 +20,42 @@ namespace folly {
 namespace detail {
 
 #if FOLLY_ALLOW_TFO
-ssize_t tfo_sendto(
-    int sockfd,
-    const void* buf,
-    size_t len,
-    int flags,
-    const struct sockaddr* dest_addr,
-    socklen_t addrlen) {
+
+#include <netinet/tcp.h>
+#include <stdio.h>
+
+// Sometimes these flags are not present in the headers,
+// so define them if not present.
+#if !defined(MSG_FASTOPEN)
+#define MSG_FASTOPEN 0x20000000
+#endif
+
+#if !defined(TCP_FASTOPEN)
+#define TCP_FASTOPEN 23
+#endif
+
+ssize_t tfo_sendmsg(int sockfd, const struct msghdr* msg, int flags) {
   flags |= MSG_FASTOPEN;
-  return sendto(sockfd, buf, len, flags, dest_addr, addrlen);
+  return sendmsg(sockfd, msg, flags);
 }
 
 int tfo_enable(int sockfd, size_t max_queue_size) {
   return setsockopt(
       sockfd, SOL_TCP, TCP_FASTOPEN, &max_queue_size, sizeof(max_queue_size));
 }
+
+#else
+
+ssize_t tfo_sendmsg(int sockfd, const struct msghdr* msg, int flags) {
+  errno = EOPNOTSUPP;
+  return -1;
+}
+
+int tfo_enable(int sockfd, size_t max_queue_size) {
+  errno = ENOPROTOOPT;
+  return -1;
+}
+
 #endif
 }
 }

folly/detail/SocketFastOpen.h

@@ -19,31 +19,23 @@
 #include <folly/portability/Sockets.h>
 #include <sys/types.h>
 
-#if !defined(FOLLY_ALLOW_TFO) && defined(TCP_FASTOPEN) && defined(MSG_FASTOPEN)
+#if !defined(FOLLY_ALLOW_TFO) && defined(__linux__) && !defined(__ANDROID__)
+// only allow for linux right now
 #define FOLLY_ALLOW_TFO 1
 #endif
 
 namespace folly {
 namespace detail {
 
-#if FOLLY_ALLOW_TFO
-
 /**
- * tfo_sendto has the same semantics as sendto, but is used to
+ * tfo_sendto has the same semantics as sendmsg, but is used to
  * send with TFO data.
  */
-ssize_t tfo_sendto(
-    int sockfd,
-    const void* buf,
-    size_t len,
-    int flags,
-    const struct sockaddr* dest_addr,
-    socklen_t addlen);
+ssize_t tfo_sendmsg(int sockfd, const struct msghdr* msg, int flags);
 
 /**
  * Enable TFO on a listening socket.
  */
 int tfo_enable(int sockfd, size_t max_queue_size);
-#endif
 }
 }

folly/io/async/AsyncSocket.h

@@ -18,6 +18,7 @@
 
 #include <folly/Optional.h>
 #include <folly/SocketAddress.h>
+#include <folly/detail/SocketFastOpen.h>
 #include <folly/io/IOBuf.h>
 #include <folly/io/ShutdownSocketSet.h>
 #include <folly/io/async/AsyncSocketException.h>
@@ -416,6 +417,20 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     return connectTimeout_;
   }
 
+  bool getTFOAttempted() const {
+    return tfoAttempted_;
+  }
+
+  /**
+   * Returns whether or not the attempt to use TFO
+   * finished successfully. This does not necessarily
+   * mean TFO worked, just that trying to use TFO
+   * succeeded.
+   */
+  bool getTFOFinished() const {
+    return tfoFinished_;
+  }
+
   // Methods controlling socket options
 
   /**
@@ -509,12 +524,24 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     peek_ = peek;
   }
 
+  /**
+   * Enables TFO behavior on the AsyncSocket if FOLLY_ALLOW_TFO
+   * is set.
+   */
+  void enableTFO() {
+    // No-op if folly does not allow tfo
+#if FOLLY_ALLOW_TFO
+    tfoEnabled_ = true;
+#endif
+  }
+
   enum class StateEnum : uint8_t {
     UNINIT,
     CONNECTING,
     ESTABLISHED,
     CLOSED,
-    ERROR
+    ERROR,
+    FAST_OPEN,
   };
 
   void setBufferCallback(BufferCallback* cb);
@@ -784,6 +811,20 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
       uint32_t* countWritten,
       uint32_t* partialWritten);
 
+  /**
+   * Sends the message over the socket using sendmsg
+   *
+   * @param msg       Message to send
+   * @param msg_flags Flags to pass to sendmsg
+   */
+  AsyncSocket::WriteResult sendSocketMessage(struct msghdr* msg, int msg_flags);
+
+  virtual ssize_t tfoSendMsg(int fd, struct msghdr* msg, int msg_flags);
+
+  int socketConnect(const struct sockaddr* addr, socklen_t len);
+
+  void scheduleConnectTimeoutAndRegisterForEvents();
+
   bool updateEventRegistration();
 
   /**
@@ -854,6 +895,9 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   std::chrono::milliseconds connectTimeout_{0};
 
   BufferCallback* bufferCallback_{nullptr};
+  bool tfoEnabled_{false};
+  bool tfoAttempted_{false};
+  bool tfoFinished_{false};
 };
 #ifdef _MSC_VER
 #pragma vtordisp(pop)

folly/io/async/test/AsyncSocketTest.h

@@ -72,6 +72,7 @@ class WriteCallback : public folly::AsyncTransportWrapper::WriteCallback {
 
   void writeErr(size_t bytesWritten,
                 const folly::AsyncSocketException& ex) noexcept override {
+    LOG(ERROR) << ex.what();
     state = STATE_FAILED;
     this->bytesWritten = bytesWritten;
     exception = ex;
@@ -205,8 +206,7 @@ class TestServer {
  public:
   // Create a TestServer.
   // This immediately starts listening on an ephemeral port.
-  TestServer()
-    : fd_(-1) {
+  explicit TestServer(bool enableTFO = false) : fd_(-1) {
     fd_ = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
     if (fd_ < 0) {
       throw folly::AsyncSocketException(
@@ -221,6 +221,11 @@ class TestServer {
           "non-blocking mode",
           errno);
     }
+    if (enableTFO) {
+#if FOLLY_ALLOW_TFO
+      folly::detail::tfo_enable(fd_, 100);
+#endif
+    }
     if (listen(fd_, 10) != 0) {
       throw folly::AsyncSocketException(
           folly::AsyncSocketException::INTERNAL_ERROR,

folly/io/async/test/BlockingSocket.h

@@ -40,6 +40,10 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
     sock_->attachEventBase(&eventBase_);
   }
 
+  void setAddress(folly::SocketAddress address) {
+    address_ = address;
+  }
+
   void open() {
     sock_->connect(this, address_);
     eventBase_.loop();
@@ -110,11 +114,15 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
   }
 
   int32_t readHelper(uint8_t *buf, size_t len, bool all) {
+    if (!sock_->good()) {
+      return 0;
+    }
+
     readBuf_ = buf;
     readLen_ = len;
     sock_->setReadCB(this);
     while (!err_ && sock_->good() && readLen_ > 0) {
-      eventBase_.loop();
+      eventBase_.loopOnce();
       if (!all) {
         break;
       }

folly/io/async/test/SocketClient.cpp

+/*
+ * Copyright 2016 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <folly/io/async/test/BlockingSocket.h>
+
+#include <folly/ExceptionWrapper.h>
+#include <gflags/gflags.h>
+
+using namespace folly;
+
+DEFINE_string(host, "localhost", "Host");
+DEFINE_int32(port, 0, "port");
+DEFINE_bool(tfo, false, "enable tfo");
+DEFINE_string(msg, "", "Message to send");
+
+int main(int argc, char** argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+
+  if (FLAGS_port == 0) {
+    LOG(ERROR) << "Must specify port";
+    exit(EXIT_FAILURE);
+  }
+
+  // Prep the socket
+  EventBase evb;
+  AsyncSocket::UniquePtr socket(new AsyncSocket(&evb));
+  socket->detachEventBase();
+
+  if (FLAGS_tfo) {
+#if FOLLY_ALLOW_TFO
+    socket->enableTFO();
+#endif
+  }
+
+  // Keep this around
+  auto sockAddr = socket.get();
+
+  BlockingSocket sock(std::move(socket));
+  SocketAddress addr;
+  addr.setFromHostPort(FLAGS_host, FLAGS_port);
+  sock.setAddress(addr);
+  sock.open();
+  LOG(INFO) << "connected to " << addr.getAddressStr();
+
+  sock.write((const uint8_t*)FLAGS_msg.data(), FLAGS_msg.size());
+
+  LOG(ERROR) << "TFO attempted: " << sockAddr->getTFOAttempted();
+  LOG(ERROR) << "TFO finished: " << sockAddr->getTFOFinished();
+
+  std::array<char, 1024> buf;
+  int32_t bytesRead = 0;
+  while ((bytesRead = sock.read((uint8_t*)buf.data(), buf.size())) != 0) {
+    std::cout << std::string(buf.data(), bytesRead);
+  }
+
+  sock.close();
+  return 0;
+}