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

@@ -8,6 +8,7 @@ namespace Envoy {
 namespace {
 // The minimal fill rate will be one second every year.
 constexpr double kMinFillRate = 1.0 / (365 * 24 * 60 * 60);
+constexpr auto MILLISECONDS_PER_SECOND = 1000;
 
 } // namespace
 
@@ -59,13 +60,16 @@ void TokenBucketImpl::maybeReset(uint64_t num_tokens) {
 }
 
 AtomicTokenBucketImpl::AtomicTokenBucketImpl(uint64_t max_tokens, TimeSource& time_source,
+                                             std::chrono::milliseconds fill_interval,
                                              double fill_rate, bool init_fill)
-    : AtomicTokenBucketImpl::AtomicTokenBucketImpl(max_tokens, time_source, fill_rate,
-                                                   (init_fill) ? max_tokens : 0) {}
+    : AtomicTokenBucketImpl::AtomicTokenBucketImpl(max_tokens, time_source, fill_interval,
+                                                   fill_rate, (init_fill) ? max_tokens : 0) {}
 
 AtomicTokenBucketImpl::AtomicTokenBucketImpl(uint64_t max_tokens, TimeSource& time_source,
+                                             std::chrono::milliseconds fill_interval,
                                              double fill_rate, uint64_t initial_tokens)
     : max_tokens_(max_tokens), fill_rate_(std::max(std::abs(fill_rate), kMinFillRate)),
+      fill_interval_(std::chrono::duration<double>(fill_interval).count()),
       time_source_(time_source) {
   auto time_in_seconds = timeNowInSeconds();
   if (initial_tokens) {
@@ -103,4 +107,18 @@ 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);
+  }
+
+  // Calculate time since the last fill
+  double current_time = timeNowInSeconds();
+  double time_since_last_fill = std::fmod(current_time, fill_interval_);
+  double time_until_next_fill = fill_interval_ - time_since_last_fill;
+  return std::chrono::milliseconds(
+      static_cast<uint64_t>(time_until_next_fill * MILLISECONDS_PER_SECOND));
+}
+
 } // namespace Envoy

source/common/common/token_bucket_impl.h

@@ -48,8 +48,10 @@ class AtomicTokenBucketImpl {
    * @param init_fill supplies whether the bucket should be initialized with max_tokens.
    */
   explicit AtomicTokenBucketImpl(uint64_t max_tokens, TimeSource& time_source,
-                                 double fill_rate = 1.0, bool init_fill = true);
-  explicit AtomicTokenBucketImpl(uint64_t max_tokens, TimeSource& time_source, double fill_rate,
+                                 std::chrono::milliseconds fill_interval, double fill_rate = 1.0,
+                                 bool init_fill = true);
+  explicit AtomicTokenBucketImpl(uint64_t max_tokens, TimeSource& time_source,
+                                 std::chrono::milliseconds fill_interval, double fill_rate,
                                  uint64_t initial_tokens);
 
   // This reference https://github.com/facebook/folly/blob/main/folly/TokenBucket.h.
@@ -115,13 +117,20 @@ 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;
 
   const double max_tokens_;
   const double fill_rate_;
+  const double fill_interval_;
 
-  std::atomic<double> time_in_seconds_{};
+  std::atomic<double> time_in_seconds_;
   TimeSource& time_source_;
 };
 

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

