#pragma once

#include <memory>

#include "envoy/http/conn_pool.h"

#include "envoy/http/header_map.h"

#include "envoy/network/connection.h"

#include "envoy/router/router_ratelimit.h"

#include "envoy/tcp/conn_pool.h"

#include "envoy/tcp/upstream.h"

#include "envoy/upstream/load_balancer.h"

#include "envoy/upstream/thread_local_cluster.h"

#include "envoy/upstream/upstream.h"

#include "source/common/buffer/buffer_impl.h"

#include "source/common/common/dump_state_utils.h"

#include "source/common/http/codec_client.h"

#include "source/common/http/hash_policy.h"

#include "source/common/http/null_route_impl.h"

#include "source/common/http/response_decoder_impl_base.h"

#include "source/common/network/utility.h"

#include "source/common/router/config_impl.h"

#include "source/common/router/header_parser.h"

#include "source/common/router/router.h"

#include "source/extensions/early_data/default_early_data_policy.h"

namespace Envoy {

namespace TcpProxy {

constexpr absl::string_view DisableTunnelingFilterStateKey = "envoy.tcp_proxy.disable_tunneling";

class TcpConnPool : public GenericConnPool, public Tcp::ConnectionPool::Callbacks {

public:

  TcpConnPool(Upstream::HostConstSharedPtr host, Upstream::ThreadLocalCluster& thread_local_cluster,

              Upstream::LoadBalancerContext* context,

              Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks,

              StreamInfo::StreamInfo& downstream_info);

  ~TcpConnPool() override;

  // GenericConnPool

  void newStream(GenericConnectionPoolCallbacks& callbacks) override;

  // Tcp::ConnectionPool::Callbacks

  void onPoolFailure(ConnectionPool::PoolFailureReason reason,

                     absl::string_view transport_failure_reason,

                     Upstream::HostDescriptionConstSharedPtr host) override;

  void onPoolReady(Tcp::ConnectionPool::ConnectionDataPtr&& conn_data,

                   Upstream::HostDescriptionConstSharedPtr host) override;

private:

  absl::optional<Upstream::TcpPoolData> conn_pool_data_{};

  Tcp::ConnectionPool::Cancellable* upstream_handle_{};

  GenericConnectionPoolCallbacks* callbacks_{};

  Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks_;

  StreamInfo::StreamInfo& downstream_info_;

};

class HttpUpstream;

class CombinedUpstream;

// This class is specific to TCP proxy connection pool and enables TCP proxying mode

// for HTTP upstreams. This is currently only needed for HTTP/1 client codec that half closes

// upstream network connection after encoding end_stream in TCP proxy (i.e. via CONNECT).

class RouterUpstreamRequest : public Router::UpstreamRequest {

public:

  using Router::UpstreamRequest::UpstreamRequest;

  void onPoolReady(std::unique_ptr<Router::GenericUpstream>&& upstream,

                   Upstream::HostDescriptionConstSharedPtr host,

                   const Network::ConnectionInfoProvider& address_provider,

};

class HttpConnPool : public GenericConnPool, public Http::ConnectionPool::Callbacks {

public:

  HttpConnPool(Upstream::HostConstSharedPtr host,

               Upstream::ThreadLocalCluster& thread_local_cluster,

               Upstream::LoadBalancerContext* context, const TunnelingConfigHelper& config,

               Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks,

               Http::StreamDecoderFilterCallbacks&, Http::CodecType type,

               StreamInfo::StreamInfo& downstream_info);

  using RouterUpstreamRequestPtr = std::unique_ptr<RouterUpstreamRequest>;

  ~HttpConnPool() override;

  std::unique_ptr<Router::GenericConnPool> createConnPool(Upstream::HostConstSharedPtr host,

                                                          Upstream::ThreadLocalCluster&,

                                                          Upstream::LoadBalancerContext* context,

                                                          absl::optional<Http::Protocol> protocol);

  // GenericConnPool

  void newStream(GenericConnectionPoolCallbacks& callbacks) override;

  // Http::ConnectionPool::Callbacks,

  void onPoolFailure(ConnectionPool::PoolFailureReason reason,

                     absl::string_view transport_failure_reason,

                     Upstream::HostDescriptionConstSharedPtr host) override;

