TT-6075: Update rate limit implementation

cli/linter/schema.json

@@ -1610,6 +1610,12 @@
         }
       }
     },
+    "rate_limit_response_headers": {
+      "description": "defines the source for the rate limit headers",
+      "type": ["string"],
+      "enum": ["", "quotas", "rate_limits"],
+      "default": "quotas"
+    },
     "allow_unsafe_policy_ids": {
       "type": ["boolean", "null"],
       "additionalProperties": false

config/rate_limit.go

@@ -37,8 +37,20 @@ type RateLimit struct {
 
 	// Controls which algorthm to use as a fallback when your distributed rate limiter can't be used.
 	DRLEnableSentinelRateLimiter bool `json:"drl_enable_sentinel_rate_limiter"`
+
+	// RateLimitResponseHeaders specifies the data source for rate limit headers in HTTP responses.
+	// This controls whether rate limit headers (X-RateLimit-Limit, X-RateLimit-Remaining, etc.)
+	// are populated from quota data or rate limit data. Valid values: "quotas", "rate_limits".
+	RateLimitResponseHeaders RateLimitSource `json:"rate_limit_response_headers"`
 }
 
+type RateLimitSource string
+
+const (
+	SourceQuotas     RateLimitSource = "quotas"
+	SourceRateLimits RateLimitSource = "rate_limits"
+)
+
 // String returns a readable setting for the rate limiter in effect.
 func (r *RateLimit) String() string {
 	info := "using transactions"

gateway/middleware_test.go

@@ -420,7 +420,7 @@ func TestSessionLimiter_RedisQuotaExceeded_ExpiredAtReset(t *testing.T) {
 		}
 
 		beforeTime := time.Now()
-		blocked := limiter.RedisQuotaExceeded(req, session, quotaKey, "", limit, g.Gw.GlobalSessionManager.Store(), false)
+		blocked := limiter.RedisQuotaExceeded(req, session, quotaKey, "", limit, false, false)
 		afterTime := time.Now()
 
 		assert.Equal(t, quotaMax-1, session.QuotaRemaining, "Quota remaining should be quotaMax - 1 after increment")
@@ -475,7 +475,7 @@ func TestSessionLimiter_RedisQuotaExceeded_ExpiredAtReset(t *testing.T) {
 		}
 
 		beforeTime := time.Now()
-		blocked := limiter.RedisQuotaExceeded(req, session, quotaKey, scope, limit, g.Gw.GlobalSessionManager.Store(), false)
+		blocked := limiter.RedisQuotaExceeded(req, session, quotaKey, scope, limit, false, false)
 		afterTime := time.Now()
 
 		accessDef := session.AccessRights["api1"]

gateway/model.go

@@ -12,8 +12,20 @@ import (
 
 type EventMetaDefault = model.EventMetaDefault
 
+type CtxData = map[string]any
+
+const (
+	ctxDataKeyRateLimitLimit     = "rate_limit_limit"
+	ctxDataKeyRateLimitRemaining = "rate_limit_remaining"
+	ctxDataKeyRateLimitReset     = "rate_limit_reset"
+
+	ctxDataKeyQuotaLimit     = "quota_limit"
+	ctxDataKeyQuotaRemaining = "quota_remaining"
+	ctxDataKeyQuotaReset     = "quota_reset"
+)
+
 var (
-	ctxData = httpctx.NewValue[map[string]any](ctx.ContextData)
+	ctxData = httpctx.NewValue[CtxData](ctx.ContextData)
 
 	ctxGetData = ctxData.Get
 	ctxSetData = ctxData.Set
@@ -28,3 +40,14 @@ var (
 
 	EncodeRequestToEvent = event.EncodeRequestToEvent
 )
+
+func ctxGetOrCreateData(r *http.Request) CtxData {
+	data := ctxGetData(r)
+
+	if data == nil {
+		data = CtxData{}
+		ctxSetData(r, data)
+	}
+
+	return data
+}

gateway/model_apispec.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"net/http"
 	"net/url"
-	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -16,15 +15,13 @@ import (
 	"github.com/TykTechnologies/tyk/apidef/oas"
 	"github.com/TykTechnologies/tyk/config"
 	"github.com/TykTechnologies/tyk/ctx"
-	"github.com/TykTechnologies/tyk/header"
 	"github.com/TykTechnologies/tyk/internal/agentprotocol"
 	"github.com/TykTechnologies/tyk/internal/certcheck"
 	"github.com/TykTechnologies/tyk/internal/errors"
 	"github.com/TykTechnologies/tyk/internal/graphengine"
 	"github.com/TykTechnologies/tyk/internal/httpctx"
 	"github.com/TykTechnologies/tyk/internal/httputil"
 	"github.com/TykTechnologies/tyk/internal/jsonrpc"
-	"github.com/TykTechnologies/tyk/user"
 
 	_ "github.com/TykTechnologies/tyk/internal/mcp" // registers MCP VEM prefixes
 )
@@ -340,19 +337,3 @@ func (a *APISpec) APIType() string {
 		return "classic"
 	}
 }
-
-func (a *APISpec) sendRateLimitHeaders(session *user.SessionState, dest *http.Response) {
-	quotaMax, quotaRemaining, quotaRenews := int64(0), int64(0), int64(0)
-
-	if session != nil {
-		quotaMax, quotaRemaining, _, quotaRenews = session.GetQuotaLimitByAPIID(a.APIID)
-	}
-
-	if dest.Header == nil {
-		dest.Header = http.Header{}
-	}
-
-	dest.Header.Set(header.XRateLimitLimit, strconv.Itoa(int(quotaMax)))
-	dest.Header.Set(header.XRateLimitRemaining, strconv.Itoa(int(quotaRemaining)))
-	dest.Header.Set(header.XRateLimitReset, strconv.Itoa(int(quotaRenews)))
-}

gateway/model_test.go

@@ -0,0 +1,32 @@
+package gateway
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_ctxGetOrCreateData(t *testing.T) {
+	t.Run("returns data if already exists", func(t *testing.T) {
+		req, err := http.NewRequestWithContext(t.Context(), "GET", "/", nil)
+		require.NoError(t, err)
+
+		ctxSetData(req, CtxData{"hello": "world0"})
+
+		assert.Equal(t, CtxData{"hello": "world0"}, ctxGetOrCreateData(req))
+	})
+
+	t.Run("create new data if not exists", func(t *testing.T) {
+		req, err := http.NewRequestWithContext(t.Context(), "GET", "/", nil)
+		require.NoError(t, err)
+
+		data1 := ctxGetOrCreateData(req)
+		data1["hello"] = "world1"
+
+		data2 := ctxGetOrCreateData(req)
+
+		assert.Equal(t, data1, data2)
+	})
+}

gateway/mw_api_rate_limit.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/TykTechnologies/tyk/ctx"
+	"github.com/TykTechnologies/tyk/header"
 	tykerrors "github.com/TykTechnologies/tyk/internal/errors"
 	"github.com/TykTechnologies/tyk/internal/event"
 	"github.com/TykTechnologies/tyk/storage"
@@ -94,30 +95,38 @@ func (k *RateLimitForAPI) EnabledForSpec() bool {
 }
 
 // ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
-func (k *RateLimitForAPI) ProcessRequest(_ http.ResponseWriter, r *http.Request, _ interface{}) (error, int) {
+//
+//nolint:staticcheck
+func (k *RateLimitForAPI) ProcessRequest(rw http.ResponseWriter, r *http.Request, _ interface{}) (error, int) {
 	// Skip rate limiting and quotas for looping
 	if !ctxCheckLimits(r) {
 		return nil, http.StatusOK
 	}
 
-	storeRef := k.Gw.GlobalSessionManager.Store()
+	session := k.getSession(r)
+
+	limitHeaderSender := k.Gw.limitHeaderFactory(rw.Header())
+	// Only inject API-level rate limit headers if personal rate limit headers
+	// haven't already been injected by RateLimitAndQuotaCheck.
+	if rw.Header().Get(header.XRateLimitLimit) != "" {
+		limitHeaderSender = nil
+	}
 
 	reason := k.Gw.SessionLimiter.ForwardMessage(
 		r,
-		k.getSession(r),
+		session,
 		k.keyName,
 		k.quotaKey,
-		storeRef,
 		true,
 		false,
 		k.Spec,
 		false,
+		limitHeaderSender,
 	)
 
 	k.emitRateLimitEvents(r, k.keyName)
 
 	if reason == sessionFailRateLimit {
-		// Set error classification for access logs
 		ctx.SetErrorClassification(r, tykerrors.ClassifyRateLimitError(tykerrors.ErrTypeAPIRateLimit, k.Name()))
 		return k.handleRateLimitFailure(r, event.RateLimitExceeded, "API Rate Limit Exceeded", k.keyName)
 	}

gateway/mw_api_rate_limit_test.go

@@ -1,13 +1,16 @@
 package gateway
 
 import (
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"testing"
 	"time"
 
 	"github.com/TykTechnologies/tyk/apidef"
+	"github.com/TykTechnologies/tyk/config"
+	"github.com/TykTechnologies/tyk/header"
 
 	"github.com/stretchr/testify/assert"
 
@@ -75,6 +78,82 @@
 	assert.False(t, rlDisabled.EnabledForSpec())
 }
 
+func TestAPIRateLimitResponseHeaders(t *testing.T) {
+	limiters := []string{"Redis", "Sentinel", "DRL", "FixedWindow"}
+
+	for _, limiter := range limiters {
+		t.Run("API Rate limit headers for "+limiter, func(t *testing.T) {
+			ts := StartTest(func(globalConf *config.Config) {
+				globalConf.RateLimitResponseHeaders = config.SourceRateLimits
+
+				switch limiter {
+				case "Redis":
+					globalConf.EnableRedisRollingLimiter = true
+				case "Sentinel":
+					globalConf.EnableSentinelRateLimiter = true
+				case "DRL":
+					globalConf.DRLEnableSentinelRateLimiter = true
+				case "FixedWindow":
+					globalConf.EnableFixedWindowRateLimiter = true
+				}
+			})
+			defer ts.Close()
+
+			var (
+				rateLimitRate float64 = 2
+				rateLimitPer  float64 = 10
+			)
+
+			_ = ts.Gw.BuildAndLoadAPI(func(spec *APISpec) {
+				spec.APIID = "api-rate-limit-headers-test-" + limiter
+				spec.Proxy.ListenPath = "/api-rate-limit-headers-test"
+				spec.UseKeylessAccess = true
+				spec.GlobalRateLimit = apidef.GlobalRateLimit{
+					Disabled: false,
+					Rate:     rateLimitRate,
+					Per:      rateLimitPer,
+				}
+			})[0]
+
+			expectedRemaining1 := fmt.Sprintf("%d", int(rateLimitRate)-1)
+			expectedRemaining2 := fmt.Sprintf("%d", int(rateLimitRate)-2)
+
+			headersMatch1 := map[string]string{
+				header.XRateLimitLimit: fmt.Sprintf("%d", int(rateLimitRate)),
+			}
+			headersMatch2 := map[string]string{
+				header.XRateLimitLimit: fmt.Sprintf("%d", int(rateLimitRate)),
+			}
+
+			// For limiters that don't support Remaining (Sentinel, FixedWindow), it should be assigned to 0.
+			if limiter == "Redis" || limiter == "DRL" {
+				headersMatch1[header.XRateLimitRemaining] = expectedRemaining1
+				headersMatch2[header.XRateLimitRemaining] = expectedRemaining2
+			} else {
+				headersMatch1[header.XRateLimitRemaining] = "0"
+				headersMatch2[header.XRateLimitRemaining] = "0"
+			}
+
+			_, _ = ts.Run(t, []test.TestCase{
+				{
+					Path:         "/api-rate-limit-headers-test",
+					Code:         http.StatusOK,
+					HeadersMatch: headersMatch1,
+				},
+				{
+					Path:         "/api-rate-limit-headers-test",
+					Code:         http.StatusOK,
+					HeadersMatch: headersMatch2,
+				},
+				{
+					Path: "/api-rate-limit-headers-test",
+					Code: http.StatusTooManyRequests,
+				},
+			}...)
+		})
+	}
+}
+
 func TestRLOpen(t *testing.T) {
 	ts := StartTest(nil)
 	defer ts.Close()

gateway/mw_mock_response.go

@@ -168,7 +168,7 @@ func (m *mockResponseMiddleware) mockResponse(r *http.Request) (
 		res.Header.Set(header.Connection, "close")
 	}
 
-	m.Spec.sendRateLimitHeaders(ctxGetSession(r), res)
+	m.Gw.limitHeaderFactory(res.Header).SendQuotas(ctxGetSession(r), m.Spec.APIID)
 
 	return res, internal, nil
 }

gateway/mw_organisation_activity.go

@@ -120,7 +120,13 @@ func (k *OrganizationMonitor) refreshOrgSession(orgID string) {
 }
 
 // ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
-func (k *OrganizationMonitor) ProcessRequestLive(r *http.Request, orgSession *user.SessionState) (error, int) {
+//
+//nolint:staticcheck
+func (k *OrganizationMonitor) ProcessRequestLive(
+	r *http.Request,
+	orgSession *user.SessionState,
+) (error, int) {
+
 	logger := k.Logger()
 
 	if orgSession.IsInactive {
@@ -135,11 +141,11 @@ func (k *OrganizationMonitor) ProcessRequestLive(r *http.Request, orgSession *us
 		orgSession,
 		k.Spec.OrgID,
 		"",
-		k.Spec.OrgSessionManager.Store(),
 		orgSession.Per > 0 && orgSession.Rate > 0,
 		true,
 		k.Spec,
 		false,
+		nil,
 	)
 
 	sessionLifeTime := orgSession.Lifetime(k.Spec.GetSessionLifetimeRespectsKeyExpiration(), k.Spec.SessionLifetime, k.Gw.GetConfig().ForceGlobalSessionLifetime, k.Gw.GetConfig().GlobalSessionLifetime)
@@ -211,7 +217,12 @@ func (k *OrganizationMonitor) SetOrgSentinel(orgChan chan bool, orgId string) {
 	}
 }
 
-func (k *OrganizationMonitor) ProcessRequestOffThread(r *http.Request, orgSession *user.SessionState) (error, int) {
+//nolint:staticcheck
+func (k *OrganizationMonitor) ProcessRequestOffThread(
+	r *http.Request,
+	orgSession *user.SessionState,
+) (error, int) {
+
 	orgChanMap.Lock()
 	orgChan, ok := orgChanMap.channels[k.Spec.OrgID]
 	if !ok {
@@ -251,7 +262,8 @@ func (k *OrganizationMonitor) AllowAccessNext(
 	path string,
 	IP string,
 	r *http.Request,
-	session *user.SessionState) {
+	session *user.SessionState,
+) {
 
 	// Is it active?
 	logEntry := k.Gw.getExplicitLogEntryForRequest(k.Logger(), path, IP, k.Spec.OrgID, nil)
@@ -269,11 +281,11 @@ func (k *OrganizationMonitor) AllowAccessNext(
 		session,
 		k.Spec.OrgID,
 		customQuotaKey,
-		k.Spec.OrgSessionManager.Store(),
 		session.Per > 0 && session.Rate > 0,
 		true,
 		k.Spec,
 		false,
+		nil,
 	)
 
 	sessionLifeTime := session.Lifetime(k.Spec.GetSessionLifetimeRespectsKeyExpiration(), k.Spec.SessionLifetime, k.Gw.GetConfig().ForceGlobalSessionLifetime, k.Gw.GetConfig().GlobalSessionLifetime)