From 4a8ed8ab1b4f930427fc2e38d431b82bb4f35810 Mon Sep 17 00:00:00 2001 From: Vishal Rana Date: Sat, 13 Jun 2026 13:41:49 -0700 Subject: [PATCH 1/2] feat(middleware): optional RateLimiterStoreContext for response headers (#2961) Adds an optional RateLimiterStoreContext interface. When the configured store implements AllowContext(c, identifier), the rate limiter calls it instead of Allow, giving the store access to the request context so it can set response headers such as Retry-After / X-RateLimit-*. Fully backward compatible: stores implementing only Allow are unchanged. This is the optional-interface approach proposed by the maintainer in the issue thread; it does not alter the existing Allow interface or the built-in store. Co-Authored-By: Claude Opus 4.8 (1M context) --- middleware/rate_limiter.go | 17 +++++++- middleware/rate_limiter_context_test.go | 54 +++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) create mode 100644 middleware/rate_limiter_context_test.go diff --git a/middleware/rate_limiter.go b/middleware/rate_limiter.go index e80c51df1..091e9dea8 100644 --- a/middleware/rate_limiter.go +++ b/middleware/rate_limiter.go @@ -19,6 +19,14 @@ type RateLimiterStore interface { Allow(identifier string) (bool, error) } +// RateLimiterStoreContext is an optional interface a RateLimiterStore may implement. +// When the configured store implements it, the rate limiter calls AllowContext +// (with the request context) instead of Allow, allowing the store to set response +// headers such as Retry-After or X-RateLimit-* on the allow/deny decision. +type RateLimiterStoreContext interface { + AllowContext(c *echo.Context, identifier string) (bool, error) +} + // RateLimiterConfig defines the configuration for the rate limiter type RateLimiterConfig struct { Skipper Skipper @@ -136,7 +144,14 @@ func (config RateLimiterConfig) ToMiddleware() (echo.MiddlewareFunc, error) { return config.ErrorHandler(c, err) } - if allow, allowErr := config.Store.Allow(identifier); !allow { + var allow bool + var allowErr error + if sc, ok := config.Store.(RateLimiterStoreContext); ok { + allow, allowErr = sc.AllowContext(c, identifier) + } else { + allow, allowErr = config.Store.Allow(identifier) + } + if !allow { return config.DenyHandler(c, identifier, allowErr) } return next(c) diff --git a/middleware/rate_limiter_context_test.go b/middleware/rate_limiter_context_test.go new file mode 100644 index 000000000..929994faf --- /dev/null +++ b/middleware/rate_limiter_context_test.go @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: © 2015 LabStack LLC and Echo contributors + +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/labstack/echo/v5" + "github.com/stretchr/testify/assert" +) + +// ctxAwareStore implements both Allow and the optional AllowContext. AllowContext +// gives the store the request context so it can set response headers (e.g. +// Retry-After / X-RateLimit-*) — see #2961. +type ctxAwareStore struct { + allowCalled bool + ctxAllowCalled bool + allow bool +} + +func (s *ctxAwareStore) Allow(identifier string) (bool, error) { + s.allowCalled = true + return s.allow, nil +} + +func (s *ctxAwareStore) AllowContext(c *echo.Context, identifier string) (bool, error) { + s.ctxAllowCalled = true + c.Response().Header().Set("Retry-After", "42") + return s.allow, nil +} + +// When the store implements AllowContext, the middleware must call it instead of +// Allow, so the store can set rate-limit headers on the response. +func TestRateLimiter_storeAllowContextIsPreferred(t *testing.T) { + e := echo.New() + store := &ctxAwareStore{allow: true} + mw := RateLimiterWithConfig(RateLimiterConfig{ + Store: store, + IdentifierExtractor: func(c *echo.Context) (string, error) { return "id", nil }, + }) + handler := mw(func(c *echo.Context) error { return c.String(http.StatusOK, "ok") }) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + assert.NoError(t, handler(c)) + assert.True(t, store.ctxAllowCalled, "AllowContext should be called when implemented") + assert.False(t, store.allowCalled, "Allow should not be called when AllowContext is implemented") + assert.Equal(t, "42", rec.Header().Get("Retry-After"), "store should be able to set headers via the context") +} From 11fdbda073a48418a8b1de354b09c5007c16e867 Mon Sep 17 00:00:00 2001 From: Vishal Rana Date: Sat, 13 Jun 2026 21:43:13 -0700 Subject: [PATCH 2/2] feat(middleware): set X-RateLimit-* / Retry-After from built-in store (#2961) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements AllowContext on RateLimiterMemoryStore so the default store sets X-RateLimit-Limit, X-RateLimit-Remaining, and (on deny) Retry-After headers out of the box — mirroring the v4 PR #2985 by @leno23 on the v5 line. Allow() is refactored to share an internal allow() with AllowContext; the optional RateLimiterStoreContext interface (added earlier in this PR) routes the middleware to AllowContext when the store implements it. Co-Authored-By: Claude Opus 4.8 (1M context) --- middleware/rate_limiter.go | 45 +++++++++++++++++++++++-- middleware/rate_limiter_context_test.go | 35 +++++++++++++++++++ 2 files changed, 77 insertions(+), 3 deletions(-) diff --git a/middleware/rate_limiter.go b/middleware/rate_limiter.go index 091e9dea8..9756daf50 100644 --- a/middleware/rate_limiter.go +++ b/middleware/rate_limiter.go @@ -7,6 +7,7 @@ import ( "errors" "math" "net/http" + "strconv" "sync" "time" @@ -14,6 +15,12 @@ import ( "golang.org/x/time/rate" ) +// Rate limit response headers set by stores that implement RateLimiterStoreContext. +const ( + HeaderXRateLimitLimit = "X-RateLimit-Limit" + HeaderXRateLimitRemaining = "X-RateLimit-Remaining" +) + // RateLimiterStore is the interface to be implemented by custom stores. type RateLimiterStore interface { Allow(identifier string) (bool, error) @@ -247,7 +254,22 @@ var DefaultRateLimiterMemoryStoreConfig = RateLimiterMemoryStoreConfig{ // Allow implements RateLimiterStore.Allow func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) { + _, allowed := store.allow(identifier) + return allowed, nil +} + +// AllowContext implements RateLimiterStoreContext: it makes the allow/deny decision +// and sets the X-RateLimit-* (and Retry-After when denied) response headers. +func (store *RateLimiterMemoryStore) AllowContext(c *echo.Context, identifier string) (bool, error) { + limiter, allowed := store.allow(identifier) + store.setRateLimitHeaders(c, limiter, allowed) + return allowed, nil +} + +func (store *RateLimiterMemoryStore) allow(identifier string) (*rate.Limiter, bool) { store.mutex.Lock() + defer store.mutex.Unlock() + limiter, exists := store.visitors[identifier] if !exists { limiter = new(Visitor) @@ -259,9 +281,26 @@ func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) { if now.Sub(store.lastCleanup) > store.expiresIn { store.cleanupStaleVisitors(now) } - allowed := limiter.AllowN(now, 1) - store.mutex.Unlock() - return allowed, nil + return limiter.Limiter, limiter.AllowN(now, 1) +} + +func (store *RateLimiterMemoryStore) setRateLimitHeaders(c *echo.Context, limiter *rate.Limiter, allowed bool) { + header := c.Response().Header() + header.Set(HeaderXRateLimitLimit, strconv.Itoa(store.burst)) + + remaining := int(math.Floor(limiter.Tokens())) + if remaining < 0 { + remaining = 0 + } + header.Set(HeaderXRateLimitRemaining, strconv.Itoa(remaining)) + + if !allowed { + reservation := limiter.ReserveN(store.timeNow(), 1) + if delay := reservation.Delay(); delay > 0 { + header.Set(echo.HeaderRetryAfter, strconv.Itoa(int(math.Ceil(delay.Seconds())))) + } + reservation.Cancel() + } } /* diff --git a/middleware/rate_limiter_context_test.go b/middleware/rate_limiter_context_test.go index 929994faf..629c01e47 100644 --- a/middleware/rate_limiter_context_test.go +++ b/middleware/rate_limiter_context_test.go @@ -6,6 +6,7 @@ package middleware import ( "net/http" "net/http/httptest" + "strconv" "testing" "github.com/labstack/echo/v5" @@ -52,3 +53,37 @@ func TestRateLimiter_storeAllowContextIsPreferred(t *testing.T) { assert.False(t, store.allowCalled, "Allow should not be called when AllowContext is implemented") assert.Equal(t, "42", rec.Header().Get("Retry-After"), "store should be able to set headers via the context") } + +// The built-in memory store implements AllowContext, so it sets X-RateLimit-Limit / +// X-RateLimit-Remaining on every request and Retry-After when the limit is hit (#2961). +func TestRateLimiterMemoryStore_AllowContextSetsHeaders(t *testing.T) { + store := NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{Rate: 1, Burst: 3}) + e := echo.New() + e.GET("/", func(c *echo.Context) error { return c.String(http.StatusOK, "ok") }, + RateLimiterWithConfig(RateLimiterConfig{ + Store: store, + IdentifierExtractor: func(c *echo.Context) (string, error) { return "id", nil }, + })) + + do := func() *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodGet, "/", nil) + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + return rec + } + + // Burst of 3: each allowed request advertises the limit and decreasing remaining. + for i := 0; i < 3; i++ { + rec := do() + assert.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, "3", rec.Header().Get(HeaderXRateLimitLimit)) + assert.Equal(t, strconv.Itoa(2-i), rec.Header().Get(HeaderXRateLimitRemaining)) + assert.Empty(t, rec.Header().Get(echo.HeaderRetryAfter)) + } + + // 4th request is denied: 429, remaining 0, and a Retry-After hint. + rec := do() + assert.Equal(t, http.StatusTooManyRequests, rec.Code) + assert.Equal(t, "0", rec.Header().Get(HeaderXRateLimitRemaining)) + assert.NotEmpty(t, rec.Header().Get(echo.HeaderRetryAfter)) +}