ratelimit: support returning custom response bodies for non-OK respon…

…ses from the external ratelimit service (envoyproxy#14189)

Signed-off-by: John Esmet <john.esmet@gmail.com>

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

@@ -51,6 +51,7 @@ message RateLimitRequest {
 }
 
 // A response from a ShouldRateLimit call.
+// [#next-free-field: 6]
 message RateLimitResponse {
   option (udpa.annotations.versioning).previous_message_type =
       "envoy.service.ratelimit.v2.RateLimitResponse";
@@ -131,4 +132,7 @@ message RateLimitResponse {
 
   // A list of headers to add to the request when forwarded
   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.
+  bytes raw_body = 5;
 }

docs/root/version_history/current.rst

@@ -81,6 +81,7 @@ New Features
 * overload: add :ref:`envoy.overload_actions.reduce_timeouts <config_overload_manager_overload_actions>` overload action to enable scaling timeouts down with load. Scaling support :ref:`is limited <envoy_v3_api_enum_config.overload.v3.ScaleTimersOverloadActionConfig.TimerType>` to the HTTP connection and stream idle timeouts.
 * ratelimit: added support for use of various :ref:`metadata <envoy_v3_api_field_config.route.v3.RateLimit.Action.metadata>` as a ratelimit action.
 * ratelimit: added :ref:`disable_x_envoy_ratelimited_header <envoy_v3_api_msg_extensions.filters.http.ratelimit.v3.RateLimit>` option to disable `X-Envoy-RateLimited` header.
+* ratelimit: added :ref:`body <envoy_v3_api_field_service.ratelimit.v3.RateLimitResponse.raw_body>` field to support custom response bodies for non-OK responses from the external ratelimit service.
 * router: added support for regex rewrites during HTTP redirects using :ref:`regex_rewrite <envoy_v3_api_field_config.route.v3.RedirectAction.regex_rewrite>`.
 * sds: improved support for atomic :ref:`key rotations <xds_certificate_rotation>` and added configurable rotation triggers for
   :ref:`TlsCertificate <envoy_v3_api_field_extensions.transport_sockets.tls.v3.TlsCertificate.watched_directory>` and

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

@@ -44,12 +44,20 @@ class RequestCallbacks {
   virtual ~RequestCallbacks() = default;
 
   /**
-   * Called when a limit request is complete. The resulting status,
-   * response headers and request headers to be forwarded to the upstream are supplied.
+   * Called when a limit request is complete. The resulting status, response headers
+   * and request headers to be forwarded to the upstream are supplied.
+   *
+   * @status The ratelimit status
+   * @descriptor_statuses The descriptor statuses
+   * @response_headers_to_add The headers to add to the downstream response, for non-OK statuses
+   * @request_headers_to_add The headers to add to the upstream request, if not ratelimited
+   * @response_body The response body to use for the downstream response, for non-OK statuses. May
+   * contain non UTF-8 values (e.g. binary data).
    */
   virtual void complete(LimitStatus status, DescriptorStatusListPtr&& descriptor_statuses,
                         Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                        Http::RequestHeaderMapPtr&& request_headers_to_add) PURE;
+                        Http::RequestHeaderMapPtr&& request_headers_to_add,
+                        const std::string& response_body) PURE;
 };
 
 /**

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

@@ -106,14 +106,14 @@ void GrpcClientImpl::onSuccess(
   DescriptorStatusListPtr descriptor_statuses = std::make_unique<DescriptorStatusList>(
       response->statuses().begin(), response->statuses().end());
   callbacks_->complete(status, std::move(descriptor_statuses), std::move(response_headers_to_add),
-                       std::move(request_headers_to_add));
+                       std::move(request_headers_to_add), response->raw_body());
   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, nullptr);
+  callbacks_->complete(LimitStatus::Error, nullptr, nullptr, nullptr, EMPTY_STRING);
   callbacks_ = nullptr;
 }
 

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

@@ -115,7 +115,7 @@ Http::FilterHeadersStatus Filter::encode100ContinueHeaders(Http::ResponseHeaderM
 }
 
 Http::FilterHeadersStatus Filter::encodeHeaders(Http::ResponseHeaderMap& headers, bool) {
-  populateResponseHeaders(headers);
+  populateResponseHeaders(headers, /*from_local_reply=*/false);
   return Http::FilterHeadersStatus::Continue;
 }
 
