#pragma once

#include <chrono>

#include <cstdint>

#include <memory>

#include <string>

#include <vector>

#include "envoy/access_log/access_log.h"

#include "envoy/common/random_generator.h"

#include "envoy/event/timer.h"

#include "envoy/extensions/filters/network/tcp_proxy/v3/tcp_proxy.pb.h"

#include "envoy/http/codec.h"

#include "envoy/http/header_evaluator.h"

#include "envoy/http/request_id_extension.h"

#include "envoy/network/connection.h"

#include "envoy/network/filter.h"

#include "envoy/runtime/runtime.h"

#include "envoy/server/filter_config.h"

#include "envoy/stats/scope.h"

#include "envoy/stats/stats_macros.h"

#include "envoy/stats/timespan.h"

#include "envoy/stream_info/filter_state.h"

#include "envoy/upstream/cluster_manager.h"

#include "envoy/upstream/upstream.h"

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

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

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

#include "source/common/formatter/substitution_format_string.h"

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

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

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

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

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

#include "source/common/stream_info/stream_info_impl.h"

#include "source/common/tcp_proxy/upstream.h"

#include "source/common/upstream/load_balancer_context_base.h"

#include "source/common/upstream/od_cds_api_impl.h"

#include "absl/container/node_hash_map.h"

namespace Envoy {

namespace TcpProxy {

constexpr absl::string_view PerConnectionIdleTimeoutMs =

    "envoy.tcp_proxy.per_connection_idle_timeout_ms";

/**

 * ReceiveBeforeConnectKey is the key for the receive_before_connect filter state. The

 * filter state value is a ``StreamInfo::BoolAccessor`` indicating whether the

 * receive_before_connect functionality should be enabled. Network filters setting this filter

 * state should return `StopIteration` in their `onNewConnection` and `onData` methods until they

 * have read the data they need before the upstream connection establishment, and only then allow

 * the filter chain to proceed to the TCP_PROXY filter.

 */

constexpr absl::string_view ReceiveBeforeConnectKey = "envoy.tcp_proxy.receive_before_connect";

/**

 * All tcp proxy stats. @see stats_macros.h

 */

#define ALL_TCP_PROXY_STATS(COUNTER, GAUGE)                                                        \

/**

 * Tcp proxy stats for on-demand. These stats are generated only if the tcp proxy enables on demand.

 */

#define ON_DEMAND_TCP_PROXY_STATS(COUNTER)                                                         \

/**

 * Struct definition for all tcp proxy stats. @see stats_macros.h

 */

struct TcpProxyStats {

  ALL_TCP_PROXY_STATS(GENERATE_COUNTER_STRUCT, GENERATE_GAUGE_STRUCT)

};

/**

 * Struct definition for on-demand related tcp proxy stats. @see stats_macros.h

 * These stats are available if and only if the tcp proxy enables on-demand.

 * Note that these stats has the same prefix as `TcpProxyStats`.

 */

struct OnDemandStats {

  ON_DEMAND_TCP_PROXY_STATS(GENERATE_COUNTER_STRUCT)

};

class Drainer;

class UpstreamDrainManager;

/**

 * Route is an individual resolved route for a connection.

 */

class Route {

public:

  /**

   * Check whether this route matches a given connection.

   * @param connection supplies the connection to test against.

   * @return bool true if this route matches a given connection.

   */

  virtual bool matches(Network::Connection& connection) const PURE;

  /**

   * @return const std::string& the upstream cluster that owns the route.

   */

  virtual const std::string& clusterName() const PURE;

  /**

   * @return MetadataMatchCriteria* the metadata that a subset load balancer should match when

   * selecting an upstream host

   */

  virtual const Router::MetadataMatchCriteria* metadataMatchCriteria() const PURE;

};

using RouteConstSharedPtr = std::shared_ptr<const Route>;

using TunnelingConfig =

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

/**

 * Response headers for the tunneling connections.

 */

class TunnelResponseHeaders : public Http::TunnelResponseHeadersOrTrailersImpl {

public:

  TunnelResponseHeaders(Http::ResponseHeaderMapPtr&& response_headers)

  static const std::string& key();

private:

  const Http::ResponseHeaderMapPtr response_headers_;

};

/**

 * Response trailers for the tunneling connections.

 */

class TunnelResponseTrailers : public Http::TunnelResponseHeadersOrTrailersImpl {

public:

  TunnelResponseTrailers(Http::ResponseTrailerMapPtr&& response_trailers)

  static const std::string& key();

private:

  const Http::ResponseTrailerMapPtr response_trailers_;

};

class Config;

class TunnelingConfigHelperImpl : public TunnelingConfigHelper,

                                  protected Logger::Loggable<Logger::Id::filter> {

public:

  TunnelingConfigHelperImpl(

      Stats::Scope& scope,

      const envoy::extensions::filters::network::tcp_proxy::v3::TcpProxy& config_message,

      Server::Configuration::FactoryContext& context);

  std::string host(const StreamInfo::StreamInfo& stream_info) const override;

  void

  propagateResponseHeaders(Http::ResponseHeaderMapPtr&& headers,

                           const StreamInfo::FilterStateSharedPtr& filter_state) const override;

  void

  propagateResponseTrailers(Http::ResponseTrailerMapPtr&& trailers,

                            const StreamInfo::FilterStateSharedPtr& filter_state) const override;

private:

  std::unique_ptr<Envoy::Router::HeaderParser> header_parser_;

  Formatter::FormatterPtr hostname_fmt_;

  const bool propagate_response_headers_;

  const bool propagate_response_trailers_;

  std::string post_path_;

  // Request ID extension for tunneling requests. If null, no request ID is generated.

  Envoy::Http::RequestIDExtensionSharedPtr request_id_extension_;

  // Optional overrides for request ID header name and metadata key.

  std::string request_id_header_;

  std::string request_id_metadata_key_;

  Stats::StatNameManagedStorage route_stat_name_storage_;

  const Router::FilterConfig router_config_;

  Server::Configuration::ServerFactoryContext& server_factory_context_;

};

/**

 * On demand configurations.

 */

class OnDemandConfig {

public:

  OnDemandConfig(const envoy::extensions::filters::network::tcp_proxy::v3::TcpProxy_OnDemand&

                     on_demand_message,

                 Server::Configuration::FactoryContext& context, Stats::Scope& scope);

private:

  static OnDemandStats generateStats(Stats::Scope& scope);

  Upstream::OdCdsApiHandlePtr odcds_;

  // The timeout of looking up the on-demand cluster.

  std::chrono::milliseconds lookup_timeout_;

  // On demand stats.

  OnDemandStats stats_;

};

using OnDemandConfigOptConstRef = OptRef<const OnDemandConfig>;

/**

 * Filter configuration.

 *

 * This configuration holds a TLS slot, and therefore it must be destructed

 * on the main thread.

 */

class Config {

public:

  /**

   * Configuration that can be shared and have an arbitrary lifetime safely.

   */

  class SharedConfig {

  public:

    SharedConfig(const envoy::extensions::filters::network::tcp_proxy::v3::TcpProxy& config,

                 Server::Configuration::FactoryContext& context);

    envoy::extensions::filters::network::tcp_proxy::v3::ProxyProtocolTlvMergePolicy

    // Evaluate dynamic TLV formatters and combine with static TLVs.

    Network::ProxyProtocolTLVVector

    evaluateDynamicTLVs(const StreamInfo::StreamInfo& stream_info) const;

  private:

    // Structure to hold TLV formatter information.

    struct TlvFormatter {

      uint8_t type;

      Formatter::FormatterPtr formatter;

    };

    static TcpProxyStats generateStats(Stats::Scope& scope);

    static Network::ProxyProtocolTLVVector

