PeekingAcceptorHandshakeHelper.h

Go to the documentation of this file.

/*
 * Copyright 2017-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <wangle/acceptor/AcceptorHandshakeManager.h>
#include <wangle/acceptor/SocketPeeker.h>

namespace wangle {

class PeekingAcceptorHandshakeHelper : public AcceptorHandshakeHelper,
                                       public SocketPeeker::Callback {
 public:
  class PeekCallback {
   public:
    explicit PeekCallback(size_t bytesRequired):
      bytesRequired_(bytesRequired) {}
    virtual ~PeekCallback() = default;

    size_t getBytesRequired() const {
      return bytesRequired_;
    }

    virtual AcceptorHandshakeHelper::UniquePtr getHelper(
        const std::vector<uint8_t>& peekedBytes,
        const folly::SocketAddress& clientAddr,
        std::chrono::steady_clock::time_point acceptTime,
        TransportInfo& tinfo) = 0;

   private:
    const size_t bytesRequired_;
  };

  PeekingAcceptorHandshakeHelper(
      const folly::SocketAddress& clientAddr,
      std::chrono::steady_clock::time_point acceptTime,
      TransportInfo& tinfo,
      const std::vector<PeekCallback*>& peekCallbacks,
      size_t numBytes)
      : clientAddr_(clientAddr),
        acceptTime_(acceptTime),
        tinfo_(tinfo),
        peekCallbacks_(peekCallbacks),
        numBytes_(numBytes) {}

  // From AcceptorHandshakeHelper
  void start(
      folly::AsyncSSLSocket::UniquePtr sock,
      AcceptorHandshakeHelper::Callback* callback) noexcept override {
    socket_ = std::move(sock);
    callback_ = callback;
    CHECK_EQ(
        socket_->getSSLState(),
        folly::AsyncSSLSocket::SSLStateEnum::STATE_UNENCRYPTED);
    peeker_.reset(new SocketPeeker(*socket_, this, numBytes_));
    peeker_->start();
  }

  void dropConnection(SSLErrorEnum reason = SSLErrorEnum::NO_ERROR) override {
    CHECK_NE(socket_.get() == nullptr, helper_.get() == nullptr);
    if (socket_) {
      socket_->closeNow();
    } else if (helper_) {
      helper_->dropConnection(reason);
    }
  }

  void peekSuccess(std::vector<uint8_t> peekBytes) noexcept override {
    folly::DelayedDestruction::DestructorGuard dg(this);
    peeker_ = nullptr;

    for (auto& peekCallback : peekCallbacks_) {
      helper_ = peekCallback->getHelper(
          peekBytes, clientAddr_, acceptTime_, tinfo_);
      if (helper_) {
        break;
      }
    }

    if (!helper_) {
      // could not get a helper, report error.
      auto type =
          folly::AsyncSocketException::AsyncSocketExceptionType::CORRUPTED_DATA;
      return peekError(
          folly::AsyncSocketException(type, "Unrecognized protocol"));
    }

    auto callback = callback_;
    callback_ = nullptr;
    helper_->start(std::move(socket_), callback);
    CHECK(!socket_);
  }

  void peekError(const folly::AsyncSocketException& ex) noexcept override {
    peeker_ = nullptr;
    auto callback = callback_;
    callback_ = nullptr;
    callback->connectionError(
        socket_.get(), folly::exception_wrapper(ex), folly::none);
  }

 private:
  ~PeekingAcceptorHandshakeHelper() override = default;

  folly::AsyncSSLSocket::UniquePtr socket_;
  AcceptorHandshakeHelper::UniquePtr helper_;
  SocketPeeker::UniquePtr peeker_;

  AcceptorHandshakeHelper::Callback* callback_;
  const folly::SocketAddress& clientAddr_;
  std::chrono::steady_clock::time_point acceptTime_;
  TransportInfo& tinfo_;
  const std::vector<PeekCallback*>& peekCallbacks_;
  size_t numBytes_;
};

using PeekingCallbackPtr = PeekingAcceptorHandshakeHelper::PeekCallback*;

class PeekingAcceptorHandshakeManager : public AcceptorHandshakeManager {
 public:
  PeekingAcceptorHandshakeManager(
        Acceptor* acceptor,
        const folly::SocketAddress& clientAddr,
        std::chrono::steady_clock::time_point acceptTime,
        TransportInfo tinfo,
        const std::vector<PeekingCallbackPtr>& peekCallbacks,
        size_t numBytes):
      AcceptorHandshakeManager(
          acceptor,
          clientAddr,
          acceptTime,
          std::move(tinfo)),
      peekCallbacks_(peekCallbacks),
      numBytes_(numBytes) {}

 protected:
  void startHelper(folly::AsyncSSLSocket::UniquePtr sock) override {
    helper_.reset(new PeekingAcceptorHandshakeHelper(
        clientAddr_,
        acceptTime_,
        tinfo_,
        peekCallbacks_,
        numBytes_));
    helper_->start(std::move(sock), this);
  }

  const std::vector<PeekingCallbackPtr>& peekCallbacks_;
  size_t numBytes_;
};

}