ratelimit: backport raw_body and dynamic_metadata changes

api/envoy/service/ratelimit/v2/rls.proto

@@ -5,6 +5,8 @@ package envoy.service.ratelimit.v2;
 import "envoy/api/v2/core/base.proto";
 import "envoy/api/v2/ratelimit/ratelimit.proto";
 
+import "google/protobuf/struct.proto";
+
 import "udpa/annotations/migrate.proto";
 import "udpa/annotations/status.proto";
 import "validate/validate.proto";
@@ -46,7 +48,7 @@ message RateLimitRequest {
 }
 
 // A response from a ShouldRateLimit call.
-// [#next-free-field: 6]
+// [#next-free-field: 7]
 message RateLimitResponse {
   enum Code {
     // The response code is not known.
@@ -116,5 +118,13 @@ message RateLimitResponse {
   repeated api.v2.core.HeaderValue request_headers_to_add = 4;
 
   // A response body to send to the downstream client when the response code is not OK.
-  string body = 5;
+  bytes raw_body = 5;
+
+  // Optional response metadata that will be emitted as dynamic metadata to be consumed by the next
+  // filter. This metadata lives in a namespace specified by the canonical name of extension filter
+  // that requires it:
+  //
+  // - :ref:`envoy.filters.http.ratelimit <config_http_filters_ratelimit_dynamic_metadata>` for HTTP filter.
+  // - :ref:`envoy.filters.network.ratelimit <config_network_filters_ratelimit_dynamic_metadata>` for network filter.
+  google.protobuf.Struct dynamic_metadata = 6;
 }

api/envoy/service/ratelimit/v3/rls.proto

@@ -5,6 +5,8 @@ package envoy.service.ratelimit.v3;
 import "envoy/config/core/v3/base.proto";
 import "envoy/extensions/common/ratelimit/v3/ratelimit.proto";
 
+import "google/protobuf/struct.proto";
+
 import "udpa/annotations/status.proto";
 import "udpa/annotations/versioning.proto";
 import "validate/validate.proto";
@@ -67,7 +69,7 @@ message RateLimitRequest {
 }
 
 // A response from a ShouldRateLimit call.
-// [#next-free-field: 6]
+// [#next-free-field: 7]
 message RateLimitResponse {
   option (udpa.annotations.versioning).previous_message_type =
       "envoy.service.ratelimit.v2.RateLimitResponse";
@@ -147,5 +149,13 @@ message RateLimitResponse {
   repeated config.core.v3.HeaderValue request_headers_to_add = 4;
 
   // A response body to send to the downstream client when the response code is not OK.
-  string body = 5;
+  bytes raw_body = 5;
+
+  // Optional response metadata that will be emitted as dynamic metadata to be consumed by the next
+  // filter. This metadata lives in a namespace specified by the canonical name of extension filter
+  // that requires it:
+  //
+  // - :ref:`envoy.filters.http.ratelimit <config_http_filters_ratelimit_dynamic_metadata>` for HTTP filter.
+  // - :ref:`envoy.filters.network.ratelimit <config_network_filters_ratelimit_dynamic_metadata>` for network filter.
+  google.protobuf.Struct dynamic_metadata = 6;
 }

source/extensions/filters/common/ratelimit/ratelimit.h

@@ -30,6 +30,8 @@ enum class LimitStatus {
   OverLimit
 };
 
+using DynamicMetadataPtr = std::unique_ptr<ProtobufWkt::Struct>;
+
 /**
  * Async callbacks used during limit() calls.
  */
@@ -43,7 +45,8 @@ class RequestCallbacks {
    */
   virtual void complete(LimitStatus status, Http::ResponseHeaderMapPtr&& response_headers_to_add,
                         Http::RequestHeaderMapPtr&& request_headers_to_add,
-                        const std::string& response_body) PURE;
+                        const std::string& response_body,
+                        DynamicMetadataPtr&& dynamic_metadata) PURE;
 };
 
 /**

source/extensions/filters/common/ratelimit/ratelimit_impl.cc

@@ -103,15 +103,23 @@ void GrpcClientImpl::onSuccess(
       request_headers_to_add->addCopy(Http::LowerCaseString(h.key()), h.value());
     }
   }
+
+  // TODO(esmet): This dynamic metadata copy is probably unnecessary, but let's just following the
+  // existing pattern of copying parameters over as unique pointers for now.
+  DynamicMetadataPtr dynamic_metadata =
+      response->has_dynamic_metadata()
+          ? std::make_unique<ProtobufWkt::Struct>(response->dynamic_metadata())
+          : nullptr;
   callbacks_->complete(status, std::move(response_headers_to_add),
-                       std::move(request_headers_to_add), response->body());
+                       std::move(request_headers_to_add), response->raw_body(),
+                       std::move(dynamic_metadata));
   callbacks_ = nullptr;
 }
 
 void GrpcClientImpl::onFailure(Grpc::Status::GrpcStatus status, const std::string&,
                                Tracing::Span&) {
   ASSERT(status != Grpc::Status::WellKnownGrpcStatus::Ok);
-  callbacks_->complete(LimitStatus::Error, nullptr, nullptr, "");
+  callbacks_->complete(LimitStatus::Error, nullptr, nullptr, EMPTY_STRING, nullptr);
   callbacks_ = nullptr;
 }
 

source/extensions/filters/http/ratelimit/ratelimit.cc

@@ -12,6 +12,8 @@
 #include "common/http/header_utility.h"
 #include "common/router/config_impl.h"
 
+#include "extensions/filters/http/well_known_names.h"
+
 namespace Envoy {
 namespace Extensions {
 namespace HttpFilters {
@@ -99,7 +101,7 @@ Http::FilterHeadersStatus Filter::encode100ContinueHeaders(Http::ResponseHeaderM
 }
 
 Http::FilterHeadersStatus Filter::encodeHeaders(Http::ResponseHeaderMap& headers, bool) {
-  populateResponseHeaders(headers, false);
+  populateResponseHeaders(headers, /*from_local_reply=*/false);
   return Http::FilterHeadersStatus::Continue;
 }
 