    parseTLVs(absl::Span<const envoy::config::core::v3::TlvEntry* const> tlvs,

              Server::Configuration::GenericFactoryContext& context,

              std::vector<TlvFormatter>& dynamic_tlvs);

    // Hold a Scope for the lifetime of the configuration because connections in

    // the UpstreamDrainManager can live longer than the listener.

    const Stats::ScopeSharedPtr stats_scope_;

    const TcpProxyStats stats_;

    bool flush_access_log_on_connected_ : 1;

    const bool flush_access_log_on_start_ : 1;

    absl::optional<std::chrono::milliseconds> idle_timeout_;

    absl::optional<std::chrono::milliseconds> max_downstream_connection_duration_;

    absl::optional<double> max_downstream_connection_duration_jitter_percentage_;

    absl::optional<std::chrono::milliseconds> access_log_flush_interval_;

    std::unique_ptr<TunnelingConfigHelper> tunneling_config_helper_;

    std::unique_ptr<OnDemandConfig> on_demand_config_;

    BackOffStrategyPtr backoff_strategy_;

    Network::ProxyProtocolTLVVector proxy_protocol_tlvs_;

    std::vector<TlvFormatter> dynamic_tlv_formatters_;

    envoy::extensions::filters::network::tcp_proxy::v3::ProxyProtocolTlvMergePolicy

        proxy_protocol_tlv_merge_policy_{

            envoy::extensions::filters::network::tcp_proxy::v3::ADD_IF_ABSENT};

  };

  using SharedConfigSharedPtr = std::shared_ptr<SharedConfig>;

  Config(const envoy::extensions::filters::network::tcp_proxy::v3::TcpProxy& config,

         Server::Configuration::FactoryContext& context);

  /**

   * Find out which cluster an upstream connection should be opened to based on the

   * parameters of a downstream connection.

   * @param connection supplies the parameters of the downstream connection for

   * which the proxy needs to open the corresponding upstream.

   * @return the route to be used for the upstream connection.

   * If no route applies, returns nullptr.

   */

  RouteConstSharedPtr getRouteFromEntries(Network::Connection& connection);

  RouteConstSharedPtr getRegularRouteFromEntries(Network::Connection& connection);

  const absl::optional<std::chrono::milliseconds>

  calculateMaxDownstreamConnectionDurationWithJitter();

  // Return nullptr if there is no tunneling config.

  UpstreamDrainManager& drainManager();

  // This function must not be called if on demand is disabled.

  // This function must not be called if on demand is disabled.

  envoy::extensions::filters::network::tcp_proxy::v3::UpstreamConnectMode

private:

  struct SimpleRouteImpl : public Route {

    SimpleRouteImpl(const Config& parent, absl::string_view cluster_name);

    // Route

    const Config& parent_;

    std::string cluster_name_;

  };

  class WeightedClusterEntry : public Route {

  public:

    WeightedClusterEntry(const Config& parent,

                         const envoy::extensions::filters::network::tcp_proxy::v3::TcpProxy::

                             WeightedCluster::ClusterWeight& config);

    // Route

  private:

    const Config& parent_;

    const std::string cluster_name_;

    const uint64_t cluster_weight_;

    Router::MetadataMatchCriteriaConstPtr metadata_match_criteria_;

  };

  using WeightedClusterEntryConstSharedPtr = std::shared_ptr<const WeightedClusterEntry>;

  RouteConstSharedPtr default_route_;

  std::vector<WeightedClusterEntryConstSharedPtr> weighted_clusters_;

  uint64_t total_cluster_weight_;

  AccessLog::InstanceSharedPtrVector access_logs_;

  const uint32_t max_connect_attempts_;

  ThreadLocal::SlotPtr upstream_drain_manager_slot_;

  SharedConfigSharedPtr shared_config_;

  std::unique_ptr<const Router::MetadataMatchCriteria> cluster_metadata_match_criteria_;

