ratelimit: Add ratelimit custom response headers

api/envoy/service/ratelimit/v2/BUILD

@@ -7,6 +7,7 @@ api_proto_library_internal(
     srcs = ["rls.proto"],
     has_services = 1,
     deps = [
+        "//envoy/api/v2/core:base",
         "//envoy/api/v2/core:grpc_service",
         "//envoy/api/v2/ratelimit",
     ],
@@ -16,6 +17,7 @@ api_go_grpc_library(
     name = "rls",
     proto = ":rls",
     deps = [
+        "//envoy/api/v2/core:base_go_proto",
         "//envoy/api/v2/core:grpc_service_go_proto",
         "//envoy/api/v2/ratelimit:ratelimit_go_proto",
     ],

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

@@ -3,6 +3,7 @@ syntax = "proto3";
 package envoy.service.ratelimit.v2;
 option go_package = "v2";
 
+import "envoy/api/v2/core/base.proto";
 import "envoy/api/v2/ratelimit/ratelimit.proto";
 
 import "validate/validate.proto";
@@ -75,4 +76,6 @@ message RateLimitResponse {
   // in the RateLimitRequest. This can be used by the caller to determine which individual
   // descriptors failed and/or what the currently configured limits are for all of them.
   repeated DescriptorStatus statuses = 2;
+  // A list of headers to add to the response
+  repeated envoy.api.v2.core.HeaderValue headers = 3;
 }

include/envoy/ratelimit/ratelimit.h

@@ -33,9 +33,10 @@ class RequestCallbacks {
   virtual ~RequestCallbacks() {}
 
   /**
-   * Called when a limit request is complete. The resulting status is supplied.
+   * Called when a limit request is complete. The resulting status and
+   * response headers are supplied.
    */
-  virtual void complete(LimitStatus status) PURE;
+  virtual void complete(LimitStatus status, Http::HeaderMapPtr&& headers) PURE;
 };
 
 /**

source/common/http/header_utility.cc

@@ -2,6 +2,7 @@
 
 #include "common/common/utility.h"
 #include "common/config/rds_json.h"
+#include "common/http/header_map_impl.h"
 #include "common/protobuf/utility.h"
 
 #include "absl/strings/match.h"
@@ -115,5 +116,18 @@ bool HeaderUtility::matchHeaders(const Http::HeaderMap& request_headers,
   return match != header_data.invert_match_;
 }
 
+void HeaderUtility::addHeaders(Http::HeaderMap& headers, const Http::HeaderMap& headers_to_add) {
+  headers_to_add.iterate(
+      [](const Http::HeaderEntry& header, void* context) -> Http::HeaderMap::Iterate {
+        Http::HeaderString k;
+        k.setCopy(header.key().c_str(), header.key().size());
+        Http::HeaderString v;
+        v.setCopy(header.value().c_str(), header.value().size());
+        static_cast<Http::HeaderMapImpl*>(context)->addViaMove(std::move(k), std::move(v));
+        return Http::HeaderMap::Iterate::Continue;
+      },
+      &headers);
+}
+
 } // namespace Http
 } // namespace Envoy

source/common/http/header_utility.h

@@ -44,6 +44,13 @@ class HeaderUtility {
                            const std::vector<HeaderData>& config_headers);
 
   static bool matchHeaders(const Http::HeaderMap& request_headers, const HeaderData& config_header);
+
+  /**
+   * Add headers from one HeaderMap to another
+   * @param headers target where headers will be added
+   * @param headers_to_add supplies the headers to be added
+   */
+  static void addHeaders(Http::HeaderMap& headers, const Http::HeaderMap& headers_to_add);
 };
 } // namespace Http
 } // namespace Envoy

source/common/ratelimit/ratelimit_impl.cc

@@ -8,6 +8,7 @@
 #include "envoy/api/v2/ratelimit/ratelimit.pb.h"
 
 #include "common/common/assert.h"
+#include "common/http/header_map_impl.h"
 #include "common/http/headers.h"
 
 namespace Envoy {
@@ -66,14 +67,23 @@ void GrpcClientImpl::onSuccess(
     span.setTag(Constants::get().TraceStatus, Constants::get().TraceOk);
   }
 
-  callbacks_->complete(status);
+  if (response->headers_size()) {
+    Http::HeaderMapPtr headers = std::make_unique<Http::HeaderMapImpl>();
+    for (const auto& h : response->headers()) {
+      headers->addCopy(Http::LowerCaseString(h.key()), h.value());
+    }
+    callbacks_->complete(status, std::move(headers));
+  } else {
+    callbacks_->complete(status, nullptr);
+  }
+
   callbacks_ = nullptr;
 }
 
 void GrpcClientImpl::onFailure(Grpc::Status::GrpcStatus status, const std::string&,
                                Tracing::Span&) {
   ASSERT(status != Grpc::Status::GrpcStatus::Ok);
-  callbacks_->complete(LimitStatus::Error);
+  callbacks_->complete(LimitStatus::Error, nullptr);
   callbacks_ = nullptr;
 }
 

