Local reply and upstream response rewrites using the FilterManager

source/common/http/filter_manager.cc

@@ -879,8 +879,16 @@ void FilterManager::sendLocalReplyViaFilterChain(
                  absl::string_view& content_type) -> void {
             // TODO(snowp): This &get() business isn't nice, rework LocalReply and others to accept
             // opt refs.
+            ENVOY_STREAM_LOG(
+                trace, "sendLocalReply local_reply_.rewrite code={}, body={}, content_type={}",
+                *this, code, body, content_type);
             local_reply_.rewrite(filter_manager_callbacks_.requestHeaders().ptr(), response_headers,
                                  stream_info_, code, body, content_type);
+            // Note that we did a local reply rewrite so that we don't try to do it again in
+            // encodeHeaders. This isn't very clean but we're trying to support local reply rewrites
+            // as well as upstream rewrites, where the match + rewrite logic makes the most sense in
+            // encodeHeaders/Data.
+            did_rewrite_ = true;
           },
           [this, modify_headers](ResponseHeaderMapPtr&& headers, bool end_stream) -> void {
             filter_manager_callbacks_.setResponseHeaders(std::move(headers));
@@ -1044,9 +1052,45 @@ void FilterManager::encodeHeaders(ActiveStreamEncoderFilter* filter, ResponseHea
     }
   }
 
-  const bool modified_end_stream = (end_stream && continue_data_entry == encoder_filters_.end());
+  // See if the response would be written by local_reply_.
+  if (!did_rewrite_) {
+    do_rewrite_ =
+        local_reply_.match(filter_manager_callbacks_.requestHeaders().ptr(),
+                           *filter_manager_callbacks_.responseHeaders(),
+                           filter_manager_callbacks_.responseTrailers().ptr(), stream_info_);
+  }
+
+  bool modified_end_stream = (end_stream && continue_data_entry == encoder_filters_.end());
+  const bool original_modified_end_stream = modified_end_stream;
+  ENVOY_STREAM_LOG(
+      debug, "FilterManager::encodeHeaders: end_stream={}, modified_end_stream={}, do_rewrite_={}",
+      *this, end_stream, modified_end_stream, do_rewrite_);
+  if (do_rewrite_) {
+    // _This_ actually sets buffered_response_data_ internally.
+    rewriteResponse();
+
+    if (buffered_response_data_->length() > 0) {
+      // If we're going to rewrite the response here, then modified_end_stream can no longer be true
+      // because we have a body now.
+      modified_end_stream = false;
+    }
+  }
+
   state_.non_100_response_headers_encoded_ = true;
   filter_manager_callbacks_.encodeHeaders(headers, modified_end_stream);
