From 6f8d37dc510dfbbf8beacc67b23d937bff69182d Mon Sep 17 00:00:00 2001
From: Subodh Iyengar <subodh@fb.com>
Date: Sat, 4 Mar 2017 17:46:28 -0800
Subject: [PATCH] Add getTotalConnectTimeout method

Summary:
AsyncSSLSocket can be connected using connect with
total timeouts.

This adds a method to get the total timeout that
was set.

Reviewed By: ngoyal, yfeldblum

Differential Revision: D4656331

fbshipit-source-id: a55ad9f081449d358b8133e9598a2063f625a2e6
---
 folly/io/async/AsyncSSLSocket.cpp          |  1 +
 folly/io/async/AsyncSSLSocket.h            | 11 +++++++++++
 folly/io/async/test/AsyncSSLSocketTest.cpp |  3 ++-
 folly/io/async/test/BlockingSocket.h       |  9 +++++++++
 4 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/folly/io/async/AsyncSSLSocket.cpp b/folly/io/async/AsyncSSLSocket.cpp
index 7c72a584..7b4e16bb 100644
--- a/folly/io/async/AsyncSSLSocket.cpp
+++ b/folly/io/async/AsyncSSLSocket.cpp
@@ -666,6 +666,7 @@ void AsyncSSLSocket::connect(
   assert(state_ == StateEnum::UNINIT);
   assert(sslState_ == STATE_UNINIT);
   noTransparentTls_ = true;
+  totalConnectTimeout_ = totalConnectTimeout;
   AsyncSSLSocketConnector* connector =
       new AsyncSSLSocketConnector(this, callback, totalConnectTimeout.count());
   AsyncSocket::connect(
diff --git a/folly/io/async/AsyncSSLSocket.h b/folly/io/async/AsyncSSLSocket.h
index 4edbec58..99728ddd 100644
--- a/folly/io/async/AsyncSSLSocket.h
+++ b/folly/io/async/AsyncSSLSocket.h
@@ -693,6 +693,16 @@ class AsyncSSLSocket : public virtual AsyncSocket {
     return sessionResumptionAttempted_;
   }
 
+  /**
+   * If the SSL socket was used to connect as well
+   * as establish an SSL connection, this gives the total
+   * timeout for the connect + SSL connection that was
+   * set.
+   */
+  std::chrono::milliseconds getTotalConnectTimeout() const {
+    return totalConnectTimeout_;
+  }
+
  private:
 
   void init();
@@ -845,6 +855,7 @@ class AsyncSSLSocket : public virtual AsyncSocket {
   std::chrono::steady_clock::time_point handshakeEndTime_;
   std::chrono::milliseconds handshakeConnectTimeout_{0};
   bool sessionResumptionAttempted_{false};
+  std::chrono::milliseconds totalConnectTimeout_{0};
 
   std::unique_ptr<IOBuf> preReceivedData_;
 };
diff --git a/folly/io/async/test/AsyncSSLSocketTest.cpp b/folly/io/async/test/AsyncSSLSocketTest.cpp
index b5c6430a..41284b6e 100644
--- a/folly/io/async/test/AsyncSSLSocketTest.cpp
+++ b/folly/io/async/test/AsyncSSLSocketTest.cpp
@@ -152,7 +152,7 @@ TEST(AsyncSSLSocketTest, ConnectWriteReadClose) {
   // connect
   auto socket = std::make_shared<BlockingSocket>(server.getAddress(),
                                                  sslContext);
-  socket->open();
+  socket->open(std::chrono::milliseconds(10000));
 
   // write()
   uint8_t buf[128];
@@ -169,6 +169,7 @@ TEST(AsyncSSLSocketTest, ConnectWriteReadClose) {
   socket->close();
 
   cerr << "ConnectWriteReadClose test completed" << endl;
+  EXPECT_EQ(socket->getSSLSocket()->getTotalConnectTimeout().count(), 10000);
 }
 
 /**
diff --git a/folly/io/async/test/BlockingSocket.h b/folly/io/async/test/BlockingSocket.h
index 7f72ab24..0aeb4d0e 100644
--- a/folly/io/async/test/BlockingSocket.h
+++ b/folly/io/async/test/BlockingSocket.h
@@ -56,6 +56,7 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
       throw err_.value();
     }
   }
+
   void close() {
     sock_->close();
   }
@@ -86,6 +87,14 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
     return sock_->getFd();
   }
 
+  folly::AsyncSocket* getSocket() {
+    return sock_.get();
+  }
+
+  folly::AsyncSSLSocket* getSSLSocket() {
+    return dynamic_cast<folly::AsyncSSLSocket*>(sock_.get());
+  }
+
  private:
   folly::EventBase eventBase_;
   folly::AsyncSocket::UniquePtr sock_;
-- 
2.34.1