@@ -143,7 +143,8 @@ void Filter::onDestroy() {
 void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
                       Filters::Common::RateLimit::DescriptorStatusListPtr&& descriptor_statuses,
                       Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                      Http::RequestHeaderMapPtr&& request_headers_to_add) {
+                      Http::RequestHeaderMapPtr&& request_headers_to_add,
+                      const std::string& response_body) {
   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);
@@ -195,8 +196,10 @@ void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
       config_->runtime().snapshot().featureEnabled("ratelimit.http_filter_enforcing", 100)) {
     state_ = State::Responded;
     callbacks_->sendLocalReply(
-        Http::Code::TooManyRequests, "",
-        [this](Http::HeaderMap& headers) { populateResponseHeaders(headers); },
+        Http::Code::TooManyRequests, response_body,
+        [this](Http::HeaderMap& headers) {
+          populateResponseHeaders(headers, /*from_local_reply=*/true);
+        },
         config_->rateLimitedGrpcStatus(), RcDetails::get().RateLimited);
     callbacks_->streamInfo().setResponseFlag(StreamInfo::ResponseFlag::RateLimited);
   } else if (status == Filters::Common::RateLimit::LimitStatus::Error) {
@@ -208,8 +211,8 @@ void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
       }
     } else {
       state_ = State::Responded;
-      callbacks_->sendLocalReply(Http::Code::InternalServerError, "", nullptr, absl::nullopt,
-                                 RcDetails::get().RateLimitError);
+      callbacks_->sendLocalReply(Http::Code::InternalServerError, response_body, nullptr,
+                                 absl::nullopt, RcDetails::get().RateLimitError);
       callbacks_->streamInfo().setResponseFlag(StreamInfo::ResponseFlag::RateLimitServiceError);
     }
   } else if (!initiating_call_) {
@@ -236,8 +239,18 @@ void Filter::populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_li
   }
 }
 
-void Filter::populateResponseHeaders(Http::HeaderMap& response_headers) {
+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 we're
+    // populating response headers for a local reply, overwrite the existing
+    // content-type header.
+    //
+    // 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_);
     response_headers_to_add_ = nullptr;
   }

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

@@ -148,15 +148,16 @@ class Filter : public Http::StreamFilter, public Filters::Common::RateLimit::Req
   void complete(Filters::Common::RateLimit::LimitStatus status,
                 Filters::Common::RateLimit::DescriptorStatusListPtr&& descriptor_statuses,
                 Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                Http::RequestHeaderMapPtr&& request_headers_to_add) override;
+                Http::RequestHeaderMapPtr&& request_headers_to_add,
+                const std::string& response_body) 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);
+  void populateResponseHeaders(Http::HeaderMap& response_headers, bool from_local_reply);
   void appendRequestHeaders(Http::HeaderMapPtr& request_headers_to_add);
   VhRateLimitOptions getVirtualHostRateLimitOption(const Router::RouteConstSharedPtr& route);
 

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

@@ -72,7 +72,8 @@ void Filter::onEvent(Network::ConnectionEvent event) {
 
 void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
                       Filters::Common::RateLimit::DescriptorStatusListPtr&&,
-                      Http::ResponseHeaderMapPtr&&, Http::RequestHeaderMapPtr&&) {
+                      Http::ResponseHeaderMapPtr&&, Http::RequestHeaderMapPtr&&,
+                      const std::string&) {
   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,
                 Filters::Common::RateLimit::DescriptorStatusListPtr&& descriptor_statuses,
                 Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                Http::RequestHeaderMapPtr&& request_headers_to_add) override;
+                Http::RequestHeaderMapPtr&& request_headers_to_add,
+                const std::string& response_body) override;
 
 private:
   enum class Status { NotStarted, Calling, Complete };

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