  Random::RandomGenerator& random_generator_;

  std::unique_ptr<const Network::HashPolicyImpl> hash_policy_;

  Regex::Engine& regex_engine_; // Static lifetime object, safe to store as a reference

  envoy::extensions::filters::network::tcp_proxy::v3::UpstreamConnectMode upstream_connect_mode_{

      envoy::extensions::filters::network::tcp_proxy::v3::IMMEDIATE};

  absl::optional<uint32_t> max_early_data_bytes_;

};

using ConfigSharedPtr = std::shared_ptr<Config>;

/**

 * Per-connection TCP Proxy Cluster configuration.

 */

class PerConnectionCluster : public StreamInfo::FilterState::Object {

public:

  static const std::string& key();

private:

  const std::string cluster_;

};

/**

 * An implementation of a TCP (L3/L4) proxy. This filter will instantiate a new outgoing TCP

 * connection using the defined load balancing proxy for the configured cluster. All data will

 * be proxied back and forth between the two connections.

 */

class Filter : public Network::ReadFilter,

               public Upstream::LoadBalancerContextBase,

               protected Logger::Loggable<Logger::Id::filter>,

               public GenericConnectionPoolCallbacks {

public:

  Filter(ConfigSharedPtr config, Upstream::ClusterManager& cluster_manager);

  ~Filter() override;

  // Network::ReadFilter

  Network::FilterStatus onData(Buffer::Instance& data, bool end_stream) override;

  Network::FilterStatus onNewConnection() override;

  void initializeReadFilterCallbacks(Network::ReadFilterCallbacks& callbacks) override;

  bool startUpstreamSecureTransport() override;

  // GenericConnectionPoolCallbacks

  void onGenericPoolReady(StreamInfo::StreamInfo* info, std::unique_ptr<GenericUpstream>&& upstream,

                          Upstream::HostDescriptionConstSharedPtr& host,

                          const Network::ConnectionInfoProvider& address_provider,

                          Ssl::ConnectionInfoConstSharedPtr ssl_info) override;

  void onGenericPoolFailure(ConnectionPool::PoolFailureReason reason,

                            absl::string_view failure_reason,

                            Upstream::HostDescriptionConstSharedPtr host) override;

  // Upstream::LoadBalancerContext

  const Router::MetadataMatchCriteria* metadataMatchCriteria() override;

  // These two functions allow enabling/disabling reads on the upstream and downstream connections.

  // They are called by the Downstream/Upstream Watermark callbacks to limit buffering.

  void readDisableUpstream(bool disable);

  void readDisableDownstream(bool disable);

  struct UpstreamCallbacks : public Tcp::ConnectionPool::UpstreamCallbacks {

    // Tcp::ConnectionPool::UpstreamCallbacks

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

    void onEvent(Network::ConnectionEvent event) override;

    void onAboveWriteBufferHighWatermark() override;

    void onBelowWriteBufferLowWatermark() override;

    void onBytesSent();

    void onIdleTimeout();

    void drain(Drainer& drainer);

    // Either parent_ or drainer_ will be non-NULL, but never both. This could be

    // logically be represented as a union, but saving one pointer of memory is

    // outweighed by more type safety/better error handling.

    //

    // Parent starts out as non-NULL. If the downstream connection is closed while

    // the upstream connection still has buffered data to flush, drainer_ becomes

    // non-NULL and parent_ is set to NULL.

    Filter* parent_{};

    Drainer* drainer_{};

    bool on_high_watermark_called_{false};

  };

  StreamInfo::StreamInfo& getStreamInfo();

  class HttpStreamDecoderFilterCallbacks : public Http::StreamDecoderFilterCallbacks,

                                           public ScopeTrackedObject {

  public:

    HttpStreamDecoderFilterCallbacks(Filter* parent);

    // Http::StreamDecoderFilterCallbacks

    void sendLocalReply(Http::Code, absl::string_view,

                        std::function<void(Http::ResponseHeaderMap& headers)>,

                        const absl::optional<Grpc::Status::GrpcStatus>,

    absl::optional<Upstream::LoadBalancerContext::OverrideHost>

    // absl::optional<absl::string_view> upstreamOverrideHost() const override {

    //   return absl::nullopt;

    // }

    // ScopeTrackedObject

    Filter* parent_{};

    Http::RequestTrailerMapPtr request_trailer_map_;

    std::shared_ptr<Http::NullRouteImpl> route_;

  };

  Tracing::NullSpan active_span_;

  const Tracing::Config& tracing_config_;

protected:

  struct DownstreamCallbacks : public Network::ConnectionCallbacks {

    // Network::ConnectionCallbacks

    void onAboveWriteBufferHighWatermark() override;

    void onBelowWriteBufferLowWatermark() override;

    Filter& parent_;

    bool on_high_watermark_called_{false};

  };

  enum class UpstreamFailureReason {

    ConnectFailed,

    NoHealthyUpstream,

    ResourceLimitExceeded,

    NoRoute,

  };

  // Callbacks for different error and success states during connection establishment

  virtual void onInitFailure(UpstreamFailureReason reason);

  void initialize(Network::ReadFilterCallbacks& callbacks, bool set_connection_stats);

  // Create connection to the upstream cluster. This function can be repeatedly called on upstream

  // connection failure.

  Network::FilterStatus establishUpstreamConnection();

  // The callback upon on demand cluster discovery response.

  void onClusterDiscoveryCompletion(Upstream::ClusterDiscoveryStatus cluster_status);

  bool maybeTunnel(Upstream::ThreadLocalCluster& cluster);

  void onConnectMaxAttempts();

  void onConnectTimeout();

  void onDownstreamEvent(Network::ConnectionEvent event);

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

  void onUpstreamEvent(Network::ConnectionEvent event);

  void onUpstreamConnection();

  void onIdleTimeout();

  void resetIdleTimer();

  void disableIdleTimer();

  void onMaxDownstreamConnectionDuration();

  void onAccessLogFlushInterval();

  void resetAccessLogFlushTimer();

  void flushAccessLog(AccessLog::AccessLogType access_log_type);

  void disableAccessLogFlushTimer();

  void onRetryTimer();

  void enableRetryTimer();

  void disableRetryTimer();

public:

  // Public for testing purposes

  void onDownstreamTlsHandshakeComplete();

protected:

  const ConfigSharedPtr config_;

  Upstream::ClusterManager& cluster_manager_;

  Network::ReadFilterCallbacks* read_callbacks_{};

  DownstreamCallbacks downstream_callbacks_;

  Event::TimerPtr idle_timer_;

  Event::TimerPtr connection_duration_timer_;

  Event::TimerPtr access_log_flush_timer_;

  Event::TimerPtr retry_timer_;

  // A pointer to the on demand cluster lookup when lookup is in flight.

  Upstream::ClusterDiscoveryCallbackHandlePtr cluster_discovery_handle_;

  std::shared_ptr<UpstreamCallbacks> upstream_callbacks_; // shared_ptr required for passing as a

                                                          // read filter.

  // The upstream handle (either TCP or HTTP). This is set in onGenericPoolReady and should persist

  // until either the upstream or downstream connection is terminated.

  std::unique_ptr<GenericUpstream> upstream_;

  // The connection pool used to set up |upstream_|.

  // This will be non-null from when an upstream connection is attempted until

  // it either succeeds or fails.

  std::unique_ptr<GenericConnPool> generic_conn_pool_;

  // Time the filter first attempted to connect to the upstream after the

  // cluster is discovered. Capture the first time as the filter may try multiple times to connect

  // to the upstream.

  absl::optional<MonotonicTime> initial_upstream_connection_start_time_;

  RouteConstSharedPtr route_;

  Router::MetadataMatchCriteriaConstPtr metadata_match_criteria_;

  Network::TransportSocketOptionsConstSharedPtr transport_socket_options_;

  Network::Socket::OptionsSharedPtr upstream_options_;

  absl::optional<std::chrono::milliseconds> idle_timeout_;

  uint32_t connect_attempts_{};

  bool connecting_{};

  bool downstream_closed_{};

  // Stores the ReceiveBeforeConnect filter state value which can be set by preceding

  // filters in the filter chain. When the filter state is set, TCP_PROXY doesn't disable

  // downstream read during initialization. This feature can hence be used by preceding filters

  // in the filter chain to read data from the downstream connection (for eg: to parse SNI) before

  // the upstream connection is established.

  bool receive_before_connect_{false};

  bool early_data_end_stream_{false};

  Buffer::OwnedImpl early_data_buffer_{};

  HttpStreamDecoderFilterCallbacks upstream_decoder_filter_callbacks_;

  // Connection establishment mode configuration.

  envoy::extensions::filters::network::tcp_proxy::v3::UpstreamConnectMode connect_mode_{

      envoy::extensions::filters::network::tcp_proxy::v3::IMMEDIATE};

  bool waiting_for_tls_handshake_{false};

  bool tls_handshake_complete_{false};

  bool initial_data_received_{false};

  bool read_disabled_due_to_buffer_{false}; // Track if we disabled reading due to buffer overflow.

  uint32_t max_buffered_bytes_{65536};      // Default 64KB.

};

// This class deals with an upstream connection that needs to finish flushing, when the downstream

// connection has been closed. The TcpProxy is destroyed when the downstream connection is closed,

// so handling the upstream connection here allows it to finish draining or timeout.

class Drainer : public Event::DeferredDeletable, protected Logger::Loggable<Logger::Id::filter> {

public:

  Drainer(UpstreamDrainManager& parent, const Config::SharedConfigSharedPtr& config,

          const std::shared_ptr<Filter::UpstreamCallbacks>& callbacks,

          Tcp::ConnectionPool::ConnectionDataPtr&& conn_data, Event::TimerPtr&& idle_timer,

          absl::optional<std::chrono::milliseconds> idle_timeout,

          const Upstream::HostDescriptionConstSharedPtr& upstream_host);

  void onEvent(Network::ConnectionEvent event);

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

  void onIdleTimeout();

  void onBytesSent();

  void cancelDrain();

  Event::Dispatcher& dispatcher();

private:

  UpstreamDrainManager& parent_;

  std::shared_ptr<Filter::UpstreamCallbacks> callbacks_;

  Tcp::ConnectionPool::ConnectionDataPtr upstream_conn_data_;

  Event::TimerPtr idle_timer_;

  absl::optional<std::chrono::milliseconds> idle_timeout_;

  Upstream::HostDescriptionConstSharedPtr upstream_host_;

  Config::SharedConfigSharedPtr config_;

};

using DrainerPtr = std::unique_ptr<Drainer>;

class UpstreamDrainManager : public ThreadLocal::ThreadLocalObject {

public:

  ~UpstreamDrainManager() override;

  void add(const Config::SharedConfigSharedPtr& config,

           Tcp::ConnectionPool::ConnectionDataPtr&& upstream_conn_data,

           const std::shared_ptr<Filter::UpstreamCallbacks>& callbacks,

           Event::TimerPtr&& idle_timer, absl::optional<std::chrono::milliseconds> idle_timeout,

           const Upstream::HostDescriptionConstSharedPtr& upstream_host);

  void remove(Drainer& drainer, Event::Dispatcher& dispatcher);

private:

  // This must be a map instead of set because there is no way to move elements

  // out of a set, and these elements get passed to deferredDelete() instead of

  // being deleted in-place. The key and value will always be equal.

  absl::node_hash_map<Drainer*, DrainerPtr> drainers_;

};

} // namespace TcpProxy

} // namespace Envoy