From: Subodh Iyengar <subodh@fb.com>
Date: Tue, 18 Aug 2015 04:28:08 +0000 (-0700)
Subject: Refactor HandshakeHelper and add a peeking handshake helper
X-Git-Tag: v0.55.0~16
X-Git-Url: http://plrg.eecs.uci.edu/git/?a=commitdiff_plain;h=663d3a7b260b9991e519310fb304c71edcff2498;p=folly.git

Refactor HandshakeHelper and add a peeking handshake helper

Summary: This adds support to Acceptor to be able to switch between
multiple protcols when SSL is being negotiated using MSG_PEEK.

The motivation for this is to be able to try out multiple protocols.

Reviewed By: @djwatson

Differential Revision: D2327946
---

diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index 9df09eb2..e7c9c451 100644
--- a/folly/io/async/AsyncSocket.cpp
+++ b/folly/io/async/AsyncSocket.cpp
@@ -1235,7 +1235,12 @@ ssize_t AsyncSocket::performRead(void** buf, size_t* buflen, size_t* offset) {
   VLOG(5) << "AsyncSocket::performRead() this=" << this
           << ", buf=" << *buf << ", buflen=" << *buflen;
 
-  ssize_t bytes = recv(fd_, *buf, *buflen, MSG_DONTWAIT);
+  int recvFlags = 0;
+  if (peek_) {
+    recvFlags |= MSG_PEEK;
+  }
+
+  ssize_t bytes = recv(fd_, *buf, *buflen, MSG_DONTWAIT | recvFlags);
   if (bytes < 0) {
     if (errno == EAGAIN || errno == EWOULDBLOCK) {
       // No more data to read right now.
diff --git a/folly/io/async/AsyncSocket.h b/folly/io/async/AsyncSocket.h
index 38939ee7..17e87796 100644
--- a/folly/io/async/AsyncSocket.h
+++ b/folly/io/async/AsyncSocket.h
@@ -475,6 +475,10 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     return setsockopt(fd_, level, optname, optval, sizeof(T));
   }
 
+  virtual void setPeek(bool peek) {
+    peek_ = peek;
+  }
+
   enum class StateEnum : uint8_t {
     UNINIT,
     CONNECTING,
@@ -764,6 +768,8 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   size_t appBytesReceived_;             ///< Num of bytes received from socket
   size_t appBytesWritten_;              ///< Num of bytes written to socket
   bool isBufferMovable_{false};
+
+  bool peek_{false}; // Peek bytes.
 };
 
 
diff --git a/folly/io/async/test/MockAsyncSSLSocket.h b/folly/io/async/test/MockAsyncSSLSocket.h
index 47df2f77..ff9456c6 100644
--- a/folly/io/async/test/MockAsyncSSLSocket.h
+++ b/folly/io/async/test/MockAsyncSSLSocket.h
@@ -24,8 +24,9 @@ class MockAsyncSSLSocket : public AsyncSSLSocket {
  public:
   MockAsyncSSLSocket(
    const std::shared_ptr<SSLContext>& ctx,
-   EventBase* base) :
-    AsyncSSLSocket(ctx, base) {
+   EventBase* base,
+   bool deferSecurityNegotiation = false) :
+    AsyncSSLSocket(ctx, base, deferSecurityNegotiation) {
   }
 
   GMOCK_METHOD5_(, noexcept, ,
@@ -47,6 +48,8 @@ class MockAsyncSSLSocket : public AsyncSSLSocket {
   MOCK_CONST_METHOD2(
    getSelectedNextProtocolNoThrow,
    bool(const unsigned char**, unsigned*));
+  MOCK_METHOD1(setPeek, void(bool));
+  MOCK_METHOD1(setReadCB, void(ReadCallback*));
 
   void sslConn(
     AsyncSSLSocket::HandshakeCB* cb,
@@ -63,10 +66,32 @@ class MockAsyncSSLSocket : public AsyncSSLSocket {
 
     sslConnectMockable(cb, timeout, verify);
   }
+
+  void sslAccept(
+    AsyncSSLSocket::HandshakeCB* cb,
+    uint32_t timeout,
+    const SSLContext::SSLVerifyPeerEnum& verify)
+      override {
+    if (timeout > 0) {
+      handshakeTimeout_.scheduleTimeout(timeout);
+    }
+
+    state_ = StateEnum::ESTABLISHED;
+    sslState_ = STATE_ACCEPTING;
+    handshakeCallback_ = cb;
+
+    sslAcceptMockable(cb, timeout, verify);
+  }
+
   MOCK_METHOD3(
    sslConnectMockable,
    void(AsyncSSLSocket::HandshakeCB*, uint64_t,
      const SSLContext::SSLVerifyPeerEnum&));
+
+  MOCK_METHOD3(
+   sslAcceptMockable,
+   void(AsyncSSLSocket::HandshakeCB*, uint32_t,
+     const SSLContext::SSLVerifyPeerEnum&));
 };
 
 }}