@@ -79,7 +79,7 @@ ShareProviderManagerSharedPtr ShareProviderManager::singleton(Event::Dispatcher&
 RateLimitTokenBucket::RateLimitTokenBucket(uint64_t max_tokens, uint64_t tokens_per_fill,
                                            std::chrono::milliseconds fill_interval,
                                            TimeSource& time_source)
-    : token_bucket_(max_tokens, time_source,
+    : token_bucket_(max_tokens, time_source, fill_interval,
                     // Calculate the fill rate in tokens per second.
                     tokens_per_fill / std::chrono::duration<double>(fill_interval).count()),
       fill_interval_(fill_interval) {}

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>(std::ceil(token_bucket_.nextTokenAvailable().count() / 1000));
+  }
 
 private:
   AtomicTokenBucketImpl token_bucket_;

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

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

source/extensions/filters/http/rate_limit_quota/global_client_impl.cc

@@ -251,9 +251,9 @@ createTokenBucketFromAction(const RateLimitStrategy& strategy, TimeSource& time_
                                 ? max_tokens * (existing_token_bucket->remainingTokens() /
                                                 existing_token_bucket->maxTokens())
                                 : max_tokens;
-
-  return std::make_shared<AtomicTokenBucketImpl>(max_tokens, time_source, fill_rate_per_sec,
-                                                 initial_tokens);
+  return std::make_shared<AtomicTokenBucketImpl>(
+      max_tokens, time_source, std::chrono::milliseconds(fill_interval_sec * 1000),
+      fill_rate_per_sec, initial_tokens);
 }
 
 void GlobalRateLimitClientImpl::onReceiveMessage(RateLimitQuotaResponsePtr&& response) {

test/common/common/token_bucket_impl_test.cc

@@ -134,20 +134,22 @@ class AtomicTokenBucketImplTest : public testing::Test {
 
 // Verifies TokenBucket initialization.
 TEST_F(AtomicTokenBucketImplTest, Initialization) {
-  AtomicTokenBucketImpl token_bucket{1, time_system_, -1.0};
+  AtomicTokenBucketImpl token_bucket{1, time_system_, std::chrono::seconds(1), -1.0};
 
   EXPECT_EQ(1, token_bucket.fillRate());
   EXPECT_EQ(1, token_bucket.maxTokens());
   EXPECT_EQ(1, token_bucket.remainingTokens());
+  EXPECT_EQ(token_bucket.nextTokenAvailable(), std::chrono::milliseconds(0));
 
   EXPECT_EQ(1, token_bucket.consume(1, false));
   EXPECT_EQ(0, token_bucket.consume(1, false));
   EXPECT_EQ(false, token_bucket.consume());
+  EXPECT_EQ(token_bucket.nextTokenAvailable(), std::chrono::seconds(1));
 }
 
 // Verifies TokenBucket's maximum capacity.
 TEST_F(AtomicTokenBucketImplTest, MaxBucketSize) {
-  AtomicTokenBucketImpl token_bucket{3, time_system_, 1};
+  AtomicTokenBucketImpl token_bucket{3, time_system_, std::chrono::seconds(1), 1};
 
   EXPECT_EQ(1, token_bucket.fillRate());
   EXPECT_EQ(3, token_bucket.maxTokens());
@@ -161,7 +163,7 @@ TEST_F(AtomicTokenBucketImplTest, MaxBucketSize) {
 
 // Verifies that TokenBucket can consume tokens.
 TEST_F(AtomicTokenBucketImplTest, Consume) {
-  AtomicTokenBucketImpl token_bucket{10, time_system_, 1};
+  AtomicTokenBucketImpl token_bucket{10, time_system_, std::chrono::seconds(1), 1};
 
   EXPECT_EQ(0, token_bucket.consume(20, false));
   EXPECT_EQ(9, token_bucket.consume(9, false));
@@ -182,7 +184,7 @@ TEST_F(AtomicTokenBucketImplTest, Consume) {
 
 // Verifies that TokenBucket can refill tokens.
 TEST_F(AtomicTokenBucketImplTest, Refill) {
-  AtomicTokenBucketImpl token_bucket{1, time_system_, 0.5};
+  AtomicTokenBucketImpl token_bucket{1, time_system_, std::chrono::seconds(1), 0.5};
   EXPECT_EQ(1, token_bucket.consume(1, false));
 
   time_system_.setMonotonicTime(std::chrono::milliseconds(500));
@@ -195,7 +197,7 @@ TEST_F(AtomicTokenBucketImplTest, Refill) {
 
 // Test partial consumption of tokens.
 TEST_F(AtomicTokenBucketImplTest, PartialConsumption) {
-  AtomicTokenBucketImpl token_bucket{16, time_system_, 16};
+  AtomicTokenBucketImpl token_bucket{16, time_system_, std::chrono::seconds(1), 16};
   EXPECT_EQ(16, token_bucket.consume(18, true));
   time_system_.advanceTimeWait(std::chrono::milliseconds(62));
   EXPECT_EQ(0, token_bucket.consume(1, true));
@@ -207,7 +209,8 @@ TEST_F(AtomicTokenBucketImplTest, PartialConsumption) {
 TEST_F(AtomicTokenBucketImplTest, YearlyMinRefillRate) {
   constexpr uint64_t seconds_per_year = 365 * 24 * 60 * 60;
   // Set the fill rate to be 2 years.
-  AtomicTokenBucketImpl token_bucket{1, time_system_, 1.0 / (seconds_per_year * 2)};
+  AtomicTokenBucketImpl token_bucket{1, time_system_, std::chrono::seconds(seconds_per_year),
+                                     1.0 / (seconds_per_year * 2)};
 
   // Consume first token.
   EXPECT_EQ(1, token_bucket.consume(1, false));
@@ -220,7 +223,7 @@ TEST_F(AtomicTokenBucketImplTest, YearlyMinRefillRate) {
 }
 
 TEST_F(AtomicTokenBucketImplTest, ConsumeNegativeTokens) {
-  AtomicTokenBucketImpl token_bucket{10, time_system_, 1};
+  AtomicTokenBucketImpl token_bucket{10, time_system_, std::chrono::seconds(1), 1};
 
   EXPECT_EQ(3, token_bucket.consume([](double) { return 3; }));
   EXPECT_EQ(7, token_bucket.remainingTokens());
@@ -229,7 +232,7 @@ TEST_F(AtomicTokenBucketImplTest, ConsumeNegativeTokens) {
 }
 
 TEST_F(AtomicTokenBucketImplTest, ConsumeSuperLargeTokens) {
-  AtomicTokenBucketImpl token_bucket{10, time_system_, 1};
+  AtomicTokenBucketImpl token_bucket{10, time_system_, std::chrono::seconds(1), 1};
 
   EXPECT_EQ(100, token_bucket.consume([](double) { return 100; }));
   EXPECT_EQ(-90, token_bucket.remainingTokens());
@@ -239,7 +242,7 @@ TEST_F(AtomicTokenBucketImplTest, MultipleThreadsConsume) {
   // Real time source to ensure we will not fall into endless loop.
   Event::TestRealTimeSystem real_time_source;
 
-  AtomicTokenBucketImpl token_bucket{1200, time_system_, 1.0};
+  AtomicTokenBucketImpl token_bucket{1200, time_system_, std::chrono::seconds(1), 1.0};
 
   // Exhaust all tokens.
   EXPECT_EQ(1200, token_bucket.consume(1200, false));

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

@@ -1,4 +1,7 @@
+#include <chrono>
+
 #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"
@@ -159,6 +162,29 @@ class LocalRateLimitFilterIntegrationTest : public Event::TestUsingSimulatedTime
     EXPECT_EQ(expected_status, response->headers().getStatusValue());
     EXPECT_EQ(expected_body_size, response->body().size());
   }
+  void verifyResponse(IntegrationStreamDecoderPtr response, const std::string& expected_status,
+                      size_t expected_body_size, const std::string& expected_limit,
+                      const std::string& expected_remaining, const std::string& expected_reset) {
+    ASSERT_TRUE(response->waitForEndStream());
+    EXPECT_TRUE(response->complete());
+    EXPECT_EQ(expected_status, response->headers().getStatusValue());
+    EXPECT_EQ(expected_body_size, response->body().size());
+    EXPECT_THAT(
+        response->headers(),
+        Http::HeaderValueOf(
+            Extensions::HttpFilters::Common::RateLimit::XRateLimitHeaders::get().XRateLimitLimit,
+            expected_limit));
+    EXPECT_THAT(
+        response->headers(),
+        Http::HeaderValueOf(Extensions::HttpFilters::Common::RateLimit::XRateLimitHeaders::get()
+                                .XRateLimitRemaining,
+                            expected_remaining));
+    EXPECT_THAT(
+        response->headers(),
+        Http::HeaderValueOf(
+            Extensions::HttpFilters::Common::RateLimit::XRateLimitHeaders::get().XRateLimitReset,
+            expected_reset));
+  }
 
   void sendAndVerifyRequest(const std::string& cluster, const std::string& expected_status,
                             size_t expected_body_size) {
@@ -169,11 +195,29 @@ class LocalRateLimitFilterIntegrationTest : public Event::TestUsingSimulatedTime
     EXPECT_TRUE(upstream_request_->complete());
     EXPECT_EQ(0U, upstream_request_->bodyLength());
   }
+  void sendAndVerifyRequest(const std::string& expected_limit,
+                            const std::string& expected_remaining,
+                            const std::string& expected_reset) {
+    auto response = codec_client_->makeRequestWithBody(default_request_headers_, 0);
+    waitForNextUpstreamRequest();
+    upstream_request_->encodeHeaders(default_response_headers_, 1);
+    verifyResponse(std::move(response), "200", 0, expected_limit, expected_remaining,
+                   expected_reset);
+    EXPECT_TRUE(upstream_request_->complete());
+    EXPECT_EQ(0U, upstream_request_->bodyLength());
+  }
   void sendRateLimitedRequest(const std::string& cluster) {
     auto response = makeRequest(cluster);
     verifyResponse(std::move(response), "429",
                    18); // 18 is the expected body size for rate-limited responses.
   }
+  void sendRateLimitedRequest(const std::string& expected_limit,
+                              const std::string& expected_remaining,
+                              const std::string& expected_reset) {
+    auto response = codec_client_->makeRequestWithBody(default_request_headers_, 0);
+    verifyResponse(std::move(response), "429", 18, expected_limit, expected_remaining,
+                   expected_reset);
+  }
 
   static constexpr absl::string_view filter_config_ =
       R"EOF(
@@ -203,6 +247,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: 2
+    tokens_per_fill: 2
+    fill_interval: 4s
+  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 +562,30 @@ TEST_P(LocalRateLimitFilterIntegrationTest, DenyRequestWithinSameConnection) {
   EXPECT_EQ(18, response->body().size());
 }
 
+TEST_P(LocalRateLimitFilterIntegrationTest, HeaderTest) {
+  initializeFilter(fmt::format(limit_header_filter_config_, "false"));
+
+  // The first request should be allowed and the response should contain
+  codec_client_ = makeHttpConnection(lookupPort("http"));
+  sendAndVerifyRequest("2", "1", "0");
+  cleanupUpstreamAndDownstream();
+
+  // Max tokens is 2, the second request should be allowed.
+  codec_client_ = makeHttpConnection(lookupPort("http"));
+  sendAndVerifyRequest("2", "0", "4");
+  cleanupUpstreamAndDownstream();
+
+  // The third request should be rate limited, x-ratelimit-reset should be 4s.
+  codec_client_ = makeHttpConnection(lookupPort("http"));
+  sendRateLimitedRequest("2", "0", "4");
+  cleanupUpstreamAndDownstream();
+
+  // After 1s, the forth request should be rate limited, x-ratelimit-reset should be 3s.
+  simTime().advanceTimeWait(std::chrono::seconds(1));
+  codec_client_ = makeHttpConnection(lookupPort("http"));
+  sendRateLimitedRequest("2", "0", "3");
+}
+
 TEST_P(LocalRateLimitFilterIntegrationTest, PermitRequestAcrossDifferentConnections) {
   initializeFilter(fmt::format(filter_config_, "true"));
 
@@ -562,7 +659,6 @@ TEST_P(LocalRateLimitFilterIntegrationTest, BasicTestPerRouteAndRds) {
   EXPECT_TRUE(response->complete());
   EXPECT_EQ("200", response->headers().getStatusValue());
   EXPECT_EQ(0, response->body().size());
-
   cleanupUpstreamAndDownstream();
 
   cleanUpXdsConnection();

test/extensions/filters/http/rate_limit_quota/filter_test.cc

@@ -598,7 +598,8 @@ TEST_F(FilterTest, DecodeHeaderWithTokenBucketAllow) {
   token_bucket->mutable_fill_interval()->set_seconds(60);
   // 100 available tokens so the test doesn't get throttled.
   std::shared_ptr<AtomicTokenBucketImpl> token_bucket_limiter =
-      std::make_shared<AtomicTokenBucketImpl>(100, dispatcher_.timeSource(), 100 / 60);
+      std::make_shared<AtomicTokenBucketImpl>(100, dispatcher_.timeSource(),
+                                              std::chrono::seconds(60), 100 / 60);
 
   RateLimitQuotaResponse::BucketAction no_assignment_action;
   no_assignment_action.mutable_quota_assignment_action()
@@ -643,7 +644,8 @@ TEST_F(FilterTest, DecodeHeaderWithTokenBucketDeny) {
   token_bucket->mutable_tokens_per_fill()->set_value(1);
   token_bucket->mutable_fill_interval()->set_seconds(60);
   std::shared_ptr<AtomicTokenBucketImpl> token_bucket_limiter =
-      std::make_shared<AtomicTokenBucketImpl>(1, dispatcher_.timeSource(), 1 / 60);
+      std::make_shared<AtomicTokenBucketImpl>(1, dispatcher_.timeSource(), std::chrono::seconds(60),
+                                              1 / 60);
   // All subsequent requests should deny for 60 (mock) seconds.
   EXPECT_TRUE(token_bucket_limiter->consume());