From 26dab879f5c8ce50f7abdc9373c8d485698da986 Mon Sep 17 00:00:00 2001 From: Samay Varshney Date: Mon, 13 Jul 2026 22:39:20 +0530 Subject: [PATCH] Return all rate limit headers via env var Signed-off-by: Samay Varshney --- src/service/ratelimit.go | 56 +++++++++++- src/settings/settings.go | 5 ++ test/service/ratelimit_test.go | 160 +++++++++++++++++++++++++++++++++ 3 files changed, 220 insertions(+), 1 deletion(-) diff --git a/src/service/ratelimit.go b/src/service/ratelimit.go index 9bd9241d..9b69f5b0 100644 --- a/src/service/ratelimit.go +++ b/src/service/ratelimit.go @@ -54,6 +54,7 @@ type service struct { globalShadowMode bool globalQuotaMode bool responseDynamicMetadataEnabled bool + allDescriptorsHeadersEnabled bool } func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWithAtLeastOneConfigLoad bool) { @@ -100,6 +101,7 @@ func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWi this.customHeaderResetHeader = rlSettings.HeaderRatelimitReset } + this.allDescriptorsHeadersEnabled = rlSettings.RateLimitAllDescriptorsHeadersEnabled this.configLock.Unlock() logger.Info("Successfully loaded new configuration") } @@ -255,7 +257,9 @@ func (this *service) shouldRateLimitWorker( } // Add Headers if requested - if this.customHeadersEnabled && minimumDescriptor != nil { + if this.allDescriptorsHeadersEnabled { + response.ResponseHeadersToAdd = this.allDescriptorsHeaders(responseDescriptorStatuses) + } else if this.customHeadersEnabled && minimumDescriptor != nil { response.ResponseHeadersToAdd = []*core.HeaderValue{ this.rateLimitLimitHeader(minimumDescriptor), this.rateLimitRemainingHeader(minimumDescriptor), @@ -397,6 +401,56 @@ func (this *service) rateLimitResetHeader( } } +// unitToHeaderSuffix converts a rate limit unit enum to a lowercase plural suffix +// used in per-unit response header names (e.g., "seconds", "minutes"). +func unitToHeaderSuffix(unit pb.RateLimitResponse_RateLimit_Unit) string { + switch unit { + case pb.RateLimitResponse_RateLimit_SECOND: + return "seconds" + case pb.RateLimitResponse_RateLimit_MINUTE: + return "minutes" + case pb.RateLimitResponse_RateLimit_HOUR: + return "hours" + case pb.RateLimitResponse_RateLimit_DAY: + return "days" + case pb.RateLimitResponse_RateLimit_WEEK: + return "weeks" + case pb.RateLimitResponse_RateLimit_MONTH: + return "months" + case pb.RateLimitResponse_RateLimit_YEAR: + return "years" + default: + return strings.ToLower(unit.String()) + "s" + } +} + +// allDescriptorsHeaders generates per-unit limit and remaining response headers +// for every descriptor status that has a configured limit. This produces headers like: +// - ratelimit-limit-seconds: 10 +// - ratelimit-remaining-seconds: 9 +// - ratelimit-limit-minutes: 1000 +// - ratelimit-remaining-minutes: 999 +func (this *service) allDescriptorsHeaders( + descriptorStatuses []*pb.RateLimitResponse_DescriptorStatus, +) []*core.HeaderValue { + var headers []*core.HeaderValue + for _, status := range descriptorStatuses { + if status.CurrentLimit == nil { + continue + } + unitSuffix := unitToHeaderSuffix(status.CurrentLimit.Unit) + headers = append(headers, &core.HeaderValue{ + Key: "ratelimit-limit-" + unitSuffix, + Value: strconv.FormatUint(uint64(status.CurrentLimit.RequestsPerUnit), 10), + }) + headers = append(headers, &core.HeaderValue{ + Key: "ratelimit-remaining-" + unitSuffix, + Value: strconv.FormatUint(uint64(status.LimitRemaining), 10), + }) + } + return headers +} + func (this *service) ShouldRateLimit( ctx context.Context, request *pb.RateLimitRequest, diff --git a/src/settings/settings.go b/src/settings/settings.go index 2129cf80..9f1c08f3 100644 --- a/src/settings/settings.go +++ b/src/settings/settings.go @@ -120,6 +120,11 @@ type Settings struct { // value: remaining seconds HeaderRatelimitReset string `envconfig:"LIMIT_RESET_HEADER" default:"RateLimit-Reset"` + // When enabled, response headers include per-unit limit and remaining headers for ALL descriptors + // (e.g., ratelimit-limit-seconds, ratelimit-remaining-seconds, ratelimit-limit-minutes, ratelimit-remaining-minutes) + // instead of only returning headers for the single descriptor closest to hitting the rate limit. + RateLimitAllDescriptorsHeadersEnabled bool `envconfig:"LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED" default:"false"` + // Health-check settings HealthyWithAtLeastOneConfigLoaded bool `envconfig:"HEALTHY_WITH_AT_LEAST_ONE_CONFIG_LOADED" default:"false"` diff --git a/test/service/ratelimit_test.go b/test/service/ratelimit_test.go index dc4ab8a7..a70d7d48 100644 --- a/test/service/ratelimit_test.go +++ b/test/service/ratelimit_test.go @@ -1292,3 +1292,163 @@ func TestServiceMixedModeWithShadowMode(test *testing.T) { // Verify global shadow mode counter is incremented t.assert.EqualValues(1, t.statStore.NewCounter("global_shadow_mode").Value()) } + +func TestServiceWithAllDescriptorsHeaders(test *testing.T) { + os.Setenv("LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED", "true") + defer func() { + os.Unsetenv("LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + // Config reload to pick up the env var. + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + // Two descriptors with different time units (SECOND and MINUTE) + request := common.NewRateLimitRequest( + "test-domain", [][][2]string{{{"remote_address", "10.0.0.1"}}, {{"remote_address", "10.0.0.1"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, t.statsManager.NewStats("key1"), false, false, false, "", nil, false), + config.NewRateLimit(1000, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key2"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "test-domain", request.Descriptors[0]).Return(limits[0]) + t.config.EXPECT().GetLimit(context.Background(), "test-domain", request.Descriptors[1]).Return(limits[1]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 9, DurationUntilReset: nil}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 999, DurationUntilReset: nil}, + }) + + response, err := service.ShouldRateLimit(context.Background(), request) + common.AssertProtoEqual( + t.assert, + &pb.RateLimitResponse{ + OverallCode: pb.RateLimitResponse_OK, + Statuses: []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 9}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 999}, + }, + ResponseHeadersToAdd: []*core.HeaderValue{ + {Key: "ratelimit-limit-seconds", Value: "10"}, + {Key: "ratelimit-remaining-seconds", Value: "9"}, + {Key: "ratelimit-limit-minutes", Value: "1000"}, + {Key: "ratelimit-remaining-minutes", Value: "999"}, + }, + }, + response) + t.assert.Nil(err) +} + +func TestServiceWithAllDescriptorsHeadersOverLimit(test *testing.T) { + os.Setenv("LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED", "true") + defer func() { + os.Unsetenv("LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + // Config reload to pick up the env var. + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + // Two descriptors: SECOND is over limit, MINUTE is OK + request := common.NewRateLimitRequest( + "test-domain", [][][2]string{{{"remote_address", "10.0.0.1"}}, {{"remote_address", "10.0.0.1"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, t.statsManager.NewStats("key1"), false, false, false, "", nil, false), + config.NewRateLimit(1000, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key2"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "test-domain", request.Descriptors[0]).Return(limits[0]) + t.config.EXPECT().GetLimit(context.Background(), "test-domain", request.Descriptors[1]).Return(limits[1]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 500}, + }) + + response, err := service.ShouldRateLimit(context.Background(), request) + common.AssertProtoEqual( + t.assert, + &pb.RateLimitResponse{ + OverallCode: pb.RateLimitResponse_OVER_LIMIT, + Statuses: []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 500}, + }, + ResponseHeadersToAdd: []*core.HeaderValue{ + {Key: "ratelimit-limit-seconds", Value: "10"}, + {Key: "ratelimit-remaining-seconds", Value: "0"}, + {Key: "ratelimit-limit-minutes", Value: "1000"}, + {Key: "ratelimit-remaining-minutes", Value: "500"}, + }, + }, + response) + t.assert.Nil(err) +} + +func TestServiceWithAllDescriptorsHeadersSkipsNilLimits(test *testing.T) { + os.Setenv("LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED", "true") + defer func() { + os.Unsetenv("LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + // Config reload to pick up the env var. + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + // One descriptor with limit, one without (nil limit) + request := common.NewRateLimitRequest( + "test-domain", [][][2]string{{{"remote_address", "10.0.0.1"}}, {{"other_key", "value"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, t.statsManager.NewStats("key1"), false, false, false, "", nil, false), + nil, + } + t.config.EXPECT().GetLimit(context.Background(), "test-domain", request.Descriptors[0]).Return(limits[0]) + t.config.EXPECT().GetLimit(context.Background(), "test-domain", request.Descriptors[1]).Return(limits[1]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 9}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: nil, LimitRemaining: 0}, + }) + + response, err := service.ShouldRateLimit(context.Background(), request) + common.AssertProtoEqual( + t.assert, + &pb.RateLimitResponse{ + OverallCode: pb.RateLimitResponse_OK, + Statuses: []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 9}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: nil, LimitRemaining: 0}, + }, + ResponseHeadersToAdd: []*core.HeaderValue{ + {Key: "ratelimit-limit-seconds", Value: "10"}, + {Key: "ratelimit-remaining-seconds", Value: "9"}, + }, + }, + response) + t.assert.Nil(err) +}