+
+  // Encode the rewritten response right away if the original `modified_end_stream` was true.
+  // If it wasn't, then we know a body will be encoded later, and we'll let that function
+  // take care of encoding the rewritten response.
+  if (do_rewrite_ && original_modified_end_stream && buffered_response_data_->length() > 0) {
+    ASSERT(!did_rewrite_);
+    ENVOY_STREAM_LOG(trace,
+                     "FilterManager::encodeData calling filter_manager_callbacks_ with {} bytes "
+                     "from buffered_response_data and modified_end_stream={}",
+                     *this, buffered_response_data_->length(), modified_end_stream);
+    filter_manager_callbacks_.encodeData(*buffered_response_data_, modified_end_stream);
+  }
+
   maybeEndEncode(modified_end_stream);
 
   if (!modified_end_stream) {
@@ -1180,7 +1224,24 @@ void FilterManager::encodeData(ActiveStreamEncoderFilter* filter, Buffer::Instan
   }
 
   const bool modified_end_stream = end_stream && trailers_added_entry == encoder_filters_.end();
-  filter_manager_callbacks_.encodeData(data, modified_end_stream);
+  if (do_rewrite_ && buffered_response_data_->length() > 0) {
+    // If we're doing a rewrite and modified_end_stream=true, encode the buffered_response_data_
+    // that was set by rewriteResponse() earlier in encodeHeaders().
+    if (modified_end_stream) {
+      ENVOY_STREAM_LOG(trace,
+                       "FilterManager::encodeData calling filter_manager_callbacks_ with {} bytes "
+                       "from buffered_response_data and modified_end_stream={}",
+                       *this, buffered_response_data_->length(), modified_end_stream);
+      filter_manager_callbacks_.encodeData(*buffered_response_data_, modified_end_stream);
+    }
+  } else {
+    // We're not rewriting the response, so encode the data as is.
+    ENVOY_STREAM_LOG(trace,
+                     "FilterManager::encodeData calling filter_manager_callbacks_.encodeData with "
+                     "{} bytes and modified_end_stream={}",
+                     *this, data.length(), modified_end_stream);
+    filter_manager_callbacks_.encodeData(data, modified_end_stream);
+  }
   maybeEndEncode(modified_end_stream);
 
   // If trailers were adding during encodeData we need to trigger decodeTrailers in order
@@ -1225,6 +1286,35 @@ void FilterManager::maybeEndEncode(bool end_stream) {
   }
 }
 
