local_ratelimit: support return x-ratelimit-reset header

changelogs/current.yaml

@@ -149,5 +149,8 @@ new_features:
     added :ref:`enable_ja4_fingerprinting
     <envoy_v3_api_field_extensions.filters.listener.tls_inspector.v3.TlsInspector.enable_ja4_fingerprinting>` to create
     a JA4 fingerprint hash from the Client Hello message.
+- area: local_ratelimit
+  change: |
+    ``local_ratelimit`` will return ``x-ratelimit-reset`` header when the rate limit is exceeded.
 
 deprecated:

source/common/common/token_bucket_impl.cc

@@ -103,4 +103,13 @@ double AtomicTokenBucketImpl::timeNowInSeconds() const {
   return std::chrono::duration<double>(time_source_.monotonicTime().time_since_epoch()).count();
 }
 
+std::chrono::milliseconds AtomicTokenBucketImpl::nextTokenAvailable() const {
+  // If there are tokens available, return immediately.
+  if (remainingTokens() >= 1) {
+    return std::chrono::milliseconds(0);
+  }
+
+  return std::chrono::milliseconds(static_cast<uint64_t>(std::ceil((1 / fill_rate_) * 1000)));
+}
+
 } // namespace Envoy

source/common/common/token_bucket_impl.h

@@ -115,6 +115,12 @@ class AtomicTokenBucketImpl {
    */
   double remainingTokens() const;
 
+  /**
+   * Get the time to next token available. This is a snapshot and may change after the call.
+   * @return the time to next token available.
+   */
+  std::chrono::milliseconds nextTokenAvailable() const;
+
 private:
   double timeNowInSeconds() const;
 

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

@@ -105,6 +105,7 @@ class TokenBucketContext {
 
   virtual uint64_t maxTokens() const PURE;
   virtual uint64_t remainingTokens() const PURE;
+  virtual uint64_t resetSeconds() const PURE;
 };
 
 class RateLimitTokenBucket : public TokenBucketContext,
@@ -122,6 +123,9 @@ class RateLimitTokenBucket : public TokenBucketContext,
   uint64_t remainingTokens() const override {
     return static_cast<uint64_t>(token_bucket_.remainingTokens());
   }
+  uint64_t resetSeconds() const override {
+    return static_cast<uint64_t>(token_bucket_.nextTokenAvailable().count() / 1000);
+  }
 
 private:
   AtomicTokenBucketImpl token_bucket_;

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

@@ -206,6 +206,11 @@ Http::FilterHeadersStatus Filter::encodeHeaders(Http::ResponseHeaderMap& headers
     headers.addReferenceKey(
         HttpFilters::Common::RateLimit::XRateLimitHeaders::get().XRateLimitRemaining,
         token_bucket_context_->remainingTokens());
+    if (token_bucket_context_->remainingTokens() == 0) {
+      headers.addReferenceKey(
+          HttpFilters::Common::RateLimit::XRateLimitHeaders::get().XRateLimitReset,
+          token_bucket_context_->resetSeconds());
+    }
   }
 
   return Http::FilterHeadersStatus::Continue;

test/extensions/filters/http/local_ratelimit/local_ratelimit_integration_test.cc

@@ -1,4 +1,5 @@
 #include "source/extensions/filters/common/local_ratelimit/local_ratelimit_impl.h"
+#include "source/extensions/filters/http/common/ratelimit_headers.h"
 
 #include "test/integration/http_protocol_integration.h"
 #include "test/test_common/test_runtime.h"
@@ -203,6 +204,35 @@ name: envoy.filters.http.local_ratelimit
   local_rate_limit_per_downstream_connection: {}
 )EOF";
 
+  static constexpr absl::string_view limit_header_filter_config_ =
+      R"EOF(
+name: envoy.filters.http.local_ratelimit
+typed_config:
+  "@type": type.googleapis.com/envoy.extensions.filters.http.local_ratelimit.v3.LocalRateLimit
+  stat_prefix: http_local_rate_limiter
+  enableXRatelimitHeaders: DRAFT_VERSION_03
+  token_bucket:
+    max_tokens: 1
+    tokens_per_fill: 1
+    fill_interval: 1000s
+  filter_enabled:
+    runtime_key: local_rate_limit_enabled
+    default_value:
+      numerator: 100
+      denominator: HUNDRED
+  filter_enforced:
+    runtime_key: local_rate_limit_enforced
+    default_value:
+      numerator: 100
+      denominator: HUNDRED
+  response_headers_to_add:
+    - append_action: OVERWRITE_IF_EXISTS_OR_ADD
+      header:
+        key: x-local-rate-limit
+        value: 'true'
+  local_rate_limit_per_downstream_connection: {}
+)EOF";
+
   static constexpr absl::string_view filter_config_with_blank_value_descriptor_ =
       R"EOF(
 name: envoy.filters.http.local_ratelimit
@@ -489,6 +519,46 @@ TEST_P(LocalRateLimitFilterIntegrationTest, DenyRequestWithinSameConnection) {
   EXPECT_EQ(18, response->body().size());
 }
 
+TEST_P(LocalRateLimitFilterIntegrationTest, LimitHeaderTest) {
+  initializeFilter(fmt::format(limit_header_filter_config_, "true"));
+
+  codec_client_ = makeHttpConnection(lookupPort("http"));
+  auto response = codec_client_->makeRequestWithBody(default_request_headers_, 0);
+
+  waitForNextUpstreamRequest();
+  upstream_request_->encodeHeaders(default_response_headers_, 1);
+
+  ASSERT_TRUE(response->waitForEndStream());
+
+  EXPECT_TRUE(upstream_request_->complete());
+  EXPECT_EQ(0U, upstream_request_->bodyLength());
+  EXPECT_TRUE(response->complete());
+  EXPECT_EQ("200", response->headers().getStatusValue());
+  EXPECT_EQ(0, response->body().size());
+  EXPECT_THAT(
+      response->headers(),
+      Http::HeaderValueOf(
+          Extensions::HttpFilters::Common::RateLimit::XRateLimitHeaders::get().XRateLimitLimit,
+          "1"));
+
+  response = codec_client_->makeRequestWithBody(default_request_headers_, 0);
+
+  ASSERT_TRUE(response->waitForEndStream());
+  EXPECT_TRUE(response->complete());
+  EXPECT_EQ("429", response->headers().getStatusValue());
+  EXPECT_EQ(18, response->body().size());
+  EXPECT_THAT(
+      response->headers(),
+      Http::HeaderValueOf(
+          Extensions::HttpFilters::Common::RateLimit::XRateLimitHeaders::get().XRateLimitRemaining,
+          "0"));
+  EXPECT_THAT(
+      response->headers(),
+      Http::HeaderValueOf(
+          Extensions::HttpFilters::Common::RateLimit::XRateLimitHeaders::get().XRateLimitReset,
+          "1000"));
+}
+
 TEST_P(LocalRateLimitFilterIntegrationTest, PermitRequestAcrossDifferentConnections) {
   initializeFilter(fmt::format(filter_config_, "true"));
 
@@ -562,7 +632,6 @@ TEST_P(LocalRateLimitFilterIntegrationTest, BasicTestPerRouteAndRds) {
   EXPECT_TRUE(response->complete());
   EXPECT_EQ("200", response->headers().getStatusValue());
   EXPECT_EQ(0, response->body().size());
-
   cleanupUpstreamAndDownstream();
 
   cleanUpXdsConnection();