From 106a2a56bbf2dea08947919539944f582d283a20 Mon Sep 17 00:00:00 2001
From: Kevin Lewi <klewi@fb.com>
Date: Tue, 12 Jul 2016 13:11:39 -0700
Subject: [PATCH] Adding the ability to check for whether TCP fast open
 succeeded on a socket

Summary: Modifies AsyncSocket to check if TFO succeeded. Also added checks to AsyncSocketTest2 to check if TFO succeeded when it should.

Reviewed By: siyengar

Differential Revision: D3540330

fbshipit-source-id: 8a4b64fdb040dea73ba264b8e3dfff4d717fd96f
---
 folly/detail/SocketFastOpen.cpp          | 17 +++++++++++++++++
 folly/detail/SocketFastOpen.h            |  5 +++++
 folly/io/async/AsyncSocket.cpp           | 12 ++++++++++++
 folly/io/async/AsyncSocket.h             |  9 +++++++++
 folly/io/async/test/AsyncSocketTest2.cpp |  4 ++++
 5 files changed, 47 insertions(+)

diff --git a/folly/detail/SocketFastOpen.cpp b/folly/detail/SocketFastOpen.cpp
index e8a71f24..31917eab 100644
--- a/folly/detail/SocketFastOpen.cpp
+++ b/folly/detail/SocketFastOpen.cpp
@@ -46,6 +46,18 @@ int tfo_enable(int sockfd, size_t max_queue_size) {
       sockfd, SOL_TCP, TCP_FASTOPEN, &max_queue_size, sizeof(max_queue_size));
 }
 
+bool tfo_succeeded(int sockfd) {
+  // Call getsockopt to check if TFO was used.
+  struct tcp_info info;
+  socklen_t info_len = sizeof(info);
+  errno = 0;
+  if (getsockopt(sockfd, IPPROTO_TCP, TCP_INFO, &info, &info_len) != 0) {
+    // errno is set from getsockopt
+    return false;
+  }
+  return info.tcpi_options & TCPI_OPT_SYN_DATA;
+}
+
 #else
 
 ssize_t tfo_sendmsg(int sockfd, const struct msghdr* msg, int flags) {
@@ -58,6 +70,11 @@ int tfo_enable(int sockfd, size_t max_queue_size) {
   return -1;
 }
 
+bool tfo_succeeded(int sockfd) {
+  errno = EOPNOTSUPP;
+  return false;
+}
+
 #endif
 }
 }
diff --git a/folly/detail/SocketFastOpen.h b/folly/detail/SocketFastOpen.h
index d01e89b9..fcd997b6 100644
--- a/folly/detail/SocketFastOpen.h
+++ b/folly/detail/SocketFastOpen.h
@@ -37,5 +37,10 @@ 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);
+
+/**
+ * Check if TFO succeeded in being used.
+ */
+bool tfo_succeeded(int sockfd);
 }
 }
diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index 6bb9126e..be71ee73 100644
--- a/folly/io/async/AsyncSocket.cpp
+++ b/folly/io/async/AsyncSocket.cpp
@@ -1762,6 +1762,18 @@ AsyncSocket::sendSocketMessage(int fd, struct msghdr* msg, int msg_flags) {
     msg->msg_namelen = len;
     totalWritten = tfoSendMsg(fd_, msg, msg_flags);
     if (totalWritten >= 0) {
+      // Call tfo_succeeded to check if TFO was used.
+      tfoSucceeded_ = detail::tfo_succeeded(fd_);
+      if (errno != 0) {
+        auto errnoCopy = errno;
+        AsyncSocketException ex(
+            AsyncSocketException::INTERNAL_ERROR,
+            withAddr("error calling tfo_succeeded"),
+            errnoCopy);
+        return WriteResult(
+            WRITE_ERROR, folly::make_unique<AsyncSocketException>(ex));
+      }
+
       tfoFinished_ = true;
       state_ = StateEnum::ESTABLISHED;
       handleInitialReadWrite();
diff --git a/folly/io/async/AsyncSocket.h b/folly/io/async/AsyncSocket.h
index 36949725..35b3815c 100644
--- a/folly/io/async/AsyncSocket.h
+++ b/folly/io/async/AsyncSocket.h
@@ -431,6 +431,14 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     return tfoFinished_;
   }
 
+  /**
+   * Returns whether or not TFO "worked", or, succeeded
+   * in actually being used.
+   */
+  bool getTFOSucceeded() const {
+    return tfoSucceeded_;
+  }
+
   // Methods controlling socket options
 
   /**
@@ -899,6 +907,7 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   bool tfoEnabled_{false};
   bool tfoAttempted_{false};
   bool tfoFinished_{false};
+  bool tfoSucceeded_{false};
 };
 #ifdef _MSC_VER
 #pragma vtordisp(pop)
diff --git a/folly/io/async/test/AsyncSocketTest2.cpp b/folly/io/async/test/AsyncSocketTest2.cpp
index 602e9862..f7ca991a 100644
--- a/folly/io/async/test/AsyncSocketTest2.cpp
+++ b/folly/io/async/test/AsyncSocketTest2.cpp
@@ -2406,6 +2406,7 @@ TEST(AsyncSocketTest, ConnectTFO) {
   ASSERT_EQ(1, rcb.buffers.size());
   ASSERT_EQ(sizeof(buf), rcb.buffers[0].length);
   EXPECT_EQ(0, memcmp(rcb.buffers[0].buffer, buf.data(), buf.size()));
+  EXPECT_EQ(socket->getTFOSucceeded(), socket->getTFOFinished());
 }
 
 /**
@@ -2473,6 +2474,7 @@ TEST(AsyncSocketTest, ConnectWriteAndCloseNowTFO) {
   // Loop, although there shouldn't be anything to do.
   evb.loop();
 
+  EXPECT_EQ(socket->getTFOSucceeded(), socket->getTFOFinished());
   CHECK_EQ(ccb.state, STATE_SUCCEEDED);
 
   ASSERT_TRUE(socket->isClosedBySelf());
@@ -2500,6 +2502,7 @@ TEST(AsyncSocketTest, ConnectAndCloseTFO) {
 
   // Make sure the connection was aborted
   CHECK_EQ(ccb.state, STATE_SUCCEEDED);
+  EXPECT_EQ(socket->getTFOSucceeded(), socket->getTFOFinished());
 
   ASSERT_TRUE(socket->isClosedBySelf());
   ASSERT_FALSE(socket->isClosedByPeer());
@@ -2760,6 +2763,7 @@ TEST(AsyncSocketTest, ConnectTFOWithBigData) {
 
   t.join();
 
+  EXPECT_EQ(socket->getTFOSucceeded(), socket->getTFOFinished());
   EXPECT_EQ(STATE_SUCCEEDED, write.state);
   EXPECT_EQ(0, memcmp(readBuf.data(), sendBuf->data(), readBuf.size()));
   EXPECT_EQ(STATE_SUCCEEDED, rcb.state);
-- 
2.34.1