folly/io/async/test/AsyncSSLSocketTest2.cpp

   1 /*
   2  * Copyright 2016 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #include <folly/io/async/test/AsyncSSLSocketTest.h>
  17
  18 #include <gtest/gtest.h>
  19 #include <pthread.h>
  20
  21 #include <folly/io/async/AsyncSSLSocket.h>
  22 #include <folly/io/async/EventBase.h>
  23 #include <folly/io/async/SSLContext.h>
  24
  25 using std::string;
  26 using std::vector;
  27 using std::min;
  28 using std::cerr;
  29 using std::endl;
  30 using std::list;
  31
  32 namespace folly {
  33
  34 class AttachDetachClient : public AsyncSocket::ConnectCallback,
  35                            public AsyncTransportWrapper::WriteCallback,
  36                            public AsyncTransportWrapper::ReadCallback {
  37  private:
  38   EventBase *eventBase_;
  39   std::shared_ptr<AsyncSSLSocket> sslSocket_;
  40   std::shared_ptr<SSLContext> ctx_;
  41   folly::SocketAddress address_;
  42   char buf_[128];
  43   char readbuf_[128];
  44   uint32_t bytesRead_;
  45  public:
  46   AttachDetachClient(EventBase *eventBase, const folly::SocketAddress& address)
  47       : eventBase_(eventBase), address_(address), bytesRead_(0) {
  48     ctx_.reset(new SSLContext());
  49     ctx_->setOptions(SSL_OP_NO_TICKET);
  50     ctx_->ciphers("ALL:!ADH:!LOW:!EXP:!MD5:@STRENGTH");
  51   }
  52
  53   void connect() {
  54     sslSocket_ = AsyncSSLSocket::newSocket(ctx_, eventBase_);
  55     sslSocket_->connect(this, address_);
  56   }
  57
  58   void connectSuccess() noexcept override {
  59     cerr << "client SSL socket connected" << endl;
  60
  61     for (int i = 0; i < 1000; ++i) {
  62       sslSocket_->detachSSLContext();
  63       sslSocket_->attachSSLContext(ctx_);
  64     }
  65
  66     EXPECT_EQ(ctx_->getSSLCtx()->references, 2);
  67
  68     sslSocket_->write(this, buf_, sizeof(buf_));
  69     sslSocket_->setReadCB(this);
  70     memset(readbuf_, 'b', sizeof(readbuf_));
  71     bytesRead_ = 0;
  72   }
  73
  74   void connectErr(const AsyncSocketException& ex) noexcept override
  75   {
  76     cerr << "AttachDetachClient::connectError: " << ex.what() << endl;
  77     sslSocket_.reset();
  78   }
  79
  80   void writeSuccess() noexcept override {
  81     cerr << "client write success" << endl;
  82   }
  83
  84   void writeErr(size_t /* bytesWritten */,
  85                 const AsyncSocketException& ex) noexcept override {
  86     cerr << "client writeError: " << ex.what() << endl;
  87   }
  88
  89   void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
  90     *bufReturn = readbuf_ + bytesRead_;
  91     *lenReturn = sizeof(readbuf_) - bytesRead_;
  92   }
  93   void readEOF() noexcept override {
  94     cerr << "client readEOF" << endl;
  95   }
  96
  97   void readErr(const AsyncSocketException& ex) noexcept override {
  98     cerr << "client readError: " << ex.what() << endl;
  99   }
 100
 101   void readDataAvailable(size_t len) noexcept override {
 102     cerr << "client read data: " << len << endl;
 103     bytesRead_ += len;
 104     if (len == sizeof(buf_)) {
 105       EXPECT_EQ(memcmp(buf_, readbuf_, bytesRead_), 0);
 106       sslSocket_->closeNow();
 107     }
 108   }
 109 };
 110
 111 /**
 112  * Test passing contexts between threads
 113  */
 114 TEST(AsyncSSLSocketTest2, AttachDetachSSLContext) {
 115   // Start listening on a local port
 116   WriteCallbackBase writeCallback;
 117   ReadCallback readCallback(&writeCallback);
 118   HandshakeCallback handshakeCallback(&readCallback);
 119   SSLServerAcceptCallbackDelay acceptCallback(&handshakeCallback);
 120   TestSSLServer server(&acceptCallback);
 121
 122   EventBase eventBase;
 123   EventBaseAborter eba(&eventBase, 3000);
 124   std::shared_ptr<AttachDetachClient> client(
 125     new AttachDetachClient(&eventBase, server.getAddress()));
 126
 127   client->connect();
 128   eventBase.loop();
 129 }
 130
 131 }  // folly
 132
 133 int main(int argc, char *argv[]) {
 134   signal(SIGPIPE, SIG_IGN);
 135   folly::SSLContext::setSSLLockTypes({
 136       {CRYPTO_LOCK_EVP_PKEY, folly::SSLContext::LOCK_NONE},
 137       {CRYPTO_LOCK_SSL_SESSION, folly::SSLContext::LOCK_SPINLOCK},
 138       {CRYPTO_LOCK_SSL_CTX, folly::SSLContext::LOCK_NONE}});
 139   testing::InitGoogleTest(&argc, argv);
 140   gflags::ParseCommandLineFlags(&argc, &argv, true);
 141   return RUN_ALL_TESTS();
 142 }