feat(middleware): optional RateLimiterStoreContext for response headers (#2961) (#3007)

* 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) <noreply@anthropic.com>

* feat(middleware): set X-RateLimit-* / Retry-After from built-in store (#2961)

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) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: vishr

middleware/rate_limiter.go

@@ -7,18 +7,33 @@ import (
 	"errors"
 	"math"
 	"net/http"
+	"strconv"
 	"sync"
 	"time"
 
 	"github.com/labstack/echo/v5"
 	"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)
 }
 
+// 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 +151,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)
@@ -232,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)
@@ -244,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()
+	}
 }
 
 /*

middleware/rate_limiter_context_test.go

@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: MIT
+// SPDX-FileCopyrightText: © 2015 LabStack LLC and Echo contributors
+
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"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")
+}
+
+// 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))
+}