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>

Author: vishr

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)

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