source/common/ratelimit/ratelimit_impl.h

@@ -86,7 +86,7 @@ class NullClientImpl : public Client {
   void cancel() override {}
   void limit(RequestCallbacks& callbacks, const std::string&, const std::vector<Descriptor>&,
              Tracing::Span&) override {
-    callbacks.complete(LimitStatus::OK);
+    callbacks.complete(LimitStatus::OK, nullptr);
   }
 };
 

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

@@ -26,7 +26,7 @@ Http::FilterFactoryCb RateLimitFilterConfig::createFilterFactoryFromProtoTyped(
   const uint32_t timeout_ms = PROTOBUF_GET_MS_OR_DEFAULT(proto_config, timeout, 20);
   return
       [filter_config, timeout_ms, &context](Http::FilterChainFactoryCallbacks& callbacks) -> void {
-        callbacks.addStreamDecoderFilter(std::make_shared<Filter>(
+        callbacks.addStreamFilter(std::make_shared<Filter>(
             filter_config, context.rateLimitClient(std::chrono::milliseconds(timeout_ms))));
       };
 }

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

@@ -9,6 +9,7 @@
 #include "common/common/enum_to_int.h"
 #include "common/common/fmt.h"
 #include "common/http/codes.h"
+#include "common/http/header_utility.h"
 #include "common/router/config_impl.h"
 
 namespace Envoy {
@@ -87,15 +88,35 @@ void Filter::setDecoderFilterCallbacks(Http::StreamDecoderFilterCallbacks& callb
   callbacks_ = &callbacks;
 }
 
+Http::FilterHeadersStatus Filter::encode100ContinueHeaders(Http::HeaderMap&) {
+  return Http::FilterHeadersStatus::Continue;
+}
+
+Http::FilterHeadersStatus Filter::encodeHeaders(Http::HeaderMap& headers, bool) {
+  addHeaders(headers);
+  return Http::FilterHeadersStatus::Continue;
+}
+
+Http::FilterDataStatus Filter::encodeData(Buffer::Instance&, bool) {
+  return Http::FilterDataStatus::Continue;
+}
+
+Http::FilterTrailersStatus Filter::encodeTrailers(Http::HeaderMap&) {
+  return Http::FilterTrailersStatus::Continue;
+}
+
+void Filter::setEncoderFilterCallbacks(Http::StreamEncoderFilterCallbacks&) {}
+
 void Filter::onDestroy() {
   if (state_ == State::Calling) {
     state_ = State::Complete;
     client_->cancel();
   }
 }
 
-void Filter::complete(RateLimit::LimitStatus status) {
+void Filter::complete(RateLimit::LimitStatus status, Http::HeaderMapPtr&& headers) {
   state_ = State::Complete;
+  headers_to_add_ = std::move(headers);
 
   switch (status) {
   case RateLimit::LimitStatus::OK:
@@ -123,7 +144,8 @@ void Filter::complete(RateLimit::LimitStatus status) {
   if (status == RateLimit::LimitStatus::OverLimit &&
       config_->runtime().snapshot().featureEnabled("ratelimit.http_filter_enforcing", 100)) {
     state_ = State::Responded;
-    callbacks_->sendLocalReply(Http::Code::TooManyRequests, "", nullptr);
+    callbacks_->sendLocalReply(Http::Code::TooManyRequests, "",
+                               [this](Http::HeaderMap& headers) { addHeaders(headers); });
     callbacks_->requestInfo().setResponseFlag(RequestInfo::ResponseFlag::RateLimited);
   } else if (!initiating_call_) {
     callbacks_->continueDecoding();
@@ -147,6 +169,13 @@ void Filter::populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_li
   }
 }
 
+void Filter::addHeaders(Http::HeaderMap& headers) {
+  if (headers_to_add_) {
+    Http::HeaderUtility::addHeaders(headers, *headers_to_add_);
+    headers_to_add_ = nullptr;
+  }
+}
+
 } // namespace RateLimitFilter
 } // namespace HttpFilters
 } // namespace Extensions

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

@@ -73,7 +73,7 @@ typedef std::shared_ptr<FilterConfig> FilterConfigSharedPtr;
  * HTTP rate limit filter. Depending on the route configuration, this filter calls the global
  * rate limiting service before allowing further filter iteration.
  */
-class Filter : public Http::StreamDecoderFilter, public RateLimit::RequestCallbacks {
+class Filter : public Http::StreamFilter, public RateLimit::RequestCallbacks {
 public:
   Filter(FilterConfigSharedPtr config, RateLimit::ClientPtr&& client)
       : config_(config), client_(std::move(client)) {}
@@ -87,15 +87,23 @@ class Filter : public Http::StreamDecoderFilter, public RateLimit::RequestCallba
   Http::FilterTrailersStatus decodeTrailers(Http::HeaderMap& trailers) override;
   void setDecoderFilterCallbacks(Http::StreamDecoderFilterCallbacks& callbacks) override;
 
+  // Http::StreamEncoderFilter
+  Http::FilterHeadersStatus encode100ContinueHeaders(Http::HeaderMap& headers) override;
+  Http::FilterHeadersStatus encodeHeaders(Http::HeaderMap& headers, bool end_stream) override;
+  Http::FilterDataStatus encodeData(Buffer::Instance& data, bool end_stream) override;
+  Http::FilterTrailersStatus encodeTrailers(Http::HeaderMap& trailers) override;
+  void setEncoderFilterCallbacks(Http::StreamEncoderFilterCallbacks& callbacks) override;
+
   // RateLimit::RequestCallbacks
-  void complete(RateLimit::LimitStatus status) override;
+  void complete(RateLimit::LimitStatus status, Http::HeaderMapPtr&& headers) override;
 
 private:
   void initiateCall(const Http::HeaderMap& headers);
   void populateRateLimitDescriptors(const Router::RateLimitPolicy& rate_limit_policy,
                                     std::vector<RateLimit::Descriptor>& descriptors,
                                     const Router::RouteEntry* route_entry,
                                     const Http::HeaderMap& headers) const;
+  void addHeaders(Http::HeaderMap& headers);
 
   enum class State { NotStarted, Calling, Complete, Responded };
 
@@ -105,6 +113,7 @@ class Filter : public Http::StreamDecoderFilter, public RateLimit::RequestCallba
   State state_{State::NotStarted};
   Upstream::ClusterInfoConstSharedPtr cluster_;
   bool initiating_call_{};
+  Http::HeaderMapPtr headers_to_add_;
 };
 
 } // namespace RateLimitFilter

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

@@ -67,7 +67,7 @@ void Filter::onEvent(Network::ConnectionEvent event) {
   }
 }
 
-void Filter::complete(RateLimit::LimitStatus status) {
+void Filter::complete(RateLimit::LimitStatus status, Http::HeaderMapPtr&&) {
   status_ = Status::Complete;
   config_->stats().active_.dec();
 

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

@@ -87,7 +87,7 @@ class Filter : public Network::ReadFilter,
   void onBelowWriteBufferLowWatermark() override {}
 
   // RateLimit::RequestCallbacks
-  void complete(RateLimit::LimitStatus status) override;
+  void complete(RateLimit::LimitStatus status, Http::HeaderMapPtr&& headers) override;
 
 private:
   enum class Status { NotStarted, Calling, Complete };
@@ -98,7 +98,7 @@ class Filter : public Network::ReadFilter,
   Status status_{Status::NotStarted};
   bool calling_limit_{};
 };
-}
+} // namespace RateLimitFilter
 } // namespace NetworkFilters
 } // namespace Extensions
 } // namespace Envoy

