Use std::chrono for timeouts to sslAccept and sslConn in AsyncSSLSocket

[folly.git] / folly / io / async / test / AsyncSSLSocketTest2.cpp

diff --git a/folly/io/async/test/AsyncSSLSocketTest2.cpp b/folly/io/async/test/AsyncSSLSocketTest2.cpp
index e33df0a127550dae7a6a91c7ab52faaf415d717e..4d766b68dd47968053c09ec00b46cee26a412e7e 100644 (file)
--- a/folly/io/async/test/AsyncSSLSocketTest2.cpp
+++ b/folly/io/async/test/AsyncSSLSocketTest2.cpp
@@ -17,9 +17,11 @@
 
 #include <pthread.h>
 
+#include <folly/futures/Promise.h>
 #include <folly/io/async/AsyncSSLSocket.h>
 #include <folly/io/async/EventBase.h>
 #include <folly/io/async/SSLContext.h>
+#include <folly/io/async/ScopedEventBaseThread.h>
 #include <folly/portability/GTest.h>
 
 using std::string;
@@ -31,44 +33,100 @@ using std::list;
 
 namespace folly {
 
+struct EvbAndContext {
+  EvbAndContext() {
+    ctx_.reset(new SSLContext());
+    ctx_->setOptions(SSL_OP_NO_TICKET);
+    ctx_->ciphers("ALL:!ADH:!LOW:!EXP:!MD5:@STRENGTH");
+  }
+
+  std::shared_ptr<AsyncSSLSocket> createSocket() {
+    return AsyncSSLSocket::newSocket(ctx_, getEventBase());
+  }
+
+  EventBase* getEventBase() {
+    return evb_.getEventBase();
+  }
+
+  void attach(AsyncSSLSocket& socket) {
+    socket.attachEventBase(getEventBase());
+    socket.attachSSLContext(ctx_);
+  }
+
+  folly::ScopedEventBaseThread evb_;
+  std::shared_ptr<SSLContext> ctx_;
+};
+
 class AttachDetachClient : public AsyncSocket::ConnectCallback,
                            public AsyncTransportWrapper::WriteCallback,
                            public AsyncTransportWrapper::ReadCallback {
  private:
-  EventBase *eventBase_;
+  // two threads here - we'll create the socket in one, connect
+  // in the other, and then read/write in the initial one
+  EvbAndContext t1_;
+  EvbAndContext t2_;
   std::shared_ptr<AsyncSSLSocket> sslSocket_;
-  std::shared_ptr<SSLContext> ctx_;
   folly::SocketAddress address_;
   char buf_[128];
   char readbuf_[128];
   uint32_t bytesRead_;
+  // promise to fulfill when done
+  folly::Promise<bool> promise_;
+
+  void detach() {
+    sslSocket_->detachEventBase();
+    sslSocket_->detachSSLContext();
+  }
+
  public:
-  AttachDetachClient(EventBase *eventBase, const folly::SocketAddress& address)
-      : eventBase_(eventBase), address_(address), bytesRead_(0) {
-    ctx_.reset(new SSLContext());
-    ctx_->setOptions(SSL_OP_NO_TICKET);
-    ctx_->ciphers("ALL:!ADH:!LOW:!EXP:!MD5:@STRENGTH");
+  explicit AttachDetachClient(const folly::SocketAddress& address)
+      : address_(address), bytesRead_(0) {}
+
+  Future<bool> getFuture() {
+    return promise_.getFuture();
   }
 
   void connect() {
-    sslSocket_ = AsyncSSLSocket::newSocket(ctx_, eventBase_);
-    sslSocket_->connect(this, address_);
+    // create in one and then move to another
+    auto t1Evb = t1_.getEventBase();
+    t1Evb->runInEventBaseThread([this] {
+      sslSocket_ = t1_.createSocket();
+      // ensure we can detach and reattach the context before connecting
+      for (int i = 0; i < 1000; ++i) {
+        sslSocket_->detachSSLContext();
+        sslSocket_->attachSSLContext(t1_.ctx_);
+      }
+      // detach from t1 and connect in t2
+      detach();
+      auto t2Evb = t2_.getEventBase();
+      t2Evb->runInEventBaseThread([this] {
+        t2_.attach(*sslSocket_);
+        sslSocket_->connect(this, address_);
+      });
+    });
   }
 
   void connectSuccess() noexcept override {
+    auto t2Evb = t2_.getEventBase();
+    EXPECT_TRUE(t2Evb->isInEventBaseThread());
     cerr << "client SSL socket connected" << endl;
-
     for (int i = 0; i < 1000; ++i) {
       sslSocket_->detachSSLContext();
-      sslSocket_->attachSSLContext(ctx_);
+      sslSocket_->attachSSLContext(t2_.ctx_);
     }
 
-    EXPECT_EQ(ctx_->getSSLCtx()->references, 2);
-
-    sslSocket_->write(this, buf_, sizeof(buf_));
-    sslSocket_->setReadCB(this);
-    memset(readbuf_, 'b', sizeof(readbuf_));
-    bytesRead_ = 0;
+    // detach from t2 and then read/write in t1
+    t2Evb->runInEventBaseThread([this] {
+      detach();
+      auto t1Evb = t1_.getEventBase();
+      t1Evb->runInEventBaseThread([this] {
+        t1_.attach(*sslSocket_);
+        sslSocket_->write(this, buf_, sizeof(buf_));
+        sslSocket_->setReadCB(this);
+        memset(readbuf_, 'b', sizeof(readbuf_));
+        bytesRead_ = 0;
+      });
+    });
   }
 
   void connectErr(const AsyncSocketException& ex) noexcept override
@@ -96,14 +154,19 @@ class AttachDetachClient : public AsyncSocket::ConnectCallback,
 
   void readErr(const AsyncSocketException& ex) noexcept override {
     cerr << "client readError: " << ex.what() << endl;
+    promise_.setException(ex);
   }
 
   void readDataAvailable(size_t len) noexcept override {
+    EXPECT_TRUE(t1_.getEventBase()->isInEventBaseThread());
+    EXPECT_EQ(sslSocket_->getEventBase(), t1_.getEventBase());
     cerr << "client read data: " << len << endl;
     bytesRead_ += len;
     if (len == sizeof(buf_)) {
       EXPECT_EQ(memcmp(buf_, readbuf_, bytesRead_), 0);
       sslSocket_->closeNow();
+      sslSocket_.reset();
+      promise_.setValue(true);
     }
   }
 };
@@ -119,13 +182,12 @@ TEST(AsyncSSLSocketTest2, AttachDetachSSLContext) {
   SSLServerAcceptCallbackDelay acceptCallback(&handshakeCallback);
   TestSSLServer server(&acceptCallback);
 
-  EventBase eventBase;
-  EventBaseAborter eba(&eventBase, 3000);
   std::shared_ptr<AttachDetachClient> client(
-    new AttachDetachClient(&eventBase, server.getAddress()));
+      new AttachDetachClient(server.getAddress()));
 
+  auto f = client->getFuture();
   client->connect();
-  eventBase.loop();
+  EXPECT_TRUE(f.within(std::chrono::seconds(3)).get());
 }
 
 }  // folly