@@ -127,13 +129,19 @@ void Filter::onDestroy() {
 void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
                       Http::ResponseHeaderMapPtr&& response_headers_to_add,
                       Http::RequestHeaderMapPtr&& request_headers_to_add,
-                      const std::string &response_body) {
+                      const std::string& response_body,
+                      Filters::Common::RateLimit::DynamicMetadataPtr&& dynamic_metadata) {
   state_ = State::Complete;
   response_headers_to_add_ = std::move(response_headers_to_add);
   Http::HeaderMapPtr req_headers_to_add = std::move(request_headers_to_add);
   Stats::StatName empty_stat_name;
   Filters::Common::RateLimit::StatNames& stat_names = config_->statNames();
 
+  if (dynamic_metadata != nullptr && !dynamic_metadata->fields().empty()) {
+    callbacks_->streamInfo().setDynamicMetadata(HttpFilterNames::get().RateLimit,
+                                                *dynamic_metadata);
+  }
+
   switch (status) {
   case Filters::Common::RateLimit::LimitStatus::OK:
     cluster_->statsScope().counterFromStatName(stat_names.ok_).inc();
@@ -165,15 +173,10 @@ void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
   if (status == Filters::Common::RateLimit::LimitStatus::OverLimit &&
       config_->runtime().snapshot().featureEnabled("ratelimit.http_filter_enforcing", 100)) {
     state_ = State::Responded;
-    // Always overwrite the content-type automatically set by sendLocalReply whenever
-    // we're sending back a response body here. We do this because any content-type
-    // coming from the ratelimit service should be treated as authoritative, and we
-    // must discard the default text/plain type set by default.
-    const bool overwrite_content_type = response_body.length() > 0;
     callbacks_->sendLocalReply(
         Http::Code::TooManyRequests, response_body,
-        [this, overwrite_content_type](Http::HeaderMap& headers) {
-          populateResponseHeaders(headers, overwrite_content_type);
+        [this](Http::HeaderMap& headers) {
+          populateResponseHeaders(headers, /*from_local_reply=*/true);
         },
         config_->rateLimitedGrpcStatus(), RcDetails::get().RateLimited);
     callbacks_->streamInfo().setResponseFlag(StreamInfo::ResponseFlag::RateLimited);
@@ -214,19 +217,16 @@ void Filter::populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_li
   }
 }
 
-void Filter::populateResponseHeaders(Http::HeaderMap& response_headers,
-                                     bool overwrite_content_type) {
+void Filter::populateResponseHeaders(Http::HeaderMap& response_headers, bool from_local_reply) {
   if (response_headers_to_add_) {
-    // If the ratelimit service is sending back the content-type header and the caller
-    // wants to overwrite the existing content-type, do so here.
+    // If the ratelimit service is sending back the content-type header and we're
+    // populating response headers for a local reply, overwrite the existing
+    // content-type header.
     //
-    // We need to do this when using sendLocalReply because it sets the content-type header to
-    // text/plain whenever the response body is non-empty. This won't be correct for response
-    // bodies that actually have a different content-type, as expressed by the response headers
-    // coming from the ratelimit service. We fix it here to minimize impact surface for other
-    // callers of sendLocalReply that depend on the existing behavior.
-    if (overwrite_content_type &&
-        (*response_headers_to_add_).get(Http::Headers::get().ContentType) != nullptr) {
+    // We do this because sendLocalReply initially sets content-type to text/plain
+    // whenever the response body is non-empty, but we want the content-type coming
+    // from the ratelimit service to be authoritative in this case.
+    if (from_local_reply && !response_headers_to_add_->getContentTypeValue().empty()) {
       response_headers.remove(Http::Headers::get().ContentType);
     }
     Http::HeaderUtility::addHeaders(response_headers, *response_headers_to_add_);

source/extensions/filters/http/ratelimit/ratelimit.h

@@ -119,15 +119,16 @@ class Filter : public Http::StreamFilter, public Filters::Common::RateLimit::Req
   void complete(Filters::Common::RateLimit::LimitStatus status,
                 Http::ResponseHeaderMapPtr&& response_headers_to_add,
                 Http::RequestHeaderMapPtr&& request_headers_to_add,
-                const std::string &response_body) override;
+                const std::string& response_body,
+                Filters::Common::RateLimit::DynamicMetadataPtr&& dynamic_metadata) override;
 
 private:
   void initiateCall(const Http::RequestHeaderMap& headers);
   void populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_limit_policy,
                                     std::vector<Envoy::RateLimit::Descriptor>& descriptors,
                                     const Router::RouteEntry* route_entry,
                                     const Http::HeaderMap& headers) const;
-  void populateResponseHeaders(Http::HeaderMap& response_headers, bool overwrite_content_type);
+  void populateResponseHeaders(Http::HeaderMap& response_headers, bool from_local_reply);
   void appendRequestHeaders(Http::HeaderMapPtr& request_headers_to_add);
 
   Http::Context& httpContext() { return config_->httpContext(); }

source/extensions/filters/network/ratelimit/BUILD

@@ -24,6 +24,7 @@ envoy_cc_library(
         "//include/envoy/stats:stats_macros",
         "//source/common/tracing:http_tracer_lib",
         "//source/extensions/filters/common/ratelimit:ratelimit_client_interface",
+        "//source/extensions/filters/network:well_known_names",
         "@envoy_api//envoy/extensions/filters/network/ratelimit/v3:pkg_cc_proto",
     ],
 )

source/extensions/filters/network/ratelimit/ratelimit.cc

@@ -9,6 +9,8 @@
 #include "common/common/fmt.h"
 #include "common/tracing/http_tracer_impl.h"
 
+#include "extensions/filters/network/well_known_names.h"
+
 namespace Envoy {
 namespace Extensions {
 namespace NetworkFilters {
@@ -69,9 +71,14 @@ void Filter::onEvent(Network::ConnectionEvent event) {
   }
 }
 
-void Filter::complete(Filters::Common::RateLimit::LimitStatus status, Http::ResponseHeaderMapPtr&&,
-                      Http::RequestHeaderMapPtr&&,
-                      const std::string&) {
+void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
+                      Http::ResponseHeaderMapPtr&&, Http::RequestHeaderMapPtr&&, const std::string&,
+                      Filters::Common::RateLimit::DynamicMetadataPtr&& dynamic_metadata) {
+  if (dynamic_metadata != nullptr && !dynamic_metadata->fields().empty()) {
+    filter_callbacks_->connection().streamInfo().setDynamicMetadata(
+        NetworkFilterNames::get().RateLimit, *dynamic_metadata);
+  }
+
   status_ = Status::Complete;
   config_->stats().active_.dec();
 

source/extensions/filters/network/ratelimit/ratelimit.h

@@ -94,7 +94,8 @@ class Filter : public Network::ReadFilter,
   void complete(Filters::Common::RateLimit::LimitStatus status,
                 Http::ResponseHeaderMapPtr&& response_headers_to_add,
                 Http::RequestHeaderMapPtr&& request_headers_to_add,
-                const std::string &response_body) override;
+                const std::string& response_body,
+                Filters::Common::RateLimit::DynamicMetadataPtr&& dynamic_metadata) override;
 
 private:
   enum class Status { NotStarted, Calling, Complete };

source/extensions/filters/network/thrift_proxy/filters/BUILD

@@ -39,6 +39,7 @@ envoy_cc_library(
         "//source/extensions/filters/network/thrift_proxy:decoder_events_lib",
         "//source/extensions/filters/network/thrift_proxy:protocol_interface",
         "//source/extensions/filters/network/thrift_proxy:thrift_lib",
+        "//source/extensions/filters/network/thrift_proxy/filters:well_known_names",
         "//source/extensions/filters/network/thrift_proxy/router:router_interface",
     ],
 )

source/extensions/filters/network/thrift_proxy/filters/ratelimit/BUILD

@@ -22,6 +22,7 @@ envoy_cc_library(
         "//source/extensions/filters/common/ratelimit:stat_names_lib",
         "//source/extensions/filters/network/thrift_proxy:app_exception_lib",
         "//source/extensions/filters/network/thrift_proxy/filters:pass_through_filter_lib",
+        "//source/extensions/filters/network/thrift_proxy/filters:well_known_names",
         "//source/extensions/filters/network/thrift_proxy/router:router_ratelimit_interface",
         "@envoy_api//envoy/extensions/filters/network/thrift_proxy/filters/ratelimit/v3:pkg_cc_proto",
     ],

source/extensions/filters/network/thrift_proxy/filters/ratelimit/ratelimit.cc

@@ -3,6 +3,7 @@
 #include "common/tracing/http_tracer_impl.h"
 
 #include "extensions/filters/network/thrift_proxy/app_exception_impl.h"
+#include "extensions/filters/network/thrift_proxy/filters/well_known_names.h"
 #include "extensions/filters/network/thrift_proxy/router/router.h"
 #include "extensions/filters/network/thrift_proxy/router/router_ratelimit.h"
 
@@ -59,14 +60,19 @@ void Filter::onDestroy() {
 
 void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
                       Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                      Http::RequestHeaderMapPtr&& request_headers_to_add,
-                      const std::string&) {
+                      Http::RequestHeaderMapPtr&& request_headers_to_add, const std::string&,
+                      Filters::Common::RateLimit::DynamicMetadataPtr&& dynamic_metadata) {
   // TODO(zuercher): Store headers to append to a response. Adding them to a local reply (over
   // limit or error) is a matter of modifying the callbacks to allow it. Adding them to an upstream
   // response requires either response (aka encoder) filters or some other mechanism.
   UNREFERENCED_PARAMETER(response_headers_to_add);
   UNREFERENCED_PARAMETER(request_headers_to_add);
 
+  if (dynamic_metadata != nullptr && !dynamic_metadata->fields().empty()) {
+    decoder_callbacks_->streamInfo().setDynamicMetadata(
+        ThriftProxy::ThriftFilters::ThriftFilterNames::get().RATE_LIMIT, *dynamic_metadata);
+  }
+
   state_ = State::Complete;
   Filters::Common::RateLimit::StatNames& stat_names = config_->statNames();
 

source/extensions/filters/network/thrift_proxy/filters/ratelimit/ratelimit.h

@@ -79,7 +79,8 @@ class Filter : public ThriftProxy::ThriftFilters::PassThroughDecoderFilter,
   void complete(Filters::Common::RateLimit::LimitStatus status,
                 Http::ResponseHeaderMapPtr&& response_headers_to_add,
                 Http::RequestHeaderMapPtr&& request_headers_to_add,
-                const std::string& response_body) override;
+                const std::string& response_body,
+                Filters::Common::RateLimit::DynamicMetadataPtr&& dynamic_metadata) override;
 
 private:
   void initiateCall(const ThriftProxy::MessageMetadata& metadata);

test/common/network/filter_manager_impl_test.cc

@@ -415,7 +415,7 @@ stat_prefix: name
       .WillOnce(Return(&conn_pool));
 
   request_callbacks->complete(Extensions::Filters::Common::RateLimit::LimitStatus::OK, nullptr,
-                              nullptr);
+                              nullptr, "", nullptr);
 
   conn_pool.poolReady(upstream_connection);