test/common/http/header_utility_test.cc

@@ -403,5 +403,21 @@ invert_match: true
   EXPECT_FALSE(HeaderUtility::matchHeaders(unmatching_headers, header_data));
 }
 
+TEST(HeaderAddTest, HeaderAdd) {
+  TestHeaderMapImpl headers{{"myheader1", "123value"}};
+  TestHeaderMapImpl headers_to_add{{"myheader2", "456value"}};
+
+  HeaderUtility::addHeaders(headers, headers_to_add);
+
+  headers_to_add.iterate(
+      [](const Http::HeaderEntry& entry, void* context) -> Http::HeaderMap::Iterate {
+        TestHeaderMapImpl* headers = static_cast<TestHeaderMapImpl*>(context);
+        Http::LowerCaseString lower_key{entry.key().c_str()};
+        EXPECT_STREQ(entry.value().c_str(), headers->get(lower_key)->value().c_str());
+        return Http::HeaderMap::Iterate::Continue;
+      },
+      &headers);
+}
+
 } // namespace Http
 } // namespace Envoy

test/common/network/filter_manager_impl_test.cc

@@ -207,7 +207,7 @@ TEST_F(NetworkFilterManagerTest, RateLimitAndTcpProxy) {
   EXPECT_CALL(factory_context.cluster_manager_, tcpConnPoolForCluster("fake_cluster", _, _))
       .WillOnce(Return(&conn_pool));
 
-  request_callbacks->complete(RateLimit::LimitStatus::OK);
+  request_callbacks->complete(RateLimit::LimitStatus::OK, nullptr);
 
   conn_pool.poolReady(upstream_connection);
 

test/common/ratelimit/ratelimit_impl_test.cc

@@ -27,7 +27,11 @@ namespace RateLimit {
 
 class MockRequestCallbacks : public RequestCallbacks {
 public:
-  MOCK_METHOD1(complete, void(LimitStatus status));
+  void complete(LimitStatus status, Http::HeaderMapPtr&& headers) {
+    complete_(status, headers.get());
+  }
+
+  MOCK_METHOD2(complete_, void(LimitStatus status, const Http::HeaderMap* headers));
 };
 
 // TODO(junr03): legacy rate limit is deprecated. Remove the boolean parameter after 1.8.0.
@@ -89,7 +93,7 @@ TEST_P(RateLimitGrpcClientTest, Basic) {
     response.reset(new envoy::service::ratelimit::v2::RateLimitResponse());
     response->set_overall_code(envoy::service::ratelimit::v2::RateLimitResponse_Code_OVER_LIMIT);
     EXPECT_CALL(span_, setTag("ratelimit_status", "over_limit"));
-    EXPECT_CALL(request_callbacks_, complete(LimitStatus::OverLimit));
+    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OverLimit, _));
     client_->onSuccess(std::move(response), span_);
   }
 