+void FilterManager::rewriteResponse() {
+  ASSERT(do_rewrite_);
+  ASSERT(!did_rewrite_);
+
+  auto response_headers = filter_manager_callbacks_.responseHeaders();
+  ASSERT(response_headers.ptr() != nullptr);
+
+  std::string rewritten_body{};
+  absl::string_view rewritten_content_type{};
+  Http::Code rewritten_code{static_cast<Http::Code>(Utility::getResponseStatus(*response_headers))};
+
+  ENVOY_STREAM_LOG(trace, "rewriteResponse: calling local_reply_.rewrite with code={}", *this,
+                   rewritten_code);
+  local_reply_.rewrite(filter_manager_callbacks_.requestHeaders().ptr(), *response_headers,
+                       stream_info_, rewritten_code, rewritten_body, rewritten_content_type);
+  ENVOY_STREAM_LOG(
+      trace, "rewriteResponse: local_reply_.rewrite returned body=\"{}\", content_type={}, code={}",
+      *this, rewritten_body, rewritten_content_type, rewritten_code);
+
+  buffered_response_data_ = std::make_unique<Buffer::OwnedImpl>(rewritten_body);
+  if (!rewritten_body.empty()) {
+    // Since we overwrote the response body, we need to set the content-length too.
+    response_headers->setContentLength(buffered_response_data_->length());
+  }
+  if (!rewritten_content_type.empty()) {
+    response_headers->setContentType(rewritten_content_type);
+  }
+}
+
 bool FilterManager::processNewlyAddedMetadata() {
   if (request_metadata_map_vector_ == nullptr) {
     return false;

source/common/http/filter_manager.h

@@ -909,6 +909,11 @@ class FilterManager : public ScopeTrackedObject,
    */
   void maybeEndEncode(bool end_stream);
 
+  /**
+   * Rewrite the response headers and body using local_reply_.
+   */
+  void rewriteResponse();
+
   void sendLocalReply(bool is_grpc_request, Code code, absl::string_view body,
                       const std::function<void(ResponseHeaderMap& headers)>& modify_headers,
                       const absl::optional<Grpc::Status::GrpcStatus> grpc_status,
@@ -1076,6 +1081,8 @@ class FilterManager : public ScopeTrackedObject,
 
   FilterChainFactory& filter_chain_factory_;
   const LocalReply::LocalReply& local_reply_;
+  bool do_rewrite_{};
+  bool did_rewrite_{};
   OverridableRemoteSocketAddressSetterStreamInfo stream_info_;
   // TODO(snowp): Once FM has been moved to its own file we'll make these private classes of FM,
   // at which point they no longer need to be friends.

source/common/local_reply/local_reply.cc

@@ -11,6 +11,7 @@
 #include "common/formatter/substitution_format_string.h"
 #include "common/formatter/substitution_formatter.h"
 #include "common/http/header_map_impl.h"
+#include "common/http/utility.h"
 #include "common/router/header_parser.h"
 
 namespace Envoy {
@@ -80,10 +81,48 @@ class ResponseMapper {
                        StreamInfo::StreamInfo& stream_info, Http::Code& code, std::string& body,
                        BodyFormatter*& final_formatter) const {
     // If not matched, just bail out.
-    if (!filter_->evaluate(stream_info, request_headers, response_headers, response_trailers)) {
+    if (!match(&request_headers, response_headers, &response_trailers, stream_info)) {
       return false;
     }
 
+    rewrite(request_headers, response_headers, stream_info, code, body, final_formatter);
+    return true;
+  }
+
+  // Decide if a request/response pair matches this mapper.
+  bool match(const Http::RequestHeaderMap* request_headers,
+             const Http::ResponseHeaderMap& response_headers,
+             const Http::ResponseTrailerMap* response_trailers,
+             StreamInfo::StreamInfo& stream_info) const {
+    // Set response code on the stream_info because it's used by the StatusCode filter.
+    // Further, we know that the status header present on the upstream response headers
+    // is the status we want to match on. It may not be the status we send downstream
+    // to the client, though, because we may call `rewrite` later.
+    //
+    // Under normal circumstances we should have a response status by this point, because
+    // either the upstream set it or the router filter set it. If for whatever reason we
+    // don't, skip setting the stream info's response code and just let our evaluation
+    // logic do without it. We can't do much better, and we certainly don't want to throw
+    // an exception and crash here.
+    if (response_headers.Status() != nullptr) {
+      stream_info.setResponseCode(
+          static_cast<uint32_t>(Http::Utility::getResponseStatus(response_headers)));
+    }
+
+    if (request_headers == nullptr) {
+      request_headers = Http::StaticEmptyHeaders::get().request_headers.get();
+    }
+
+    if (response_trailers == nullptr) {
+      response_trailers = Http::StaticEmptyHeaders::get().response_trailers.get();
+    }
+
+    return filter_->evaluate(stream_info, *request_headers, response_headers, *response_trailers);
+  }
+
+  void rewrite(const Http::RequestHeaderMap&, Http::ResponseHeaderMap& response_headers,
+               StreamInfo::StreamInfo& stream_info, Http::Code& code, std::string& body,
+               BodyFormatter*& final_formatter) const {
     if (body_.has_value()) {
       body = body_.value();
     }
@@ -99,7 +138,6 @@ class ResponseMapper {
     if (body_formatter_) {
       final_formatter = body_formatter_.get();
     }
-    return true;
   }
 
 private:
@@ -128,6 +166,18 @@ class LocalReplyImpl : public LocalReply {
     }
   }
 
+  bool match(const Http::RequestHeaderMap* request_headers,
+             const Http::ResponseHeaderMap& response_headers,
+             const Http::ResponseTrailerMap* response_trailers,
+             StreamInfo::StreamInfo& stream_info) const override {
+    for (const auto& mapper : mappers_) {
+      if (mapper->match(request_headers, response_headers, response_trailers, stream_info)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   void rewrite(const Http::RequestHeaderMap* request_headers,
                Http::ResponseHeaderMap& response_headers, StreamInfo::StreamInfo& stream_info,
                Http::Code& code, std::string& body,

source/common/local_reply/local_reply.h

@@ -26,6 +26,11 @@ class LocalReply {
                        Http::ResponseHeaderMap& response_headers,
                        StreamInfo::StreamInfo& stream_info, Http::Code& code, std::string& body,
                        absl::string_view& content_type) const PURE;
+
+  virtual bool match(const Http::RequestHeaderMap* request_headers,
+                     const Http::ResponseHeaderMap& response_headers,
+                     const Http::ResponseTrailerMap* response_trailers,
+                     StreamInfo::StreamInfo& stream_info) const PURE;
 };
 
 using LocalReplyPtr = std::unique_ptr<LocalReply>;

test/integration/local_reply_integration_test.cc

@@ -466,4 +466,70 @@ TEST_P(LocalReplyIntegrationTest, ShouldFormatResponseToEmptyBody) {
   EXPECT_EQ(response->body(), "");
 }
 
+// Should match and rewrite an upstream response that does not contain a body.
+TEST_P(LocalReplyIntegrationTest, LocalyReplyMatchUpstreamStatusHeadersOnly) {
+  const std::string yaml = R"EOF(
+mappers:
+- filter:
+    status_code_filter:
+      comparison:
+        op: EQ
+        value:
+          default_value: 429
+          runtime_key: key_b
+  status_code: 450
+body_format:
+  text_format_source:
+    inline_string: "%RESPONSE_CODE%: %RESPONSE_CODE_DETAILS%"
+)EOF";
+  setLocalReplyConfig(yaml);
+  initialize();
+
+  codec_client_ = makeHttpConnection(lookupPort("http"));
+
+  auto response = codec_client_->makeHeaderOnlyRequest(Http::TestRequestHeaderMapImpl{
+      {":method", "GET"}, {":path", "/"}, {":scheme", "http"}, {":authority", "host"}});
+  waitForNextUpstreamRequest();
+  upstream_request_->encodeHeaders(Http::TestResponseHeaderMapImpl{{":status", "429"}}, true);
+  response->waitForHeaders();
+
+  EXPECT_EQ("450", response->headers().Status()->value().getStringView());
+  EXPECT_EQ(response->body(), "450: via_upstream");
+
+  cleanupUpstreamAndDownstream();
+}
+
+TEST_P(LocalReplyIntegrationTest, LocalyReplyMatchUpstreamStatusWithBody) {
+  const std::string yaml = R"EOF(
+mappers:
+- filter:
+    status_code_filter:
+      comparison:
+        op: EQ
+        value:
+          default_value: 429
+          runtime_key: key_b
+  status_code: 451
+body_format:
+  text_format_source:
+    inline_string: "%RESPONSE_CODE%: %RESPONSE_CODE_DETAILS%"
+)EOF";
+  setLocalReplyConfig(yaml);
+  initialize();
+
+  codec_client_ = makeHttpConnection(lookupPort("http"));
+
+  auto response = codec_client_->makeHeaderOnlyRequest(Http::TestRequestHeaderMapImpl{
+      {":method", "GET"}, {":path", "/"}, {":scheme", "http"}, {":authority", "host"}});
+  waitForNextUpstreamRequest();
+  upstream_request_->encodeHeaders(Http::TestResponseHeaderMapImpl{{":status", "429"}}, false);
+  upstream_request_->encodeData(512, true);
+  response->waitForHeaders();
+
+  EXPECT_EQ("451", response->headers().Status()->value().getStringView());
+  EXPECT_EQ(response->body(), "451: via_upstream");
+
+  cleanupUpstreamAndDownstream();
+}
+
 } // namespace Envoy

test/mocks/local_reply/mocks.h

@@ -14,6 +14,13 @@ class MockLocalReply : public LocalReply {
                Http::ResponseHeaderMap& response_headers, StreamInfo::StreamInfo& stream_info,
                Http::Code& code, std::string& body, absl::string_view& content_type),
               (const));
+
+  MOCK_METHOD(bool, match,
+              (const Http::RequestHeaderMap* request_headers,
+               const Http::ResponseHeaderMap& response_headers,
+               const Http::ResponseTrailerMap* response_trailers,
+               StreamInfo::StreamInfo& stream_info),
+              (const));
 };
 } // namespace LocalReply
 } // namespace Envoy