@@ -62,7 +62,7 @@ void Filter::onDestroy() {
 void Filter::complete(Filters::Common::RateLimit::LimitStatus status,
                       Filters::Common::RateLimit::DescriptorStatusListPtr&& descriptor_statuses,
                       Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                      Http::RequestHeaderMapPtr&& request_headers_to_add) {
+                      Http::RequestHeaderMapPtr&& request_headers_to_add, const std::string&) {
   // 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.

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,
                 Filters::Common::RateLimit::DescriptorStatusListPtr&& descriptor_statuses,
                 Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                Http::RequestHeaderMapPtr&& request_headers_to_add) override;
+                Http::RequestHeaderMapPtr&& request_headers_to_add,
+                const std::string& response_body) override;
 
 private:
   void initiateCall(const ThriftProxy::MessageMetadata& metadata);

test/common/network/filter_manager_impl_test.cc

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

test/extensions/filters/common/ratelimit/ratelimit_impl_test.cc

@@ -38,15 +38,17 @@ class MockRequestCallbacks : public RequestCallbacks {
 public:
   void complete(LimitStatus status, DescriptorStatusListPtr&& descriptor_statuses,
                 Http::ResponseHeaderMapPtr&& response_headers_to_add,
-                Http::RequestHeaderMapPtr&& request_headers_to_add) override {
+                Http::RequestHeaderMapPtr&& request_headers_to_add,
+                const std::string& response_body) override {
     complete_(status, descriptor_statuses.get(), response_headers_to_add.get(),
-              request_headers_to_add.get());
+              request_headers_to_add.get(), response_body);
   }
 
   MOCK_METHOD(void, complete_,
               (LimitStatus status, const DescriptorStatusList* descriptor_statuses,
                const Http::ResponseHeaderMap* response_headers_to_add,
-               const Http::RequestHeaderMap* request_headers_to_add));
+               const Http::RequestHeaderMap* request_headers_to_add,
+               const std::string& response_body));
 };
 
 class RateLimitGrpcClientTest : public testing::Test {
@@ -91,7 +93,7 @@ TEST_F(RateLimitGrpcClientTest, Basic) {
     response = std::make_unique<envoy::service::ratelimit::v3::RateLimitResponse>();
     response->set_overall_code(envoy::service::ratelimit::v3::RateLimitResponse::OVER_LIMIT);
     EXPECT_CALL(span_, setTag(Eq("ratelimit_status"), Eq("over_limit")));
-    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OverLimit, _, _, _));
+    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OverLimit, _, _, _, _));
     client_.onSuccess(std::move(response), span_);
   }
 
@@ -110,7 +112,7 @@ TEST_F(RateLimitGrpcClientTest, Basic) {
     response = std::make_unique<envoy::service::ratelimit::v3::RateLimitResponse>();
     response->set_overall_code(envoy::service::ratelimit::v3::RateLimitResponse::OK);
     EXPECT_CALL(span_, setTag(Eq("ratelimit_status"), Eq("ok")));
-    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OK, _, _, _));
+    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OK, _, _, _, _));
     client_.onSuccess(std::move(response), span_);
   }
 
@@ -127,7 +129,7 @@ TEST_F(RateLimitGrpcClientTest, Basic) {
                   Tracing::NullSpan::instance(), stream_info_);
 
     response = std::make_unique<envoy::service::ratelimit::v3::RateLimitResponse>();
-    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::Error, _, _, _));
+    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::Error, _, _, _, _));
     client_.onFailure(Grpc::Status::Unknown, "", span_);
   }
 
@@ -150,7 +152,7 @@ TEST_F(RateLimitGrpcClientTest, Basic) {
     response = std::make_unique<envoy::service::ratelimit::v3::RateLimitResponse>();
     response->set_overall_code(envoy::service::ratelimit::v3::RateLimitResponse::OK);
     EXPECT_CALL(span_, setTag(Eq("ratelimit_status"), Eq("ok")));
-    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OK, _, _, _));
+    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OK, _, _, _, _));
     client_.onSuccess(std::move(response), span_);
   }
 }