@@ -108,7 +112,7 @@ TEST_P(RateLimitGrpcClientTest, Basic) {
     response.reset(new envoy::service::ratelimit::v2::RateLimitResponse());
     response->set_overall_code(envoy::service::ratelimit::v2::RateLimitResponse_Code_OK);
     EXPECT_CALL(span_, setTag("ratelimit_status", "ok"));
-    EXPECT_CALL(request_callbacks_, complete(LimitStatus::OK));
+    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::OK, _));
     client_->onSuccess(std::move(response), span_);
   }
 
@@ -125,7 +129,7 @@ TEST_P(RateLimitGrpcClientTest, Basic) {
                    Tracing::NullSpan::instance());
 
     response.reset(new envoy::service::ratelimit::v2::RateLimitResponse());
-    EXPECT_CALL(request_callbacks_, complete(LimitStatus::Error));
+    EXPECT_CALL(request_callbacks_, complete_(LimitStatus::Error, _));
     client_->onFailure(Grpc::Status::Unknown, "", span_);
   }
 }
@@ -177,7 +181,7 @@ TEST(RateLimitNullFactoryTest, Basic) {
   NullFactoryImpl factory;
   ClientPtr client = factory.create(absl::optional<std::chrono::milliseconds>());
   MockRequestCallbacks request_callbacks;
-  EXPECT_CALL(request_callbacks, complete(LimitStatus::OK));
+  EXPECT_CALL(request_callbacks, complete_(LimitStatus::OK, _));
   client->limit(request_callbacks, "foo", {{{{"foo", "bar"}}}}, Tracing::NullSpan::instance());
   client->cancel();
 }

test/extensions/filters/http/ratelimit/config_test.cc

@@ -36,7 +36,7 @@ TEST(RateLimitFilterConfigTest, RateLimitFilterCorrectJson) {
   RateLimitFilterConfig factory;
   Http::FilterFactoryCb cb = factory.createFilterFactory(*json_config, "stats", context);
   Http::MockFilterChainFactoryCallbacks filter_callback;
-  EXPECT_CALL(filter_callback, addStreamDecoderFilter(_));
+  EXPECT_CALL(filter_callback, addStreamFilter(_));
   cb(filter_callback);
 }
 
@@ -56,7 +56,7 @@ TEST(RateLimitFilterConfigTest, RateLimitFilterCorrectProto) {
   RateLimitFilterConfig factory;
   Http::FilterFactoryCb cb = factory.createFilterFactoryFromProto(proto_config, "stats", context);
   Http::MockFilterChainFactoryCallbacks filter_callback;
-  EXPECT_CALL(filter_callback, addStreamDecoderFilter(_));
+  EXPECT_CALL(filter_callback, addStreamFilter(_));
   cb(filter_callback);
 }
 
@@ -79,7 +79,7 @@ TEST(RateLimitFilterConfigTest, RateLimitFilterEmptyProto) {
 
   Http::FilterFactoryCb cb = factory.createFilterFactory(*json_config, "stats", context);
   Http::MockFilterChainFactoryCallbacks filter_callback;
-  EXPECT_CALL(filter_callback, addStreamDecoderFilter(_));
+  EXPECT_CALL(filter_callback, addStreamFilter(_));
   cb(filter_callback);
 }