From d0d34d8e226197fca5313dfcef58e2ca85ee734d Mon Sep 17 00:00:00 2001
From: Daniel Sommermann <dcsommer@whatsapp.com>
Date: Wed, 25 Jan 2017 13:01:37 -0800
Subject: [PATCH] Add ability to set custom SSLContext on TestSSLServer

Summary:
This is needed if you want to test other scenarios where the
server has other OpenSSL settings.

Differential Revision: D4463587

fbshipit-source-id: ffd4019e921649dee703363b2ff028b4d8063210
---
 folly/io/async/test/TestSSLServer.cpp | 43 +++++++++++++++++++--------
 folly/io/async/test/TestSSLServer.h   | 13 ++++----
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/folly/io/async/test/TestSSLServer.cpp b/folly/io/async/test/TestSSLServer.cpp
index dabb1d53..bc127db4 100644
--- a/folly/io/async/test/TestSSLServer.cpp
+++ b/folly/io/async/test/TestSSLServer.cpp
@@ -21,15 +21,36 @@ const char* kTestCert = "folly/io/async/test/certs/tests-cert.pem";
 const char* kTestKey = "folly/io/async/test/certs/tests-key.pem";
 const char* kTestCA = "folly/io/async/test/certs/ca-cert.pem";
 
+TestSSLServer::~TestSSLServer() {
+  if (thread_.joinable()) {
+    evb_.runInEventBaseThread([&]() { socket_->stopAccepting(); });
+    LOG(INFO) << "Waiting for server thread to exit";
+    thread_.join();
+  }
+}
+
 TestSSLServer::TestSSLServer(SSLServerAcceptCallbackBase* acb, bool enableTFO)
-    : ctx_(new SSLContext),
-      acb_(acb),
-      socket_(AsyncServerSocket::newSocket(&evb_)) {
-  // Set up the SSL context
+    : acb_(acb) {
+  // Set up a default SSL context
+  ctx_ = std::make_shared<SSLContext>();
   ctx_->loadCertificate(kTestCert);
   ctx_->loadPrivateKey(kTestKey);
   ctx_->ciphers("ALL:!ADH:!LOW:!EXP:!MD5:@STRENGTH");
 
+  init(enableTFO);
+}
+
+TestSSLServer::TestSSLServer(
+    SSLServerAcceptCallbackBase* acb,
+    std::shared_ptr<SSLContext> ctx,
+    bool enableTFO)
+    : ctx_(ctx), acb_(acb) {
+  init(enableTFO);
+}
+
+void TestSSLServer::init(bool enableTFO) {
+  socket_ = AsyncServerSocket::newSocket(&evb_);
+
   acb_->ctx_ = ctx_;
   acb_->base_ = &evb_;
 
@@ -46,15 +67,11 @@ TestSSLServer::TestSSLServer(SSLServerAcceptCallbackBase* acb, bool enableTFO)
   socket_->addAcceptCallback(acb_, &evb_);
   socket_->startAccepting();
 
-  thread_ = std::thread([&] { Main(); });
+  thread_ = std::thread([&] {
+    evb_.loop();
+    acb_->detach();
+    LOG(INFO) << "Server thread exited event loop";
+  });
   LOG(INFO) << "Accepting connections on " << address_;
 }
-
-TestSSLServer::~TestSSLServer() {
-  if (thread_.joinable()) {
-    evb_.runInEventBaseThread([&]() { socket_->stopAccepting(); });
-    LOG(INFO) << "Waiting for server thread to exit";
-    thread_.join();
-  }
-}
 }
diff --git a/folly/io/async/test/TestSSLServer.h b/folly/io/async/test/TestSSLServer.h
index a710d37d..cc83a4b2 100644
--- a/folly/io/async/test/TestSSLServer.h
+++ b/folly/io/async/test/TestSSLServer.h
@@ -99,6 +99,10 @@ class TestSSLServer {
   explicit TestSSLServer(
       SSLServerAcceptCallbackBase* acb,
       bool enableTFO = false);
+  explicit TestSSLServer(
+      SSLServerAcceptCallbackBase* acb,
+      std::shared_ptr<SSLContext> ctx,
+      bool enableTFO = false);
 
   // Kills the thread.
   virtual ~TestSSLServer();
@@ -112,17 +116,14 @@ class TestSSLServer {
   }
 
  protected:
-  void Main() {
-    evb_.loop();
-    acb_->detach();
-    LOG(INFO) << "Server thread exited event loop";
-  }
-
   EventBase evb_;
   std::shared_ptr<SSLContext> ctx_;
   SSLServerAcceptCallbackBase* acb_;
   std::shared_ptr<AsyncServerSocket> socket_;
   SocketAddress address_;
   std::thread thread_;
+
+ private:
+  void init(bool);
 };
 }
-- 
2.34.1