ratelimit: add support for x-ratelimit-* headers in local rate limiting

api/envoy/extensions/common/ratelimit/v3/ratelimit.proto

@@ -16,6 +16,15 @@ option (udpa.annotations.file_status).package_version_status = ACTIVE;
 
 // [#protodoc-title: Common rate limit components]
 
+// Defines the version of the standard to use for X-RateLimit headers.
+enum XRateLimitHeadersRFCVersion {
+  // X-RateLimit headers disabled.
+  OFF = 0;
+
+  // Use `draft RFC Version 03 <https://tools.ietf.org/id/draft-polli-ratelimit-headers-03.html>`_.
+  DRAFT_VERSION_03 = 1;
+}
+
 // A RateLimitDescriptor is a list of hierarchical entries that are used by the service to
 // determine the final rate limit key and overall allowed limit. Here are some examples of how
 // they might be used for the domain "envoy".

api/envoy/extensions/filters/http/local_ratelimit/v3/local_rate_limit.proto

@@ -19,7 +19,7 @@ option (udpa.annotations.file_status).package_version_status = ACTIVE;
 // Local Rate limit :ref:`configuration overview <config_http_filters_local_rate_limit>`.
 // [#extension: envoy.filters.http.local_ratelimit]
 
-// [#next-free-field: 12]
+// [#next-free-field: 13]
 message LocalRateLimit {
   // The human readable prefix to use when emitting stats.
   string stat_prefix = 1 [(validate.rules).string = {min_len: 1}];
@@ -106,4 +106,21 @@ message LocalRateLimit {
   // one to rate limit requests on a per connection basis.
   // If unspecified, the default value is false.
   bool local_rate_limit_per_downstream_connection = 11;
+
+  // Defines the standard version to use for X-RateLimit headers emitted by the filter:
+  //
+  // * ``X-RateLimit-Limit`` - indicates the request-quota associated to the
+  //   client in the current time-window followed by the description of the
+  //   quota policy. The value is returned by the maximum tokens of the token bucket.
+  // * ``X-RateLimit-Remaining`` - indicates the remaining requests in the
+  //   current time-window. The value is returned by the remaining tokens in the token bucket.
+  // * ``X-RateLimit-Reset`` - indicates the number of seconds until reset of
+  //   the current time-window. The value is returned by the remaining fill interval of the token bucket.
+  //
+  // For more information about the headers specification see selected version of
+  // the `draft RFC <https://tools.ietf.org/id/draft-polli-ratelimit-headers-03.html>`_.
+  //
+  // Disabled by default.
+  common.ratelimit.v3.XRateLimitHeadersRFCVersion enable_x_ratelimit_headers = 12
+      [(validate.rules).enum = {defined_only: true}];
 }

api/envoy/extensions/filters/http/ratelimit/v3/BUILD

@@ -7,6 +7,7 @@ licenses(["notice"])  # Apache 2
 api_proto_package(
     deps = [
         "//envoy/config/ratelimit/v3:pkg",
+        "//envoy/extensions/common/ratelimit/v3:pkg",
         "@com_github_cncf_udpa//udpa/annotations:pkg",
     ],
 )

api/envoy/extensions/filters/http/ratelimit/v3/rate_limit.proto

@@ -3,6 +3,7 @@ syntax = "proto3";
 package envoy.extensions.filters.http.ratelimit.v3;
 
 import "envoy/config/ratelimit/v3/rls.proto";
+import "envoy/extensions/common/ratelimit/v3/ratelimit.proto";
 
 import "google/protobuf/duration.proto";
 
@@ -24,15 +25,6 @@ message RateLimit {
   option (udpa.annotations.versioning).previous_message_type =
       "envoy.config.filter.http.rate_limit.v2.RateLimit";
 
-  // Defines the version of the standard to use for X-RateLimit headers.
-  enum XRateLimitHeadersRFCVersion {
-    // X-RateLimit headers disabled.
-    OFF = 0;
-
-    // Use `draft RFC Version 03 <https://tools.ietf.org/id/draft-polli-ratelimit-headers-03.html>`_.
-    DRAFT_VERSION_03 = 1;
-  }
-
   // The rate limit domain to use when calling the rate limit service.
   string domain = 1 [(validate.rules).string = {min_len: 1}];
 
@@ -96,7 +88,7 @@ message RateLimit {
   // the `draft RFC <https://tools.ietf.org/id/draft-polli-ratelimit-headers-03.html>`_.
   //
   // Disabled by default.
-  XRateLimitHeadersRFCVersion enable_x_ratelimit_headers = 8
+  common.ratelimit.v3.XRateLimitHeadersRFCVersion enable_x_ratelimit_headers = 8
       [(validate.rules).enum = {defined_only: true}];
 
   // Disables emitting the :ref:`x-envoy-ratelimited<config_http_filters_router_x-envoy-ratelimited>` header

source/extensions/filters/common/local_ratelimit/local_ratelimit_impl.cc

@@ -1,5 +1,7 @@
 #include "source/extensions/filters/common/local_ratelimit/local_ratelimit_impl.h"
 
+#include <chrono>
+
 #include "source/common/protobuf/utility.h"
 
 namespace Envoy {
@@ -12,7 +14,8 @@ LocalRateLimiterImpl::LocalRateLimiterImpl(
     const std::chrono::milliseconds fill_interval, const uint32_t max_tokens,
     const uint32_t tokens_per_fill, Event::Dispatcher& dispatcher,
     const Protobuf::RepeatedPtrField<
-        envoy::extensions::common::ratelimit::v3::LocalRateLimitDescriptor>& descriptors)
+        envoy::extensions::common::ratelimit::v3::LocalRateLimitDescriptor>& descriptors,
+    bool init_fill_time)
     : fill_timer_(fill_interval > std::chrono::milliseconds(0)
                       ? dispatcher.createTimer([this] { onFillTimer(); })
                       : nullptr),
@@ -25,6 +28,9 @@ LocalRateLimiterImpl::LocalRateLimiterImpl(
   token_bucket_.tokens_per_fill_ = tokens_per_fill;
   token_bucket_.fill_interval_ = absl::FromChrono(fill_interval);
   tokens_.tokens_ = max_tokens;
+  if (init_fill_time) {
+    tokens_.fill_time_ = time_source_.monotonicTime();
+  }
 
   if (fill_timer_) {
     fill_timer_->enableTimer(fill_interval);
@@ -68,6 +74,7 @@ LocalRateLimiterImpl::~LocalRateLimiterImpl() {
 
 void LocalRateLimiterImpl::onFillTimer() {
   onFillTimerHelper(tokens_, token_bucket_);
+  tokens_.fill_time_ = time_source_.monotonicTime();
   onFillTimerDescriptorHelper();
   fill_timer_->enableTimer(absl::ToChronoMilliseconds(token_bucket_.fill_interval_));
 }
@@ -136,6 +143,51 @@ bool LocalRateLimiterImpl::requestAllowed(
   return requestAllowedHelper(tokens_);
 }
 
+uint32_t LocalRateLimiterImpl::maxTokens(
+    absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const {
+  if (!descriptors_.empty() && !request_descriptors.empty()) {
+    for (const auto& request_descriptor : request_descriptors) {
+      auto it = descriptors_.find(request_descriptor);
+      if (it != descriptors_.end()) {
+        return it->token_bucket_.max_tokens_;
+      }
+    }
+  }
+  return token_bucket_.max_tokens_;
+}
+
+uint32_t LocalRateLimiterImpl::remainingTokens(
+    absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const {
+  if (!descriptors_.empty() && !request_descriptors.empty()) {
+    for (const auto& request_descriptor : request_descriptors) {
+      auto it = descriptors_.find(request_descriptor);
+      if (it != descriptors_.end()) {
+        return it->token_state_->tokens_.load(std::memory_order_relaxed);
+      }
+    }
+  }
+  return tokens_.tokens_.load(std::memory_order_relaxed);
+}
+
+uint32_t LocalRateLimiterImpl::remainingFillInterval(
+    absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const {
+  using namespace std::literals;
+
+  if (!descriptors_.empty() && !request_descriptors.empty()) {
+    for (const auto& request_descriptor : request_descriptors) {
+      auto it = descriptors_.find(request_descriptor);
+      if (it != descriptors_.end()) {
+        return static_cast<uint32_t>(absl::ToInt64Seconds(
+            it->token_bucket_.fill_interval_ -
+            absl::Seconds((time_source_.monotonicTime() - it->token_state_->fill_time_) / 1s)));
+      }
+    }
+  }
+  return static_cast<uint32_t>(absl::ToInt64Seconds(
+      token_bucket_.fill_interval_ -
+      absl::Seconds((time_source_.monotonicTime() - tokens_.fill_time_) / 1s)));
+}
+
 } // namespace LocalRateLimit
 } // namespace Common
 } // namespace Filters

source/extensions/filters/common/local_ratelimit/local_ratelimit_impl.h

@@ -22,10 +22,15 @@ class LocalRateLimiterImpl {
       const std::chrono::milliseconds fill_interval, const uint32_t max_tokens,
       const uint32_t tokens_per_fill, Event::Dispatcher& dispatcher,
       const Protobuf::RepeatedPtrField<
-          envoy::extensions::common::ratelimit::v3::LocalRateLimitDescriptor>& descriptors);
+          envoy::extensions::common::ratelimit::v3::LocalRateLimitDescriptor>& descriptors,
+      bool init_fill_time);
   ~LocalRateLimiterImpl();
 
   bool requestAllowed(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const;
+  uint32_t maxTokens(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const;
+  uint32_t remainingTokens(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const;
+  uint32_t
+  remainingFillInterval(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const;
 
 private:
   struct TokenState {

source/extensions/filters/http/local_ratelimit/BUILD

@@ -28,6 +28,7 @@ envoy_cc_library(
         "//source/extensions/filters/common/local_ratelimit:local_ratelimit_lib",
         "//source/extensions/filters/common/ratelimit:ratelimit_lib",
         "//source/extensions/filters/http/common:pass_through_filter_lib",
+        "@envoy_api//envoy/extensions/common/ratelimit/v3:pkg_cc_proto",
         "@envoy_api//envoy/extensions/filters/http/local_ratelimit/v3:pkg_cc_proto",
     ],
 )

source/extensions/filters/http/local_ratelimit/local_ratelimit.cc

@@ -1,8 +1,11 @@
 #include "source/extensions/filters/http/local_ratelimit/local_ratelimit.h"
 
+#include <memory>
 #include <string>
 #include <vector>
 
+#include "envoy/extensions/common/ratelimit/v3/ratelimit.pb.h"
+#include "envoy/extensions/filters/http/local_ratelimit/v3/local_rate_limit.pb.h"
 #include "envoy/http/codes.h"
 
 #include "source/common/http/utility.h"
@@ -17,6 +20,26 @@ const std::string& PerConnectionRateLimiter::key() {
   CONSTRUCT_ON_FIRST_USE(std::string, "per_connection_local_rate_limiter");
 }
 
+LocalRateLimitRequestDescriptorsQueue::LocalRateLimitRequestDescriptorsQueue()
+    : request_descriptors_queue_(
+          std::queue<absl::optional<std::vector<RateLimit::LocalDescriptor>>>()) {}
+
+void LocalRateLimitRequestDescriptorsQueue::push(
+    absl::optional<std::vector<RateLimit::LocalDescriptor>> request_descriptors) {
+  request_descriptors_queue_.push(request_descriptors);
+}
+
+absl::optional<std::vector<RateLimit::LocalDescriptor>>
+LocalRateLimitRequestDescriptorsQueue::pop() {
+  auto request_descriptors = std::move(request_descriptors_queue_.front());
+  request_descriptors_queue_.pop();
+  return request_descriptors;
+}
+
+const std::string& LocalRateLimitRequestDescriptorsQueue::key() {
+  CONSTRUCT_ON_FIRST_USE(std::string, "local_rate_limit_requestdescriptors_queue");
+}
+
 FilterConfig::FilterConfig(
     const envoy::extensions::filters::http::local_ratelimit::v3::LocalRateLimit& config,
     const LocalInfo::LocalInfo& local_info, Event::Dispatcher& dispatcher, Stats::Scope& scope,
@@ -30,7 +53,7 @@ FilterConfig::FilterConfig(
       descriptors_(config.descriptors()),
       rate_limit_per_connection_(config.local_rate_limit_per_downstream_connection()),
       rate_limiter_(Filters::Common::LocalRateLimit::LocalRateLimiterImpl(
-          fill_interval_, max_tokens_, tokens_per_fill_, dispatcher, descriptors_)),
+          fill_interval_, max_tokens_, tokens_per_fill_, dispatcher, descriptors_, true)),
       local_info_(local_info), runtime_(runtime),
       filter_enabled_(
           config.has_filter_enabled()
@@ -47,7 +70,9 @@ FilterConfig::FilterConfig(
       request_headers_parser_(Envoy::Router::HeaderParser::configure(
           config.request_headers_to_add_when_not_enforced())),
       stage_(static_cast<uint64_t>(config.stage())),
-      has_descriptors_(!config.descriptors().empty()) {
+      has_descriptors_(!config.descriptors().empty()),
+      enable_x_rate_limit_headers_(config.enable_x_ratelimit_headers() ==
+                                   envoy::extensions::common::ratelimit::v3::DRAFT_VERSION_03) {
   // Note: no token bucket is fine for the global config, which would be the case for enabling
   //       the filter globally but disabled and then applying limits at the virtual host or
   //       route level. At the virtual or route level, it makes no sense to have an no token
@@ -63,6 +88,21 @@ bool FilterConfig::requestAllowed(
   return rate_limiter_.requestAllowed(request_descriptors);
 }
 
+uint32_t
+FilterConfig::maxTokens(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const {
+  return rate_limiter_.maxTokens(request_descriptors);
+}
+
+uint32_t FilterConfig::remainingTokens(
+    absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const {
+  return rate_limiter_.remainingTokens(request_descriptors);
+}
+
+uint32_t FilterConfig::remainingFillInterval(
+    absl::Span<const RateLimit::LocalDescriptor> request_descriptors) const {
+  return rate_limiter_.remainingFillInterval(request_descriptors);
+}
+
 LocalRateLimitStats FilterConfig::generateStats(const std::string& prefix, Stats::Scope& scope) {
   const std::string final_prefix = prefix + ".http_local_rate_limit";
   return {ALL_LOCAL_RATE_LIMIT_STATS(POOL_COUNTER_PREFIX(scope, final_prefix))};
@@ -80,6 +120,7 @@ Http::FilterHeadersStatus Filter::decodeHeaders(Http::RequestHeaderMap& headers,
   const auto* config = getConfig();
 
   if (!config->enabled()) {
+    pushRequestDescriptors(absl::nullopt);
     return Http::FilterHeadersStatus::Continue;
   }
 
@@ -89,6 +130,7 @@ Http::FilterHeadersStatus Filter::decodeHeaders(Http::RequestHeaderMap& headers,
   if (config->hasDescriptors()) {
     populateDescriptors(descriptors, headers);
   }
+  pushRequestDescriptors(descriptors);
 
   if (requestAllowed(descriptors)) {
     config->stats().ok_.inc();
@@ -115,13 +157,77 @@ Http::FilterHeadersStatus Filter::decodeHeaders(Http::RequestHeaderMap& headers,
   return Http::FilterHeadersStatus::StopIteration;
 }
 
+Http::FilterHeadersStatus Filter::encodeHeaders(Http::ResponseHeaderMap& headers, bool) {
+  const auto* config = getConfig();
+
+  if (config->enableXRateLimitHeaders()) {
+    auto descriptors = decoder_callbacks_->streamInfo()
+                           .filterState()
+                           ->getDataMutable<LocalRateLimitRequestDescriptorsQueue>(
+                               LocalRateLimitRequestDescriptorsQueue::key())
+                           .pop();
+    if (descriptors.has_value()) {
+      auto limit = maxTokens(descriptors.value());
+      auto remaining = remainingTokens(descriptors.value());
+      auto reset = remainingFillInterval(descriptors.value());
+
+      headers.addCopy(Http::LowerCaseString{"ratelimit-limit"}, limit);
+      headers.addCopy(Http::LowerCaseString{"ratelimit-remaining"}, remaining);
+      headers.addCopy(Http::LowerCaseString{"ratelimit-reset"}, reset);
+    }
+  }
+
+  return Http::FilterHeadersStatus::Continue;
+}
+
 bool Filter::requestAllowed(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) {
   const auto* config = getConfig();
   return config->rateLimitPerConnection()
              ? getPerConnectionRateLimiter().requestAllowed(request_descriptors)
              : config->requestAllowed(request_descriptors);
 }
 
+uint32_t Filter::maxTokens(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) {
+  const auto* config = getConfig();
+  return config->rateLimitPerConnection()
+             ? getPerConnectionRateLimiter().maxTokens(request_descriptors)
+             : config->maxTokens(request_descriptors);
+}
+
+uint32_t Filter::remainingTokens(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) {
+  const auto* config = getConfig();
+  return config->rateLimitPerConnection()
+             ? getPerConnectionRateLimiter().remainingTokens(request_descriptors)
+             : config->remainingTokens(request_descriptors);
+}
+
+uint32_t
+Filter::remainingFillInterval(absl::Span<const RateLimit::LocalDescriptor> request_descriptors) {
+  const auto* config = getConfig();
+  return config->rateLimitPerConnection()
+             ? getPerConnectionRateLimiter().remainingFillInterval(request_descriptors)
+             : config->remainingFillInterval(request_descriptors);
+}
+
+void Filter::pushRequestDescriptors(
+    absl::optional<std::vector<RateLimit::LocalDescriptor>> request_descriptors) {
+  if (!decoder_callbacks_->streamInfo()
+           .filterState()
+           ->hasData<LocalRateLimitRequestDescriptorsQueue>(
+               LocalRateLimitRequestDescriptorsQueue::key())) {
+    decoder_callbacks_->streamInfo().filterState()->setData(
+        LocalRateLimitRequestDescriptorsQueue::key(),
+        std::make_unique<LocalRateLimitRequestDescriptorsQueue>(),
+        StreamInfo::FilterState::StateType::Mutable, StreamInfo::FilterState::LifeSpan::Connection);
+  }
+
+  decoder_callbacks_->streamInfo()
+      .filterState()
+      ->getDataMutable<LocalRateLimitRequestDescriptorsQueue>(
+          LocalRateLimitRequestDescriptorsQueue::key())
+      .push(request_descriptors);
+}
+
 const Filters::Common::LocalRateLimit::LocalRateLimiterImpl& Filter::getPerConnectionRateLimiter() {
   const auto* config = getConfig();
   ASSERT(config->rateLimitPerConnection());