  void onPoolReady(Http::RequestEncoder& request_encoder,

                   Upstream::HostDescriptionConstSharedPtr host, StreamInfo::StreamInfo& info,

                   absl::optional<Http::Protocol>) override;

  void onUpstreamHostSelected(Upstream::HostDescriptionConstSharedPtr, bool);

  void onHttpPoolReady(Upstream::HostDescriptionConstSharedPtr& host,

                       Ssl::ConnectionInfoConstSharedPtr ssl_info);

  class Callbacks {

  public:

    Callbacks(HttpConnPool& conn_pool, Upstream::HostDescriptionConstSharedPtr host,

              Ssl::ConnectionInfoConstSharedPtr ssl_info)

  protected:

  private:

    HttpConnPool* conn_pool_{};

    Upstream::HostDescriptionConstSharedPtr host_;

    Ssl::ConnectionInfoConstSharedPtr ssl_info_;

  };

private:

  void onGenericPoolReady(Upstream::HostDescriptionConstSharedPtr& host,

                          const Network::ConnectionInfoProvider& address_provider,

                          Ssl::ConnectionInfoConstSharedPtr ssl_info);

  const TunnelingConfigHelper& config_;

  Http::CodecType type_;

  absl::optional<Upstream::HttpPoolData> conn_pool_data_{};

  Http::ConnectionPool::Cancellable* upstream_handle_{};

  GenericConnectionPoolCallbacks* callbacks_{};

  Http::StreamDecoderFilterCallbacks* decoder_filter_callbacks_;

  Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks_;

  std::unique_ptr<HttpUpstream> upstream_;

  std::unique_ptr<CombinedUpstream> combined_upstream_;

  StreamInfo::StreamInfo& downstream_info_;

  std::unique_ptr<Router::GenericConnPool> generic_conn_pool_;

};

class TcpUpstream : public GenericUpstream {

public:

  TcpUpstream(Tcp::ConnectionPool::ConnectionDataPtr&& data,

              Tcp::ConnectionPool::UpstreamCallbacks& callbacks);

  // GenericUpstream

  bool readDisable(bool disable) override;

  void encodeData(Buffer::Instance& data, bool end_stream) override;

  void addBytesSentCallback(Network::Connection::BytesSentCb cb) override;

  Tcp::ConnectionPool::ConnectionData* onDownstreamEvent(Network::ConnectionEvent event,

                                                         absl::string_view details = "") override;

  bool startUpstreamSecureTransport() override;

  Ssl::ConnectionInfoConstSharedPtr getUpstreamConnectionSslInfo() override;

  StreamInfo::DetectedCloseType detectedCloseType() const override;

  absl::string_view localCloseReason() const override;

private:

  Tcp::ConnectionPool::ConnectionDataPtr upstream_conn_data_;

};

class HttpUpstream : public GenericUpstream, protected Http::StreamCallbacks {

public:

  using TunnelingConfig =

      envoy::extensions::filters::network::tcp_proxy::v3::TcpProxy_TunnelingConfig;

  HttpUpstream(Tcp::ConnectionPool::UpstreamCallbacks& callbacks,

               const TunnelingConfigHelper& config, StreamInfo::StreamInfo& downstream_info,

               Http::CodecType type);

  ~HttpUpstream() override;

  bool isValidResponse(const Http::ResponseHeaderMap&);

  void doneReading();

  void doneWriting();

  // GenericUpstream

  bool readDisable(bool disable) override;

  void encodeData(Buffer::Instance& data, bool end_stream) override;

  void addBytesSentCallback(Network::Connection::BytesSentCb cb) override;

  Tcp::ConnectionPool::ConnectionData* onDownstreamEvent(Network::ConnectionEvent event,

                                                         absl::string_view details = "") override;

  // HTTP upstream must not implement converting upstream transport

  // socket from non-secure to secure mode.

  // Http::StreamCallbacks

  void onResetStream(Http::StreamResetReason reason,

                     absl::string_view transport_failure_reason) override;

  void onAboveWriteBufferHighWatermark() override;

  void onBelowWriteBufferLowWatermark() override;

  virtual void setRequestEncoder(Http::RequestEncoder& request_encoder, bool is_ssl);

  StreamInfo::DetectedCloseType detectedCloseType() const override;

protected:

  void resetEncoder(Network::ConnectionEvent event, bool inform_downstream = true);

  // The encoder offered by the upstream http client.

  Http::RequestEncoder* request_encoder_{};

  // The config object that is owned by the downstream network filter chain factory.

  const TunnelingConfigHelper& config_;

  // The downstream info that is owned by the downstream connection.

  StreamInfo::StreamInfo& downstream_info_;

  std::unique_ptr<Http::RequestHeaderMapImpl> downstream_headers_;

private:

  class DecoderShim : public Http::ResponseDecoderImplBase {

  public:

  private:

    HttpUpstream& parent_;

  };

  DecoderShim response_decoder_;

  Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks_;

  const Http::CodecType type_;

  bool read_half_closed_{};

  bool write_half_closed_{};

  // Used to defer onGenericPoolReady and onGenericPoolFailure to the reception

  // of the CONNECT response or the resetEncoder.

  std::unique_ptr<HttpConnPool::Callbacks> conn_pool_callbacks_;

};

class CombinedUpstream : public GenericUpstream, public Envoy::Router::RouterFilterInterface {

public:

  CombinedUpstream(HttpConnPool& http_conn_pool, Tcp::ConnectionPool::UpstreamCallbacks& callbacks,

                   Http::StreamDecoderFilterCallbacks& decoder_callbacks,

                   const TunnelingConfigHelper& config, StreamInfo::StreamInfo& downstream_info);

  using UpstreamRequest = Router::UpstreamRequest;

  void doneReading();

  void doneWriting();

  using UpstreamRequestPtr = std::unique_ptr<UpstreamRequest>;

  void setRouterUpstreamRequest(UpstreamRequestPtr);

  void newStream(GenericConnectionPoolCallbacks& callbacks);

  void encodeData(Buffer::Instance& data, bool end_stream) override;

  Tcp::ConnectionPool::ConnectionData* onDownstreamEvent(Network::ConnectionEvent event,

                                                         absl::string_view details = "") override;

  bool isValidResponse(const Http::ResponseHeaderMap&);

  bool readDisable(bool disable) override;

  void recordUpstreamSslConnection();

  // HTTP upstream must not implement converting upstream transport

  // socket from non-secure to secure mode.

  StreamInfo::DetectedCloseType detectedCloseType() const override;

  // Router::RouterFilterInterface

  void onUpstreamHeaders(uint64_t response_code, Http::ResponseHeaderMapPtr&& headers,

                         UpstreamRequest& upstream_request, bool end_stream) override;

  void onUpstreamData(Buffer::Instance& data, UpstreamRequest& upstream_request,

                      bool end_stream) override;

  void onUpstreamTrailers(Http::ResponseTrailerMapPtr&&, UpstreamRequest&) override;

  void onUpstreamReset(Http::StreamResetReason stream_reset_reason,

                       absl::string_view transport_failure_reason, UpstreamRequest&) override;

  void onUpstreamHostSelected(Upstream::HostDescriptionConstSharedPtr host,

  Http::RequestHeaderMap* downstreamHeaders() override;

protected:

  void onResetEncoder(Network::ConnectionEvent event, bool inform_downstream = true);

  // The config object that is owned by the downstream network filter chain factory.

  const TunnelingConfigHelper& config_;

  // The downstream info that is owned by the downstream connection.

  StreamInfo::StreamInfo& downstream_info_;

  std::unique_ptr<Http::RequestHeaderMapImpl> downstream_headers_;

  HttpConnPool& parent_;

private:

  Http::StreamDecoderFilterCallbacks& decoder_filter_callbacks_;

  class DecoderShim : public Http::ResponseDecoderImplBase {

  public:

    // Http::ResponseDecoder

  private:

    CombinedUpstream& parent_;

  };

  DecoderShim response_decoder_;

  Tcp::ConnectionPool::UpstreamCallbacks& upstream_callbacks_;

  std::unique_ptr<HttpConnPool::Callbacks> conn_pool_callbacks_;

  bool read_half_closed_{};

  bool write_half_closed_{};

  // upstream_request_ has to be destroyed first as they may use CombinedUpstream parent

  // during destruction.

  UpstreamRequestPtr upstream_request_;

  Http::CodecType type_;

};

} // namespace TcpProxy

} // namespace Envoy