UDP Tunneling: Optionally propagate response headers and trailers to downstream info (envoyproxy#30597)

IssaAbuKalbein · web-flow · commit 876753ad28d6 · 2023-11-15T23:55:04.000-05:00
Add support for saving upstream response headers and trailers to downstream info

Risk Level: low
Testing: integration tests
Docs Changes: API

Signed-off-by: Issa Abu Kalbein &lt;iabukalbein@microsoft.com&gt;
diff --git a/api/envoy/extensions/filters/udp/udp_proxy/v3/udp_proxy.proto b/api/envoy/extensions/filters/udp/udp_proxy/v3/udp_proxy.proto
@@ -66,7 +66,7 @@ message UdpProxyConfig {
 
   // Configuration for tunneling UDP over other transports or application layers.
   // Tunneling is currently supported over HTTP/2.
-  // [#next-free-field: 10]
+  // [#next-free-field: 12]
   message UdpTunnelingConfig {
     // Configuration for UDP datagrams buffering.
     message BufferOptions {
@@ -160,6 +160,14 @@ message UdpProxyConfig {
     // while the upstream is not ready will be dropped. In case this field is set but the options
     // are not configured, the default values will be applied as described in the ``BufferOptions``.
     BufferOptions buffer_options = 9;
+
+    // Save the response headers to the downstream info filter state for consumption
+    // by the session filters. The filter state key is ``envoy.udp_proxy.propagate_response_headers``.
+    bool propagate_response_headers = 10;
+
+    // Save the response trailers to the downstream info filter state for consumption
+    // by the session filters. The filter state key is ``envoy.udp_proxy.propagate_response_trailers``.
+    bool propagate_response_trailers = 11;
   }
 
   // The stat prefix used when emitting UDP proxy filter stats.
diff --git a/changelogs/current.yaml b/changelogs/current.yaml
@@ -189,5 +189,12 @@ new_features:
     Ratelimit supports optional additional prefix to use when emitting statistics with :ref:`stat_prefix
     <envoy_v3_api_field_extensions.filters.http.ratelimit.v3.RateLimit.stat_prefix>`
     configuration flag.
+- area: udp_proxy
+  change: |
+    added support for propagating the response headers in :ref:`UdpTunnelingConfig
+    <envoy_v3_api_field_extensions.filters.udp.udp_proxy.v3.UdpProxyConfig.UdpTunnelingConfig.propagate_response_headers>` and
+    response trailers in :ref:`UdpTunnelingConfig
+    <envoy_v3_api_field_extensions.filters.udp.udp_proxy.v3.UdpProxyConfig.UdpTunnelingConfig.propagate_response_trailers>` to
+    the downstream info filter state.
 
 deprecated:
diff --git a/source/extensions/filters/udp/udp_proxy/config.cc b/source/extensions/filters/udp/udp_proxy/config.cc
@@ -11,6 +11,25 @@ constexpr uint32_t DefaultMaxConnectAttempts = 1;
 constexpr uint32_t DefaultMaxBufferedDatagrams = 1024;
 constexpr uint64_t DefaultMaxBufferedBytes = 16384;
 
+ProtobufTypes::MessagePtr TunnelResponseHeadersOrTrailers::serializeAsProto() const {
+  auto proto_out = std::make_unique<envoy::config::core::v3::HeaderMap>();
+  value().iterate([&proto_out](const Http::HeaderEntry& e) -> Http::HeaderMap::Iterate {
+    auto* new_header = proto_out->add_headers();
+    new_header->set_key(std::string(e.key().getStringView()));
+    new_header->set_value(std::string(e.value().getStringView()));
+    return Http::HeaderMap::Iterate::Continue;
+  });
+  return proto_out;
+}
+
+const std::string& TunnelResponseHeaders::key() {
+  CONSTRUCT_ON_FIRST_USE(std::string, "envoy.udp_proxy.propagate_response_headers");
+}
+
+const std::string& TunnelResponseTrailers::key() {
+  CONSTRUCT_ON_FIRST_USE(std::string, "envoy.udp_proxy.propagate_response_trailers");
+}
+
 TunnelingConfigImpl::TunnelingConfigImpl(const TunnelingConfig& config,
                                          Server::Configuration::FactoryContext& context)
     : header_parser_(Envoy::Router::HeaderParser::configure(config.headers_to_add())),
@@ -31,7 +50,9 @@ TunnelingConfigImpl::TunnelingConfigImpl(const TunnelingConfig& config,
                               ? PROTOBUF_GET_WRAPPED_OR_DEFAULT(config.buffer_options(),
                                                                 max_buffered_bytes,
                                                                 DefaultMaxBufferedBytes)
-                              : DefaultMaxBufferedBytes) {
+                              : DefaultMaxBufferedBytes),
+      propagate_response_headers_(config.propagate_response_headers()),
+      propagate_response_trailers_(config.propagate_response_trailers()) {
   if (!post_path_.empty() && !use_post_) {
     throw EnvoyException("Can't set a post path when POST method isn't used");
   }
diff --git a/source/extensions/filters/udp/udp_proxy/config.h b/source/extensions/filters/udp/udp_proxy/config.h
@@ -14,6 +14,43 @@ namespace UdpProxy {
 using TunnelingConfig =
     envoy::extensions::filters::udp::udp_proxy::v3::UdpProxyConfig::UdpTunnelingConfig;
 
+/**
+ * Base class for both tunnel response headers and trailers.
+ */
+class TunnelResponseHeadersOrTrailers : public StreamInfo::FilterState::Object {
+public:
+  ProtobufTypes::MessagePtr serializeAsProto() const override;
+  virtual const Http::HeaderMap& value() const PURE;
+};
+
+/**
+ * Response headers for the tunneling connections.
+ */
+class TunnelResponseHeaders : public TunnelResponseHeadersOrTrailers {
+public:
+  TunnelResponseHeaders(Http::ResponseHeaderMapPtr&& response_headers)
+      : response_headers_(std::move(response_headers)) {}
+  const Http::HeaderMap& value() const override { return *response_headers_; }
+  static const std::string& key();
+
+private:
+  const Http::ResponseHeaderMapPtr response_headers_;
+};
+
+/**
+ * Response trailers for the tunneling connections.
+ */
+class TunnelResponseTrailers : public TunnelResponseHeadersOrTrailers {
+public:
+  TunnelResponseTrailers(Http::ResponseTrailerMapPtr&& response_trailers)
+      : response_trailers_(std::move(response_trailers)) {}
+  const Http::HeaderMap& value() const override { return *response_trailers_; }
+  static const std::string& key();
+
+private:
+  const Http::ResponseTrailerMapPtr response_trailers_;
+};
+
 class TunnelingConfigImpl : public UdpTunnelingConfig {
 public:
   TunnelingConfigImpl(const TunnelingConfig& config,
@@ -37,6 +74,32 @@ class TunnelingConfigImpl : public UdpTunnelingConfig {
   uint32_t maxBufferedDatagrams() const override { return max_buffered_datagrams_; };
   uint64_t maxBufferedBytes() const override { return max_buffered_bytes_; };
 
+  void
+  propagateResponseHeaders(Http::ResponseHeaderMapPtr&& headers,
+                           const StreamInfo::FilterStateSharedPtr& filter_state) const override {
+    if (!propagate_response_headers_) {
+      return;
+    }
+
+    filter_state->setData(TunnelResponseHeaders::key(),
+                          std::make_shared<TunnelResponseHeaders>(std::move(headers)),
+                          StreamInfo::FilterState::StateType::ReadOnly,
+                          StreamInfo::FilterState::LifeSpan::Connection);
+  }
+
+  void
+  propagateResponseTrailers(Http::ResponseTrailerMapPtr&& trailers,
+                            const StreamInfo::FilterStateSharedPtr& filter_state) const override {
+    if (!propagate_response_trailers_) {
+      return;
+    }
+
+    filter_state->setData(TunnelResponseTrailers::key(),
+                          std::make_shared<TunnelResponseTrailers>(std::move(trailers)),
+                          StreamInfo::FilterState::StateType::ReadOnly,
+                          StreamInfo::FilterState::LifeSpan::Connection);
+  }
+
 private:
   std::unique_ptr<Envoy::Router::HeaderParser> header_parser_;
   Formatter::FormatterPtr proxy_host_formatter_;
@@ -49,6 +112,8 @@ class TunnelingConfigImpl : public UdpTunnelingConfig {
   bool buffer_enabled_;
   uint32_t max_buffered_datagrams_;
   uint64_t max_buffered_bytes_;
+  bool propagate_response_headers_;
+  bool propagate_response_trailers_;
 };
 
 class UdpProxyFilterConfigImpl : public UdpProxyFilterConfig,
diff --git a/source/extensions/filters/udp/udp_proxy/udp_proxy_filter.h b/source/extensions/filters/udp/udp_proxy/udp_proxy_filter.h
@@ -102,6 +102,14 @@ class UdpTunnelingConfig {
   virtual bool bufferEnabled() const PURE;
   virtual uint32_t maxBufferedDatagrams() const PURE;
   virtual uint64_t maxBufferedBytes() const PURE;
+
+  virtual void
+  propagateResponseHeaders(Http::ResponseHeaderMapPtr&& headers,
+                           const StreamInfo::FilterStateSharedPtr& filter_state) const PURE;
+
+  virtual void
+  propagateResponseTrailers(Http::ResponseTrailerMapPtr&& trailers,
+                            const StreamInfo::FilterStateSharedPtr& filter_state) const PURE;
 };
 
 using UdpTunnelingConfigPtr = std::unique_ptr<const UdpTunnelingConfig>;
@@ -314,6 +322,9 @@ class HttpUpstreamImpl : public HttpUpstream, protected Http::StreamCallbacks {
         is_valid_response = Http::HeaderUtility::isConnectUdpResponse(*headers);
       }
 
+      parent_.tunnel_config_.propagateResponseHeaders(std::move(headers),
+                                                      parent_.downstream_info_.filterState());
+
       if (!is_valid_response || end_stream) {
         parent_.resetEncoder(Network::ConnectionEvent::LocalClose);
       } else if (parent_.tunnel_creation_callbacks_.has_value()) {
@@ -329,7 +340,9 @@ class HttpUpstreamImpl : public HttpUpstream, protected Http::StreamCallbacks {
       }
     }
 
-    void decodeTrailers(Http::ResponseTrailerMapPtr&&) override {
+    void decodeTrailers(Http::ResponseTrailerMapPtr&& trailers) override {
+      parent_.tunnel_config_.propagateResponseTrailers(std::move(trailers),
+                                                       parent_.downstream_info_.filterState());
       parent_.resetEncoder(Network::ConnectionEvent::LocalClose);
     }
 
diff --git a/test/extensions/filters/udp/udp_proxy/mocks.h b/test/extensions/filters/udp/udp_proxy/mocks.h
@@ -58,6 +58,14 @@ class MockUdpTunnelingConfig : public UdpTunnelingConfig {
   MOCK_METHOD(bool, bufferEnabled, (), (const));
   MOCK_METHOD(uint32_t, maxBufferedDatagrams, (), (const));
   MOCK_METHOD(uint64_t, maxBufferedBytes, (), (const));
+  MOCK_METHOD(void, propagateResponseHeaders,
+              (Http::ResponseHeaderMapPtr && headers,
+               const StreamInfo::FilterStateSharedPtr& filter_state),
+              (const));
+  MOCK_METHOD(void, propagateResponseTrailers,
+              (Http::ResponseTrailerMapPtr && trailers,
+               const StreamInfo::FilterStateSharedPtr& filter_state),
+              (const));
 
   std::string default_proxy_host_ = "default.host.com";
   std::string default_target_host_ = "default.target.host";
diff --git a/test/integration/udp_tunneling_integration_test.cc b/test/integration/udp_tunneling_integration_test.cc
@@ -351,6 +351,8 @@ class UdpTunnelingIntegrationTest : public HttpProtocolIntegrationTest {
     absl::optional<BufferOptions> buffer_options_;
     absl::optional<std::string> idle_timeout_;
     std::string session_access_log_config_ = "";
+    bool propagate_response_headers_ = false;
+    bool propagate_response_trailers_ = false;
   };
 
   void setup(const TestConfig& config) {
@@ -379,9 +381,12 @@ name: udp_proxy
     default_target_port: {}
     retry_options:
       max_connect_attempts: {}
+    propagate_response_headers: {}
+    propagate_response_trailers: {}
 )EOF",
                     config.proxy_host_, config.target_host_, config.default_target_port_,
-                    config.max_connect_attempts_);
+                    config.max_connect_attempts_, config.propagate_response_headers_,
+                    config.propagate_response_trailers_);
 
     if (config.buffer_options_.has_value()) {
       filter_config += fmt::format(R"EOF(
@@ -739,6 +744,141 @@ TEST_P(UdpTunnelingIntegrationTest, ConnectionAttemptRetry) {
   EXPECT_THAT(waitForAccessLog(access_log_filename), testing::HasSubstr("2"));
 }
 
+TEST_P(UdpTunnelingIntegrationTest, PropagateValidResponseHeaders) {
+  const std::string access_log_filename =
+      TestEnvironment::temporaryPath(TestUtility::uniqueFilename());
+
+  const std::string session_access_log_config = fmt::format(R"EOF(
+  access_log:
+  - name: envoy.access_loggers.file
+    typed_config:
+      '@type': type.googleapis.com/envoy.extensions.access_loggers.file.v3.FileAccessLog
+      path: {}
+      log_format:
+        text_format_source:
+          inline_string: "%FILTER_STATE(envoy.udp_proxy.propagate_response_headers:TYPED)%\n"
+)EOF",
+                                                            access_log_filename);
+
+  const TestConfig config{"host.com",
+                          "target.com",
+                          1,
+                          30,
+                          false,
+                          "",
+                          BufferOptions{1, 30},
+                          absl::nullopt,
+                          session_access_log_config,
+                          true,
+                          false};
+  setup(config);
+
+  const std::string datagram = "hello";
+  establishConnection(datagram);
+
+  // Wait for buffered datagram.
+  ASSERT_TRUE(upstream_request_->waitForData(*dispatcher_, expectedCapsules({datagram})));
+
+  sendCapsuleDownstream("response", true);
+  test_server_->waitForGaugeEq("udp.foo.downstream_sess_active", 0);
+
+  // Verify response header value is in the access log.
+  EXPECT_THAT(waitForAccessLog(access_log_filename), testing::HasSubstr("capsule-protocol"));
+}
+
+TEST_P(UdpTunnelingIntegrationTest, PropagateInvalidResponseHeaders) {
+  const std::string access_log_filename =
+      TestEnvironment::temporaryPath(TestUtility::uniqueFilename());
+
+  const std::string session_access_log_config = fmt::format(R"EOF(
+  access_log:
+  - name: envoy.access_loggers.file
+    typed_config:
+      '@type': type.googleapis.com/envoy.extensions.access_loggers.file.v3.FileAccessLog
+      path: {}
+      log_format:
+        text_format_source:
+          inline_string: "%FILTER_STATE(envoy.udp_proxy.propagate_response_headers:TYPED)%\n"
+)EOF",
+                                                            access_log_filename);
+
+  const TestConfig config{"host.com",
+                          "target.com",
+                          1,
+                          30,
+                          false,
+                          "",
+                          BufferOptions{1, 30},
+                          absl::nullopt,
+                          session_access_log_config,
+                          true,
+                          false};
+  setup(config);
+
+  client_->write("hello", *listener_address_);
+  ASSERT_TRUE(fake_upstreams_[0]->waitForHttpConnection(*dispatcher_, fake_upstream_connection_));
+  ASSERT_TRUE(fake_upstream_connection_->waitForNewStream(*dispatcher_, upstream_request_));
+  ASSERT_TRUE(upstream_request_->waitForHeadersComplete());
+  expectRequestHeaders(upstream_request_->headers());
+
+  Http::TestResponseHeaderMapImpl response_headers{{":status", "404"}};
+  upstream_request_->encodeHeaders(response_headers, true);
+
+  test_server_->waitForCounterEq("cluster.cluster_0.upstream_cx_connect_attempts_exceeded", 1);
+  test_server_->waitForCounterEq("cluster.cluster_0.udp.sess_tunnel_failure", 1);
+  test_server_->waitForCounterEq("cluster.cluster_0.udp.sess_tunnel_success", 0);
+  test_server_->waitForGaugeEq("udp.foo.downstream_sess_active", 0);
+
+  // Verify response header value is in the access log.
+  EXPECT_THAT(waitForAccessLog(access_log_filename), testing::HasSubstr("404"));
+}
+
+TEST_P(UdpTunnelingIntegrationTest, PropagateResponseTrailers) {
+  const std::string access_log_filename =
+      TestEnvironment::temporaryPath(TestUtility::uniqueFilename());
+
+  const std::string session_access_log_config = fmt::format(R"EOF(
+  access_log:
+  - name: envoy.access_loggers.file
+    typed_config:
+      '@type': type.googleapis.com/envoy.extensions.access_loggers.file.v3.FileAccessLog
+      path: {}
+      log_format:
+        text_format_source:
+          inline_string: "%FILTER_STATE(envoy.udp_proxy.propagate_response_trailers:TYPED)%\n"
+)EOF",
+                                                            access_log_filename);
+
+  const TestConfig config{"host.com",
+                          "target.com",
+                          1,
+                          30,
+                          false,
+                          "",
+                          BufferOptions{1, 30},
+                          absl::nullopt,
+                          session_access_log_config,
+                          false,
+                          true};
+  setup(config);
+
+  const std::string datagram = "hello";
+  establishConnection(datagram);
+
+  // Wait for buffered datagram.
+  ASSERT_TRUE(upstream_request_->waitForData(*dispatcher_, expectedCapsules({datagram})));
+  sendCapsuleDownstream("response", false);
+
+  const std::string trailer_value = "test-trailer-value";
+  Http::TestResponseTrailerMapImpl response_trailers{{"test-trailer-name", trailer_value}};
+  upstream_request_->encodeTrailers(response_trailers);
+
+  test_server_->waitForGaugeEq("udp.foo.downstream_sess_active", 0);
+
+  // Verify response trailer value is in the access log.
+  EXPECT_THAT(waitForAccessLog(access_log_filename), testing::HasSubstr(trailer_value));
+}
+
 INSTANTIATE_TEST_SUITE_P(IpAndHttpVersions, UdpTunnelingIntegrationTest,
                          testing::ValuesIn(HttpProtocolIntegrationTest::getProtocolTestParams(
                              {Http::CodecType::HTTP2}, {Http::CodecType::HTTP2})),
