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>

Author: vishr

middleware/rate_limiter.go

@@ -7,13 +7,20 @@ 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)
@@ -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()
+	}
 }
 
 /*

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))
+}