Acceptor.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/ServerSocketConfig.h>
#include <wangle/acceptor/ConnectionCounter.h>
#include <wangle/acceptor/ConnectionManager.h>
#include <wangle/acceptor/LoadShedConfiguration.h>
#include <wangle/acceptor/SecureTransportType.h>
#include <wangle/acceptor/SecurityProtocolContextManager.h>
#include <wangle/acceptor/SSLAcceptorHandshakeHelper.h>
#include <wangle/acceptor/FizzAcceptorHandshakeHelper.h>
#include <wangle/acceptor/TLSPlaintextPeekingCallback.h>

#include <wangle/ssl/SSLCacheProvider.h>
#include <wangle/acceptor/TransportInfo.h>
#include <wangle/ssl/SSLStats.h>

#include <chrono>
#include <event.h>
#include <folly/ExceptionWrapper.h>
#include <folly/io/async/AsyncSSLSocket.h>
#include <folly/io/async/AsyncServerSocket.h>
#include <folly/io/async/AsyncUDPServerSocket.h>

namespace wangle {

class AsyncTransport;
class ManagedConnection;
class SecurityProtocolContextManager;
class SSLContextManager;

class Acceptor :
  public folly::AsyncServerSocket::AcceptCallback,
  public wangle::ConnectionManager::Callback,
  public folly::AsyncUDPServerSocket::Callback  {
 public:

  enum class State : uint32_t {
    kInit,  // not yet started
    kRunning, // processing requests normally
    kDraining, // processing outstanding conns, but not accepting new ones
    kDone,  // no longer accepting, and all connections finished
  };

  explicit Acceptor(const ServerSocketConfig& accConfig);
  ~Acceptor() override;

  virtual void setSSLCacheProvider(
      const std::shared_ptr<SSLCacheProvider>& cacheProvider) {
    cacheProvider_ = cacheProvider;
  }

  virtual void init(folly::AsyncServerSocket* serverSocket,
                    folly::EventBase* eventBase,
                    SSLStats* stats = nullptr);

  virtual void resetSSLContextConfigs();

  void addSSLContextConfig(const SSLContextConfig& sslCtxConfig);

  SSLContextManager* getSSLContextManager() const {
    return sslCtxManager_.get();
  }

  virtual void setTLSTicketSecrets(
      const std::vector<std::string>& oldSecrets,
      const std::vector<std::string>& currentSecrets,
      const std::vector<std::string>& newSecrets);

  uint32_t getNumConnections() const {
    return downstreamConnectionManager_ ?
      (uint32_t)downstreamConnectionManager_->getNumConnections() : 0;
  }

  virtual folly::EventBase* getEventBase() const { return base_; }

  virtual wangle::ConnectionManager* getConnectionManager() {
    return downstreamConnectionManager_.get();
  }

  void addConnection(wangle::ManagedConnection* connection);

  State getState() const {
    return state_;
  }

  std::chrono::milliseconds getConnTimeout() const;

  const std::string& getName() const {
    return accConfig_.name;
  }

  std::chrono::milliseconds getSSLHandshakeTimeout() const {
    return accConfig_.sslHandshakeTimeout;
  }

  void setGracefulShutdownTimeout(std::chrono::milliseconds gracefulShutdown) {
    gracefulShutdownTimeout_ = gracefulShutdown;
  }

  std::chrono::milliseconds getGracefulShutdownTimeout() const {
    return gracefulShutdownTimeout_;
  }

  virtual void forceStop();

  bool isSSL() const { return accConfig_.isSSL(); }

  const ServerSocketConfig& getConfig() const { return accConfig_; }

  static uint64_t getTotalNumPendingSSLConns() {
    return totalNumPendingSSLConns_.load();
  }

  virtual void onDoneAcceptingConnection(
    int fd,
    const folly::SocketAddress& clientAddr,
    std::chrono::steady_clock::time_point acceptTime
  ) noexcept;

  void processEstablishedConnection(
    int fd,
    const folly::SocketAddress& clientAddr,
    std::chrono::steady_clock::time_point acceptTime,
    TransportInfo& tinfo
  ) noexcept;

  virtual void startHandshakeManager(
    folly::AsyncSSLSocket::UniquePtr sslSock,
    Acceptor* acceptor,
    const folly::SocketAddress& clientAddr,
    std::chrono::steady_clock::time_point acceptTime,
    TransportInfo& tinfo) noexcept;

  void drainAllConnections();

  virtual void drainConnections(double pctToDrain);

  void dropAllConnections();

  virtual void dropConnections(double pctToDrop);

   virtual void plaintextConnectionReady(
      folly::AsyncTransportWrapper::UniquePtr sock,
      const folly::SocketAddress& clientAddr,
      const std::string& nextProtocolName,
      SecureTransportType secureTransportType,
      TransportInfo& tinfo);

   void connectionReady(
      folly::AsyncTransportWrapper::UniquePtr sock,
      const folly::SocketAddress& clientAddr,
      const std::string& nextProtocolName,
      SecureTransportType secureTransportType,
      TransportInfo& tinfo);

  virtual void sslConnectionReady(folly::AsyncTransportWrapper::UniquePtr sock,
      const folly::SocketAddress& clientAddr,
      const std::string& nextProtocol,
      SecureTransportType secureTransportType,
      TransportInfo& tinfo);

  virtual void sslConnectionError(const folly::exception_wrapper& ex);

  virtual void updateSSLStats(
      const folly::AsyncTransportWrapper* /*sock*/,
      std::chrono::milliseconds /*acceptLatency*/,
      SSLErrorEnum /*error*/) noexcept {}

 protected:

  folly::EventBase* base_{nullptr};

  virtual uint64_t getConnectionCountForLoadShedding(void) const { return 0; }
  virtual uint64_t getActiveConnectionCountForLoadShedding() const { return 0; }
  virtual uint64_t getWorkerMaxConnections() const {
    return connectionCounter_->getMaxConnections();
  }

  virtual bool canAccept(const folly::SocketAddress&);

  virtual void onNewConnection(
      folly::AsyncTransportWrapper::UniquePtr /*sock*/,
      const folly::SocketAddress* /*address*/,
      const std::string& /*nextProtocolName*/,
      SecureTransportType /*secureTransportType*/,
      const TransportInfo& /*tinfo*/) {}

  void onListenStarted() noexcept override {}
  void onListenStopped() noexcept override {}
  void onDataAvailable(
      std::shared_ptr<folly::AsyncUDPSocket> /*socket*/,
      const folly::SocketAddress&,
      std::unique_ptr<folly::IOBuf>,
      bool) noexcept override {}

  virtual folly::AsyncSocket::UniquePtr makeNewAsyncSocket(
      folly::EventBase* base,
      int fd) {
    return folly::AsyncSocket::UniquePtr(
        new folly::AsyncSocket(base, fd));
  }

  virtual folly::AsyncSSLSocket::UniquePtr makeNewAsyncSSLSocket(
    const std::shared_ptr<folly::SSLContext>& ctx, folly::EventBase* base, int fd) {
    return folly::AsyncSSLSocket::UniquePtr(
        new folly::AsyncSSLSocket(
          ctx,
          base,
          fd,
          true, /* set server */
          true /* defer the security negotiation until sslAccept */));
  }

 protected:

  virtual void onConnectionsDrained() {}

  // AsyncServerSocket::AcceptCallback methods
  void connectionAccepted(
      int fd,
      const folly::SocketAddress& clientAddr) noexcept override;
  void acceptError(const std::exception& ex) noexcept override;
  void acceptStopped() noexcept override;

  // ConnectionManager::Callback methods
  void onEmpty(const wangle::ConnectionManager& cm) override;
  void onConnectionAdded(const ManagedConnection*) override {}
  void onConnectionRemoved(const ManagedConnection*) override {}

 protected:
  const ServerSocketConfig accConfig_;
  void setLoadShedConfig(
    std::shared_ptr<const LoadShedConfiguration> loadShedConfig,
    const IConnectionCounter* counter);

  // Helper function to initialize downstreamConnectionManager_
  virtual void initDownstreamConnectionManager(folly::EventBase* eventBase);
  virtual DefaultToFizzPeekingCallback* getFizzPeeker() {
    return &defaultFizzPeeker_;
  }
  virtual std::shared_ptr<fizz::server::FizzServerContext> createFizzContext();
  virtual std::shared_ptr<fizz::server::TicketCipher>
  createFizzTicketCipher(folly::Optional<std::string> = folly::none);
  void updateFizzContext(fizz::server::FizzServerContext*);

  folly::AsyncSocket::OptionMap socketOptions_;

  std::unique_ptr<SSLContextManager> sslCtxManager_;

  SecurityProtocolContextManager securityProtocolCtxManager_;

  TLSPlaintextPeekingCallback tlsPlaintextPeekingCallback_;
  DefaultToSSLPeekingCallback defaultPeekingCallback_;
  DefaultToFizzPeekingCallback defaultFizzPeeker_;

  wangle::ConnectionManager::UniquePtr downstreamConnectionManager_;

  std::shared_ptr<SSLCacheProvider> cacheProvider_;
  wangle::TLSTicketKeySeeds currentSecrets_;

 private:

  // Forbidden copy constructor and assignment opererator
  Acceptor(Acceptor const &) = delete;
  Acceptor& operator=(Acceptor const &) = delete;

  void checkDrained();

  State state_{State::kInit};
  uint64_t numPendingSSLConns_{0};

  static std::atomic<uint64_t> totalNumPendingSSLConns_;

  bool forceShutdownInProgress_{false};
  std::shared_ptr<const LoadShedConfiguration> loadShedConfig_{nullptr};
  const IConnectionCounter* connectionCounter_{nullptr};
  std::chrono::milliseconds gracefulShutdownTimeout_{5000};
};

class AcceptorFactory {
 public:
  virtual std::shared_ptr<Acceptor> newAcceptor(folly::EventBase*) = 0;
  virtual ~AcceptorFactory() = default;
};

} // namespace