Skip to content

feat: Add per-unit response headers for all descriptors via LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED #1190

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
samay-arcana wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add per-unit response headers for all descriptors via LIMIT_ALL_DESCRIPTORS_HEADERS_ENABLED #1190

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 55 additions & 1 deletion src/service/ratelimit.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ type service struct {
globalShadowMode bool
globalQuotaMode bool
responseDynamicMetadataEnabled bool
allDescriptorsHeadersEnabled bool
}

func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWithAtLeastOneConfigLoad bool) {
Expand Down Expand Up @@ -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")
}
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 5 additions & 0 deletions src/settings/settings.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"`

Expand Down
160 changes: 160 additions & 0 deletions test/service/ratelimit_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}