Add TFO support to AsyncSSLSocket

[folly.git] / folly / io / async / test / BlockingSocket.h

diff --git a/folly/io/async/test/BlockingSocket.h b/folly/io/async/test/BlockingSocket.h
index 3830648e61d43a9e3d0407857a3f8287f9fc0c1e..b3713fc241bdfe55ec484a64acce1608934f0603 100644 (file)
--- a/folly/io/async/test/BlockingSocket.h
+++ b/folly/io/async/test/BlockingSocket.h
@@ -16,36 +16,41 @@
 #pragma once
 
 #include <folly/Optional.h>
-#include <folly/io/async/SSLContext.h>
-#include <folly/io/async/AsyncSocket.h>
 #include <folly/io/async/AsyncSSLSocket.h>
+#include <folly/io/async/AsyncSocket.h>
+#include <folly/io/async/SSLContext.h>
 
 class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
                        public folly::AsyncTransportWrapper::ReadCallback,
-                       public folly::AsyncTransportWrapper::WriteCallback
-{
+                       public folly::AsyncTransportWrapper::WriteCallback {
  public:
   explicit BlockingSocket(int fd)
-    : sock_(new folly::AsyncSocket(&eventBase_, fd)) {
-  }
+      : sock_(new folly::AsyncSocket(&eventBase_, fd)) {}
 
-  BlockingSocket(folly::SocketAddress address,
-                 std::shared_ptr<folly::SSLContext> sslContext)
-    : sock_(sslContext ? new folly::AsyncSSLSocket(sslContext, &eventBase_) :
-            new folly::AsyncSocket(&eventBase_)),
-    address_(address) {}
+  BlockingSocket(
+      folly::SocketAddress address,
+      std::shared_ptr<folly::SSLContext> sslContext)
+      : sock_(
+            sslContext ? new folly::AsyncSSLSocket(sslContext, &eventBase_)
+                       : new folly::AsyncSocket(&eventBase_)),
+        address_(address) {}
 
   explicit BlockingSocket(folly::AsyncSocket::UniquePtr socket)
       : sock_(std::move(socket)) {
     sock_->attachEventBase(&eventBase_);
   }
 
+  void enableTFO() {
+    sock_->enableTFO();
+  }
+
   void setAddress(folly::SocketAddress address) {
     address_ = address;
   }
 
-  void open() {
-    sock_->connect(this, address_);
+  void open(
+      std::chrono::milliseconds timeout = std::chrono::milliseconds::zero()) {
+    sock_->connect(this, address_, timeout.count());
     eventBase_.loop();
     if (err_.hasValue()) {
       throw err_.value();
@@ -54,7 +59,9 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
   void close() {
     sock_->close();
   }
-  void closeWithReset() { sock_->closeWithReset(); }
+  void closeWithReset() {
+    sock_->closeWithReset();
+  }
 
   int32_t write(uint8_t const* buf, size_t len) {
     sock_->write(this, buf, len);
@@ -67,11 +74,11 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
 
   void flush() {}
 
-  int32_t readAll(uint8_t *buf, size_t len) {
+  int32_t readAll(uint8_t* buf, size_t len) {
     return readHelper(buf, len, true);
   }
 
-  int32_t read(uint8_t *buf, size_t len) {
+  int32_t read(uint8_t* buf, size_t len) {
     return readHelper(buf, len, false);
   }
 
@@ -83,7 +90,7 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
   folly::EventBase eventBase_;
   folly::AsyncSocket::UniquePtr sock_;
   folly::Optional<folly::AsyncSocketException> err_;
-  uint8_t *readBuf_{nullptr};
+  uint8_t* readBuf_{nullptr};
   size_t readLen_{0};
   folly::SocketAddress address_;
 
@@ -102,18 +109,18 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
       sock_->setReadCB(nullptr);
     }
   }
-  void readEOF() noexcept override {
-  }
+  void readEOF() noexcept override {}
   void readErr(const folly::AsyncSocketException& ex) noexcept override {
     err_ = ex;
   }
   void writeSuccess() noexcept override {}
-  void writeErr(size_t /* bytesWritten */,
-                const folly::AsyncSocketException& ex) noexcept override {
+  void writeErr(
+      size_t /* bytesWritten */,
+      const folly::AsyncSocketException& ex) noexcept override {
     err_ = ex;
   }
 
-  int32_t readHelper(uint8_t *buf, size_t len, bool all) {
+  int32_t readHelper(uint8_t* buf, size_t len, bool all) {
     if (!sock_->good()) {
       return 0;
     }
@@ -132,8 +139,8 @@ class BlockingSocket : public folly::AsyncSocket::ConnectCallback,
       throw err_.value();
     }
     if (all && readLen_ > 0) {
-      throw folly::AsyncSocketException(folly::AsyncSocketException::UNKNOWN,
-                                        "eof");
+      throw folly::AsyncSocketException(
+          folly::AsyncSocketException::UNKNOWN, "eof");
     }
     return len - readLen_;
   }