[ProtoApiScrubber] Fixing header and trailer propagation in field checker (#42293)

Signed-off-by: Sumit Kumar <sumitkmr@google.com>
Signed-off-by: sumitkmr2 <sumitkmr@google.com>

Author: sumitkmr2

source/extensions/filters/http/proto_api_scrubber/filter.cc

@@ -411,7 +411,9 @@ ProtoApiScrubberFilter::createRequestProtoScrubber() {
   RETURN_IF_NOT_OK(request_type_or_status.status());
 
   request_match_tree_field_checker_ = std::make_unique<FieldChecker>(
-      ScrubberContext::kRequestScrubbing, &decoder_callbacks_->streamInfo(), method_name_,
+      ScrubberContext::kRequestScrubbing, &decoder_callbacks_->streamInfo(),
+      decoder_callbacks_->requestHeaders(), OptRef<const Http::ResponseHeaderMap>{},
+      decoder_callbacks_->requestTrailers(), OptRef<const Http::ResponseTrailerMap>{}, method_name_,
       &filter_config_);
 
   return std::make_unique<ProtoScrubber>(
@@ -427,7 +429,9 @@ ProtoApiScrubberFilter::createResponseProtoScrubber() {
   RETURN_IF_NOT_OK(response_type_or_status.status());
 
   response_match_tree_field_checker_ = std::make_unique<FieldChecker>(
-      ScrubberContext::kResponseScrubbing, &encoder_callbacks_->streamInfo(), method_name_,
+      ScrubberContext::kResponseScrubbing, &encoder_callbacks_->streamInfo(),
+      decoder_callbacks_->requestHeaders(), encoder_callbacks_->responseHeaders(),
+      decoder_callbacks_->requestTrailers(), encoder_callbacks_->responseTrailers(), method_name_,
       &filter_config_);
 
   return std::make_unique<ProtoScrubber>(

source/extensions/filters/http/proto_api_scrubber/scrubbing_util_lib/field_checker.h

@@ -32,10 +32,27 @@ using proto_processing_lib::proto_scrubber::ScrubberContext;
 class FieldChecker : public FieldCheckerInterface, public Logger::Loggable<Logger::Id::filter> {
 public:
   FieldChecker(const ScrubberContext scrubber_context,
-               const Envoy::StreamInfo::StreamInfo* stream_info, const std::string& method_name,
-               const ProtoApiScrubberFilterConfig* filter_config)
+               const Envoy::StreamInfo::StreamInfo* stream_info,
+               OptRef<const Http::RequestHeaderMap> request_headers,
+               OptRef<const Http::ResponseHeaderMap> response_headers,
+               OptRef<const Http::RequestTrailerMap> request_trailers,
+               OptRef<const Http::ResponseTrailerMap> response_trailers,
+               const std::string& method_name, const ProtoApiScrubberFilterConfig* filter_config)
       : scrubber_context_(scrubber_context), matching_data_(*stream_info),
-        method_name_(method_name), filter_config_ptr_(filter_config) {}
+        method_name_(method_name), filter_config_ptr_(filter_config) {
+    if (request_headers.has_value()) {
+      matching_data_.onRequestHeaders(request_headers.ref());
+    }
+    if (response_headers.has_value()) {
+      matching_data_.onResponseHeaders(response_headers.ref());
+    }
+    if (request_trailers.has_value()) {
+      matching_data_.onRequestTrailers(request_trailers.ref());
+    }
+    if (response_trailers.has_value()) {
+      matching_data_.onResponseTrailers(response_trailers.ref());
+    }
+  }
 
   // This type is neither copyable nor movable.
   FieldChecker(const FieldChecker&) = delete;

test/extensions/filters/http/proto_api_scrubber/BUILD

@@ -77,14 +77,20 @@ envoy_cc_test(
     rbe_pool = "6gig",
     deps = [
         "//source/common/grpc:common_lib",
+        "//source/common/http/matching:inputs_lib",
+        "//source/common/router:string_accessor_lib",
         "//source/extensions/filters/http/proto_api_scrubber:config",
         "//source/extensions/matching/http/cel_input:cel_input_lib",
         "//source/extensions/matching/input_matchers/cel_matcher:config",
+        "//source/extensions/matching/network/common:inputs_lib",
         "//test/extensions/filters/http/grpc_field_extraction/message_converter:message_converter_test_lib",
         "//test/integration:http_protocol_integration_lib",
         "//test/proto:apikeys_proto_cc_proto",
+        "//test/test_common:registry_lib",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_cel_cpp//parser",
         "@envoy_api//envoy/extensions/filters/http/proto_api_scrubber/v3:pkg_cc_proto",
+        "@envoy_api//envoy/extensions/matching/common_inputs/network/v3:pkg_cc_proto",
+        "@envoy_api//envoy/type/matcher/v3:pkg_cc_proto",
     ],
 )

test/extensions/filters/http/proto_api_scrubber/integration_test.cc

@@ -1,11 +1,15 @@
 #include "envoy/extensions/filters/http/proto_api_scrubber/v3/config.pb.h"
+#include "envoy/extensions/matching/common_inputs/network/v3/network_inputs.pb.h"
 #include "envoy/grpc/status.h"
+#include "envoy/type/matcher/v3/http_inputs.pb.h"
 
+#include "source/common/router/string_accessor_impl.h"
 #include "source/extensions/filters/http/proto_api_scrubber/scrubbing_util_lib/field_checker.h"
 
 #include "test/extensions/filters/http/grpc_field_extraction/message_converter/message_converter_test_lib.h"
 #include "test/integration/http_protocol_integration.h"
 #include "test/proto/apikeys.pb.h"
+#include "test/test_common/registry.h"
 
 #include "cel/expr/syntax.pb.h"
 #include "fmt/format.h"
@@ -29,17 +33,63 @@ std::string apikeysDescriptorPath() {
 }
 
 const std::string kCreateApiKeyMethod = "/apikeys.ApiKeys/CreateApiKey";
+const std::string kFilterStateLabelKey = "filter_state_label_key";
+const std::string kFilterStateLabelValue = "LABEL1,LABEL2,LABEL3";
+
+// This filter injects data into filter_state.
+class MetadataInjectorFilter : public ::Envoy::Http::PassThroughDecoderFilter {
+public:
+  ::Envoy::Http::FilterHeadersStatus decodeHeaders(::Envoy::Http::RequestHeaderMap&,
+                                                   bool) override {
+    const std::string key = kFilterStateLabelKey;
+    const std::string value = kFilterStateLabelValue;
+    decoder_callbacks_->streamInfo().filterState()->setData(
+        key, std::make_shared<::Envoy::Router::StringAccessorImpl>(value),
+        ::Envoy::StreamInfo::FilterState::StateType::ReadOnly);
+    return ::Envoy::Http::FilterHeadersStatus::Continue;
+  }
+};
+
+class MetadataInjectorConfigFactory
+    : public ::Envoy::Server::Configuration::NamedHttpFilterConfigFactory {
+public:
+  absl::StatusOr<::Envoy::Http::FilterFactoryCb>
+  createFilterFactoryFromProto(const ::Envoy::Protobuf::Message&, const std::string&,
+                               ::Envoy::Server::Configuration::FactoryContext&) override {
+    return [](::Envoy::Http::FilterChainFactoryCallbacks& callbacks) {
+      callbacks.addStreamDecoderFilter(std::make_shared<MetadataInjectorFilter>());
+    };
+  }
+
+  ::Envoy::ProtobufTypes::MessagePtr createEmptyConfigProto() override {
+    return std::make_unique<Protobuf::Struct>();
+  }
+
+  std::string name() const override { return "test_injector"; }
+};
+
+// RegisterFactory handles the singleton lifetime automatically and avoids the destruction crash.
+static ::Envoy::Registry::RegisterFactory<
+    MetadataInjectorConfigFactory, ::Envoy::Server::Configuration::NamedHttpFilterConfigFactory>
+    register_test_injector;
 
 class ProtoApiScrubberIntegrationTest : public HttpProtocolIntegrationTest {
 public:
   void SetUp() override { HttpProtocolIntegrationTest::SetUp(); }
 
   void TearDown() override {
+    if (codec_client_) {
+      // Close the client FIRST to prevent any "connection reset" callbacks.
+      codec_client_->close();
+      codec_client_.reset();
+    }
+
     cleanupUpstreamAndDownstream();
     HttpProtocolIntegrationTest::TearDown();
   }
 
   enum class RestrictionType { Request, Response };
+  enum class StringMatchType { Exact, Regex };
 
   static xds::type::matcher::v3::Matcher::MatcherList::Predicate
   buildCelPredicate(absl::string_view cel_expression) {
@@ -64,6 +114,30 @@ class ProtoApiScrubberIntegrationTest : public HttpProtocolIntegrationTest {
     return predicate;
   }
 
+  static xds::type::matcher::v3::Matcher::MatcherList::Predicate
+  buildStringMatcherPredicate(const std::string& input_extension_name,
+                              const Protobuf::Message& input_config,
+                              const std::string& match_pattern, StringMatchType match_type) {
+    xds::type::matcher::v3::Matcher::MatcherList::Predicate predicate;
+    auto* single = predicate.mutable_single_predicate();
+
+    // Configure the Data Input (The source of the string to be matched)
+    single->mutable_input()->set_name(input_extension_name);
+    single->mutable_input()->mutable_typed_config()->PackFrom(input_config);
+
+    // Configure the String Matcher (The logic to apply)
+    auto* string_matcher = single->mutable_value_match();
+    if (match_type == StringMatchType::Regex) {
+      auto* regex = string_matcher->mutable_safe_regex();
+      regex->mutable_google_re2();
+      regex->set_regex(match_pattern);
+    } else {
+      string_matcher->set_exact(match_pattern);
+    }
+
+    return predicate;
+  }
+
   // Helper to build the configuration with a generic predicate.
   std::string
   getFilterConfig(const std::string& descriptor_path, const std::string& method_name = "",
@@ -133,15 +207,33 @@ class ProtoApiScrubberIntegrationTest : public HttpProtocolIntegrationTest {
   }
 
   template <typename T>
-  IntegrationStreamDecoderPtr sendGrpcRequest(const T& request_msg,
-                                              const std::string& method_path) {
+  IntegrationStreamDecoderPtr
+  sendGrpcRequest(const T& request_msg, const std::string& method_path,
+                  const Http::TestRequestHeaderMapImpl& custom_headers = {}) {
+    // Close the existing connection in case it exists.
+    // This can happen if this method is called more than once from a single test.
+    if (codec_client_ != nullptr) {
+      codec_client_->close();
+    }
+
     codec_client_ = makeHttpConnection(lookupPort("http"));
     auto request_buf = Grpc::Common::serializeToGrpcFrame(request_msg);
+
+    // Default headers
     auto request_headers = Http::TestRequestHeaderMapImpl{{":method", "POST"},
                                                           {":path", method_path},
                                                           {"content-type", "application/grpc"},
                                                           {":authority", "host"},
                                                           {":scheme", "http"}};
+
+    // Merge custom headers (overwriting defaults if keys match)
+    custom_headers.iterate(
+        [&request_headers](const Http::HeaderEntry& header) -> Http::HeaderMap::Iterate {
+          request_headers.setCopy(Http::LowerCaseString(header.key().getStringView()),
+                                  header.value().getStringView());
+          return Http::HeaderMap::Iterate::Continue;
+        });
+
     return codec_client_->makeRequestWithBody(request_headers, request_buf->toString());
   }
 };
@@ -304,6 +396,137 @@ TEST_P(ProtoApiScrubberIntegrationTest, ScrubNestedField_MatcherFalse) {
   ASSERT_TRUE(response->waitForEndStream());
 }
 
+// Tests scrubbing of nested fields in the request when a CEL matcher is configured to use request
+// headers.
+TEST_P(ProtoApiScrubberIntegrationTest, ScrubNestedField_CustomCelMatcher_RequestHeader) {
+  config_helper_.prependFilter(getFilterConfig(
+      apikeysDescriptorPath(), kCreateApiKeyMethod, "key.display_name", RestrictionType::Request,
+      buildCelPredicate("request.headers['api-version'] == '2025-v1'")));
+  initialize();
+
+  {
+    // Tests that the field `key.display_name` is removed from the request as the CEL expression
+    // ("request.headers['api-version'] == '2025_v1'") evaluates to true.
+    auto original_proto = makeCreateApiKeyRequest(R"pb(
+      parent: "public"
+      key { display_name: "should-be-removed" }
+    )pb");
+    auto custom_headers = Http::TestRequestHeaderMapImpl{{"api-version", "2025-v1"}};
+
+    auto response = sendGrpcRequest(original_proto, kCreateApiKeyMethod, custom_headers);
+    waitForNextUpstreamRequest();
+
+    apikeys::CreateApiKeyRequest expected = original_proto;
+    expected.mutable_key()->clear_display_name();
+
+    Buffer::OwnedImpl data;
+    data.add(upstream_request_->body());
+    checkSerializedData<apikeys::CreateApiKeyRequest>(data, {expected});
+
+    upstream_request_->encodeHeaders(Http::TestResponseHeaderMapImpl{{":status", "200"}}, true);
+    ASSERT_TRUE(response->waitForEndStream());
+  }
+
+  {
+    // Tests that the field `key.display_name` is preserved in the request as the CEL expression
+    // ("request.headers['api-version'] == '2025_v1'") evaluates to false.
+    auto original_proto = makeCreateApiKeyRequest(R"pb(
+      parent: "public"
+      key { display_name: "should-stay" }
+    )pb");
+    auto custom_headers = Http::TestRequestHeaderMapImpl{{"api-version", "2025-v2"}};
+
+    auto response = sendGrpcRequest(original_proto, kCreateApiKeyMethod, custom_headers);
+    waitForNextUpstreamRequest();
+
+    Buffer::OwnedImpl data;
+    data.add(upstream_request_->body());
+    checkSerializedData<apikeys::CreateApiKeyRequest>(data, {original_proto});
+
+    upstream_request_->encodeHeaders(Http::TestResponseHeaderMapImpl{{":status", "200"}}, true);
+    ASSERT_TRUE(response->waitForEndStream());
+  }
+}
+
+// Tests scrubbing of nested fields in the request when a String matcher is configured to use
+// request headers via exact match.
+TEST_P(ProtoApiScrubberIntegrationTest, ScrubNestedField_StringMatcher_RequestHeader_ExactMatch) {
+  envoy::type::matcher::v3::HttpRequestHeaderMatchInput header_input;
+  header_input.set_header_name("api-version");
+
+  auto predicate = buildStringMatcherPredicate("envoy.matching.inputs.request_headers",
+                                               header_input, "2025-v1", StringMatchType::Exact);
+
+  config_helper_.prependFilter(getFilterConfig(apikeysDescriptorPath(), kCreateApiKeyMethod,
+                                               "key.display_name", RestrictionType::Request,
+                                               predicate));
+
+  initialize();
+
+  auto original_proto = makeCreateApiKeyRequest(R"pb(
+      parent: "public"
+      key { display_name: "should-be-removed" }
+    )pb");
+  auto custom_headers = Http::TestRequestHeaderMapImpl{{"api-version", "2025-v1"}};
+
+  auto response = sendGrpcRequest(original_proto, kCreateApiKeyMethod, custom_headers);
+  waitForNextUpstreamRequest();
+
+  apikeys::CreateApiKeyRequest expected = original_proto;
+  expected.mutable_key()->clear_display_name();
+
+  Buffer::OwnedImpl data;
+  data.add(upstream_request_->body());
+  checkSerializedData<apikeys::CreateApiKeyRequest>(data, {expected});
+
+  upstream_request_->encodeHeaders(Http::TestResponseHeaderMapImpl{{":status", "200"}}, true);
+  ASSERT_TRUE(response->waitForEndStream());
+}
+
+// Tests scrubbing of nested fields in the request when a String matcher is configured to use filter
+// state via regex match.
+TEST_P(ProtoApiScrubberIntegrationTest, ScrubNestedField_StringMatcher_FilterState_RegexMatch) {
+  envoy::extensions::matching::common_inputs::network::v3::FilterStateInput filter_state_input;
+  filter_state_input.set_key(kFilterStateLabelKey);
+
+  // The metadata injector filter (test_injector) sets the value to: "LABEL1,LABEL2,LABEL3"
+  // We use a regex to verify that "LABEL2" is present in the list.
+  auto predicate =
+      buildStringMatcherPredicate("envoy.matching.inputs.filter_state", filter_state_input,
+                                  ".*LABEL2.*", StringMatchType::Regex);
+
+  config_helper_.prependFilter(getFilterConfig(apikeysDescriptorPath(), kCreateApiKeyMethod,
+                                               "key.display_name", RestrictionType::Request,
+                                               predicate));
+
+  config_helper_.prependFilter(R"yaml(
+    name: test_injector
+    typed_config:
+      "@type": type.googleapis.com/google.protobuf.Struct
+  )yaml");
+
+  initialize();
+
+  auto original_proto = makeCreateApiKeyRequest(R"pb(
+    parent: "public"
+    key { display_name: "should-be-removed" }
+  )pb");
+
+  auto response = sendGrpcRequest(original_proto, kCreateApiKeyMethod);
+  waitForNextUpstreamRequest();
+
+  // Since "LABEL1,LABEL2,LABEL3" matches ".*LABEL2.*", the field should be removed.
+  apikeys::CreateApiKeyRequest expected = original_proto;
+  expected.mutable_key()->clear_display_name();
+
+  Buffer::OwnedImpl data;
+  data.add(upstream_request_->body());
+  checkSerializedData<apikeys::CreateApiKeyRequest>(data, {expected});
+
+  upstream_request_->encodeHeaders(Http::TestResponseHeaderMapImpl{{":status", "200"}}, true);
+  ASSERT_TRUE(response->waitForEndStream());
+}
+
 // ============================================================================
 // TEST GROUP 3: VALIDATION & REJECTION
 // ============================================================================

test/extensions/filters/http/proto_api_scrubber/scrubbing_util_lib/BUILD

@@ -13,6 +13,7 @@ envoy_cc_test(
     srcs = ["field_checker_test.cc"],
     data = ["//test/proto:apikeys_proto_descriptor"],
     deps = [
+        "//envoy/common:optref_lib",
         "//source/common/protobuf",
         "//source/extensions/filters/http/proto_api_scrubber/scrubbing_util_lib",
         "//source/extensions/matching/http/cel_input:cel_input_lib",