original_dst_cluster: add override by filter state

source/common/network/BUILD

@@ -477,6 +477,17 @@ envoy_cc_library(
     ],
 )
 
+envoy_cc_library(
+    name = "filter_state_dst_address_lib",
+    srcs = ["filter_state_dst_address.cc"],
+    hdrs = ["filter_state_dst_address.h"],
+    deps = [
+        "//envoy/network:address_interface",
+        "//envoy/stream_info:filter_state_interface",
+        "//source/common/common:macros",
+    ],
+)
+
 envoy_cc_library(
     name = "upstream_server_name_lib",
     srcs = ["upstream_server_name.cc"],

source/common/network/filter_state_dst_address.cc

@@ -0,0 +1,11 @@
+#include "source/common/network/filter_state_dst_address.h"
+
+namespace Envoy {
+namespace Network {
+
+const std::string& DestinationAddress::key() {
+  CONSTRUCT_ON_FIRST_USE(std::string, "envoy.network.transport_socket.original_dst_address");
+}
+
+} // namespace Network
+} // namespace Envoy

source/common/network/filter_state_dst_address.h

@@ -0,0 +1,25 @@
+#pragma once
+
+#include "envoy/network/address.h"
+#include "envoy/stream_info/filter_state.h"
+
+namespace Envoy {
+namespace Network {
+
+/**
+ * Overrides the destination host address selection for ORIGINAL_DST cluster.
+ */
+class DestinationAddress : public StreamInfo::FilterState::Object {
+public:
+  // Returns the key for looking up in the FilterState.
+  static const std::string& key();
+
+  DestinationAddress(Network::Address::InstanceConstSharedPtr address) : address_(address) {}
+  Network::Address::InstanceConstSharedPtr address() const { return address_; }
+
+private:
+  const Network::Address::InstanceConstSharedPtr address_;
+};
+
+} // namespace Network
+} // namespace Envoy

source/common/upstream/BUILD

@@ -368,6 +368,7 @@ envoy_cc_library(
         "//envoy/upstream:cluster_factory_interface",
         "//source/common/common:empty_string",
         "//source/common/network:address_lib",
+        "//source/common/network:filter_state_dst_address_lib",
         "//source/common/network:utility_lib",
         "@envoy_api//envoy/config/cluster/v3:pkg_cc_proto",
         "@envoy_api//envoy/config/core/v3:pkg_cc_proto",

source/common/upstream/original_dst_cluster.cc

@@ -13,6 +13,7 @@
 
 #include "source/common/http/headers.h"
 #include "source/common/network/address_impl.h"
+#include "source/common/network/filter_state_dst_address.h"
 #include "source/common/network/utility.h"
 #include "source/common/protobuf/protobuf.h"
 #include "source/common/protobuf/utility.h"
@@ -22,8 +23,13 @@ namespace Upstream {
 
 HostConstSharedPtr OriginalDstCluster::LoadBalancer::chooseHost(LoadBalancerContext* context) {
   if (context) {
+    // Check if filter state override is present, if yes use it before headers and local address.
+    Network::Address::InstanceConstSharedPtr dst_host = filterStateOverrideHost(context);
+
     // Check if override host header is present, if yes use it otherwise check local address.
-    Network::Address::InstanceConstSharedPtr dst_host = requestOverrideHost(context);
+    if (dst_host == nullptr) {
+      dst_host = requestOverrideHost(context);
+    }
 
     if (dst_host == nullptr) {
       const Network::Connection* connection = context->downstreamConnection();
@@ -78,6 +84,21 @@ HostConstSharedPtr OriginalDstCluster::LoadBalancer::chooseHost(LoadBalancerCont
   return nullptr;
 }
 
+Network::Address::InstanceConstSharedPtr
+OriginalDstCluster::LoadBalancer::filterStateOverrideHost(LoadBalancerContext* context) {
+  const auto* conn = context->downstreamConnection();
+  if (!conn) {
+    return nullptr;
+  }
+  const auto* dst_address =
+      conn->streamInfo().filterState().getDataReadOnly<Network::DestinationAddress>(
+          Network::DestinationAddress::key());
+  if (!dst_address) {
+    return nullptr;
+  }
+  return dst_address->address();
+}
+
 Network::Address::InstanceConstSharedPtr
 OriginalDstCluster::LoadBalancer::requestOverrideHost(LoadBalancerContext* context) {
   if (!http_header_name_.has_value()) {

source/common/upstream/original_dst_cluster.h

@@ -69,6 +69,7 @@ class OriginalDstCluster : public ClusterImplBase {
       return {};
     }
 
+    Network::Address::InstanceConstSharedPtr filterStateOverrideHost(LoadBalancerContext* context);
     Network::Address::InstanceConstSharedPtr requestOverrideHost(LoadBalancerContext* context);
 
   private:

test/common/upstream/BUILD

@@ -474,6 +474,7 @@ envoy_cc_test(
     deps = [
         ":utility_lib",
         "//source/common/event:dispatcher_lib",
+        "//source/common/network:filter_state_dst_address_lib",
         "//source/common/network:utility_lib",
         "//source/common/upstream:original_dst_cluster_lib",
         "//source/common/upstream:upstream_lib",

test/common/upstream/original_dst_cluster_test.cc

@@ -9,6 +9,7 @@
 #include "envoy/stats/scope.h"
 
 #include "source/common/network/address_impl.h"
+#include "source/common/network/filter_state_dst_address.h"
 #include "source/common/network/utility.h"
 #include "source/common/singleton/manager_impl.h"
 #include "source/common/upstream/original_dst_cluster.h"
@@ -679,6 +680,41 @@ TEST_F(OriginalDstClusterTest, UseHttpHeaderDisabled) {
   EXPECT_EQ(host3, nullptr);
 }
 
+TEST_F(OriginalDstClusterTest, UseFilterState) {
+  std::string yaml = R"EOF(
+    name: name
+    connect_timeout: 1.250s
+    type: ORIGINAL_DST
+    lb_policy: CLUSTER_PROVIDED
+    original_dst_lb_config:
+      use_http_header: true
+  )EOF";
+
+  EXPECT_CALL(initialized_, ready());
+  EXPECT_CALL(*cleanup_timer_, enableTimer(_, _));
+  setupFromYaml(yaml);
+
+  OriginalDstCluster::LoadBalancer lb(cluster_);
+  Event::PostCb post_cb;
+
+  // Filter state takes priority over header override.
+  NiceMock<Network::MockConnection> connection;
+  connection.stream_info_.filterState()->setData(
+      Network::DestinationAddress::key(),
+      std::make_shared<Network::DestinationAddress>(
+          std::make_shared<Network::Address::Ipv4Instance>("10.10.11.11", 6666)),
+      StreamInfo::FilterState::StateType::ReadOnly);
+  TestLoadBalancerContext lb_context1(&connection, Http::Headers::get().EnvoyOriginalDstHost.get(),
+                                      "127.0.0.1:5555");
+
+  EXPECT_CALL(membership_updated_, ready());
+  EXPECT_CALL(dispatcher_, post(_)).WillOnce(SaveArg<0>(&post_cb));
+  HostConstSharedPtr host1 = lb.chooseHost(&lb_context1);
+  post_cb();
+  ASSERT_NE(host1, nullptr);
+  EXPECT_EQ("10.10.11.11:6666", host1->address()->asString());
+}
+
 } // namespace
 } // namespace Upstream
 } // namespace Envoy