diff --git a/README.md b/README.md index f2e8d6df..d41fe43f 100644 --- a/README.md +++ b/README.md @@ -1414,6 +1414,15 @@ The following environment variables control the custom response feature: 1. `LIMIT_REMAINING_HEADER` - The default value is "RateLimit-Remaining", setting the environment variable will specify an alternative header name 1. `LIMIT_RESET_HEADER` - The default value is "RateLimit-Reset", setting the environment variable will specify an alternative header name +## RequestHeadersToAdd + +The following environment variables control the custom request header feature. When enabled, Envoy injects these headers into the forwarded request before it reaches the upstream service, allowing upstream services to inspect rate limit state and make per-request routing decisions (e.g. skip a hot path when quota is exhausted) without the request being blocked. + +1. `LIMIT_REQUEST_HEADERS_ENABLED` - Enables the custom request headers +1. `LIMIT_REQUEST_LIMIT_HEADER` - The default value is "RateLimit-Limit", setting the environment variable will specify an alternative header name +1. `LIMIT_REQUEST_REMAINING_HEADER` - The default value is "RateLimit-Remaining", setting the environment variable will specify an alternative header name +1. `LIMIT_REQUEST_RESET_HEADER` - The default value is "RateLimit-Reset", setting the environment variable will specify an alternative header name + # Tracing Ratelimit service supports exporting spans in OLTP format. See [OpenTelemetry](https://opentelemetry.io/) for more information. diff --git a/docs/superpowers/plans/2026-06-04-request-headers-to-add.md b/docs/superpowers/plans/2026-06-04-request-headers-to-add.md new file mode 100644 index 00000000..34e18e95 --- /dev/null +++ b/docs/superpowers/plans/2026-06-04-request-headers-to-add.md @@ -0,0 +1,564 @@ +# RequestHeadersToAdd Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add `RequestHeadersToAdd` support so the ratelimit service can inject RateLimit-* headers into the forwarded request (upstream direction), mirroring the existing `ResponseHeadersToAdd` feature. + +**Architecture:** Four new env vars (`LIMIT_REQUEST_HEADERS_ENABLED`, `LIMIT_REQUEST_LIMIT_HEADER`, `LIMIT_REQUEST_REMAINING_HEADER`, `LIMIT_REQUEST_RESET_HEADER`) control the feature. When enabled, `shouldRateLimitWorker` populates `response.RequestHeadersToAdd` using the same `minimumDescriptor` logic already used for `ResponseHeadersToAdd`. The two features are independently controlled. + +**Tech Stack:** Go, `github.com/envoyproxy/go-control-plane` (proto types for `RateLimitResponse` and `HeaderValue`), `github.com/kelseyhightower/envconfig` (settings), `github.com/golang/mock` (test mocks), `github.com/stretchr/testify` (assertions). + +--- + +## File Map + +| File | Change | +|---|---| +| `src/settings/settings.go` | Add 4 new env var fields | +| `src/service/ratelimit.go` | Add 4 struct fields, wire in `SetConfig`, populate `RequestHeadersToAdd` in `shouldRateLimitWorker` | +| `test/service/ratelimit_test.go` | Add 4 new test functions | +| `README.md` | Add `Global Rate Limit Request Headers` section after existing `Custom headers` section | + +--- + +## Task 1: Add settings fields + +**Files:** +- Modify: `src/settings/settings.go:114-121` + +- [ ] **Step 1: Write a failing test that reads the new settings fields** + +In `test/service/ratelimit_test.go`, before `TestServiceWithCustomRatelimitHeaders`, add: + +```go +func TestRequestHeadersSettingsDefaults(test *testing.T) { + s := settings.NewSettings() + assert.False(test, s.RateLimitRequestHeadersEnabled) + assert.Equal(test, "RateLimit-Limit", s.HeaderRequestRatelimitLimit) + assert.Equal(test, "RateLimit-Remaining", s.HeaderRequestRatelimitRemaining) + assert.Equal(test, "RateLimit-Reset", s.HeaderRequestRatelimitReset) +} +``` + +Add `"github.com/envoyproxy/ratelimit/src/settings"` to the import block if not already present. Check with: + +```bash +head -30 test/service/ratelimit_test.go +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +go test ./test/service/... -run TestRequestHeadersSettingsDefaults -v +``` + +Expected: compile error — `s.RateLimitRequestHeadersEnabled` undefined. + +- [ ] **Step 3: Add the four new fields to `src/settings/settings.go`** + +After line 121 (`HeaderRatelimitReset string ...`), insert: + +```go +// Settings for optional returning of custom request headers +RateLimitRequestHeadersEnabled bool `envconfig:"LIMIT_REQUEST_HEADERS_ENABLED" default:"false"` +// value: the current limit +HeaderRequestRatelimitLimit string `envconfig:"LIMIT_REQUEST_LIMIT_HEADER" default:"RateLimit-Limit"` +// value: remaining count +HeaderRequestRatelimitRemaining string `envconfig:"LIMIT_REQUEST_REMAINING_HEADER" default:"RateLimit-Remaining"` +// value: remaining seconds +HeaderRequestRatelimitReset string `envconfig:"LIMIT_REQUEST_RESET_HEADER" default:"RateLimit-Reset"` +``` + +- [ ] **Step 4: Run test to verify it passes** + +```bash +go test ./test/service/... -run TestRequestHeadersSettingsDefaults -v +``` + +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/settings/settings.go test/service/ratelimit_test.go +git commit -m "feat: add RequestHeaders settings fields" +``` + +--- + +## Task 2: Wire settings into service struct + +**Files:** +- Modify: `src/service/ratelimit.go:42-57` (struct), `src/service/ratelimit.go:89-103` (SetConfig) + +- [ ] **Step 1: Write a failing test — request headers disabled by default** + +In `test/service/ratelimit_test.go`, add after `TestRequestHeadersSettingsDefaults`: + +```go +func TestRequestHeadersDisabledByDefault(test *testing.T) { + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, + }) + + response, err := service.ShouldRateLimit(context.Background(), request) + t.assert.Nil(err) + t.assert.Nil(response.RequestHeadersToAdd) +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +go test ./test/service/... -run TestRequestHeadersDisabledByDefault -v +``` + +Expected: compile error — `response.RequestHeadersToAdd` doesn't exist yet on the struct... actually this test will pass immediately since `RequestHeadersToAdd` already exists on the proto. Confirm by running — if it passes, proceed to step 3. + +- [ ] **Step 3: Add four fields to the `service` struct** + +In `src/service/ratelimit.go`, after line 53 (`customHeaderClock utils.TimeSource`), add: + +```go +requestHeadersEnabled bool +requestHeaderLimitHeader string +requestHeaderRemainingHeader string +requestHeaderResetHeader string +``` + +- [ ] **Step 4: Wire the fields in `SetConfig`** + +In `src/service/ratelimit.go`, after the closing brace of the `if rlSettings.RateLimitResponseHeadersEnabled` block (after line 102), add: + +```go +if rlSettings.RateLimitRequestHeadersEnabled { + this.requestHeadersEnabled = true + this.requestHeaderLimitHeader = rlSettings.HeaderRequestRatelimitLimit + this.requestHeaderRemainingHeader = rlSettings.HeaderRequestRatelimitRemaining + this.requestHeaderResetHeader = rlSettings.HeaderRequestRatelimitReset +} +``` + +- [ ] **Step 5: Run tests to verify nothing is broken** + +```bash +go test ./test/service/... -run "TestRequestHeadersDisabledByDefault|TestRequestHeadersSettingsDefaults" -v +``` + +Expected: both PASS. + +- [ ] **Step 6: Commit** + +```bash +git add src/service/ratelimit.go test/service/ratelimit_test.go +git commit -m "feat: wire request headers settings into service struct" +``` + +--- + +## Task 3: Populate RequestHeadersToAdd in shouldRateLimitWorker + +**Files:** +- Modify: `src/service/ratelimit.go:258-264` + +- [ ] **Step 1: Write failing test — default header names, over limit** + +In `test/service/ratelimit_test.go`, add: + +```go +func TestServiceWithDefaultRequestHeaders(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}, {{"hello", "world"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + nil, + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.config.EXPECT().GetLimit(context.Background(), "different-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: nil, LimitRemaining: 0}, + }) + + 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: nil, LimitRemaining: 0}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "RateLimit-Limit", Value: "10"}, + {Key: "RateLimit-Remaining", Value: "0"}, + {Key: "RateLimit-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} +``` + +Note: `Value: "58"` is the reset seconds computed from `MockClock{now: 2222}` with a MINUTE unit — same value used in the existing response header tests. + +- [ ] **Step 2: Run test to verify it fails** + +```bash +go test ./test/service/... -run TestServiceWithDefaultRequestHeaders -v +``` + +Expected: FAIL — `RequestHeadersToAdd` is nil. + +- [ ] **Step 3: Populate RequestHeadersToAdd in shouldRateLimitWorker** + +In `src/service/ratelimit.go`, after the `ResponseHeadersToAdd` block (after line 264, the closing `}`), add: + +```go +// Add request headers if requested +if this.requestHeadersEnabled && minimumDescriptor != nil { + response.RequestHeadersToAdd = []*core.HeaderValue{ + {Key: this.requestHeaderLimitHeader, Value: strconv.FormatUint(uint64(minimumDescriptor.CurrentLimit.RequestsPerUnit), 10)}, + {Key: this.requestHeaderRemainingHeader, Value: strconv.FormatUint(uint64(minimumDescriptor.LimitRemaining), 10)}, + {Key: this.requestHeaderResetHeader, Value: strconv.FormatInt(utils.CalculateReset(&minimumDescriptor.CurrentLimit.Unit, this.customHeaderClock).GetSeconds(), 10)}, + } +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +```bash +go test ./test/service/... -run TestServiceWithDefaultRequestHeaders -v +``` + +Expected: PASS. + +- [ ] **Step 5: Run all existing tests to verify no regressions** + +```bash +go test ./test/service/... -v 2>&1 | tail -20 +``` + +Expected: all PASS. + +- [ ] **Step 6: Commit** + +```bash +git add src/service/ratelimit.go test/service/ratelimit_test.go +git commit -m "feat: populate RequestHeadersToAdd in shouldRateLimitWorker" +``` + +--- + +## Task 4: Add remaining test cases + +**Files:** +- Modify: `test/service/ratelimit_test.go` + +- [ ] **Step 1: Add test — custom header names** + +In `test/service/ratelimit_test.go`, add after `TestServiceWithDefaultRequestHeaders`: + +```go +func TestServiceWithCustomRequestHeaders(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + os.Setenv("LIMIT_REQUEST_LIMIT_HEADER", "X-RateLimit-Limit") + os.Setenv("LIMIT_REQUEST_REMAINING_HEADER", "X-RateLimit-Remaining") + os.Setenv("LIMIT_REQUEST_RESET_HEADER", "X-RateLimit-Reset") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + os.Unsetenv("LIMIT_REQUEST_LIMIT_HEADER") + os.Unsetenv("LIMIT_REQUEST_REMAINING_HEADER") + os.Unsetenv("LIMIT_REQUEST_RESET_HEADER") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}, {{"hello", "world"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + nil, + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.config.EXPECT().GetLimit(context.Background(), "different-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: nil, LimitRemaining: 0}, + }) + + 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: nil, LimitRemaining: 0}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "X-RateLimit-Limit", Value: "10"}, + {Key: "X-RateLimit-Remaining", Value: "0"}, + {Key: "X-RateLimit-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} +``` + +- [ ] **Step 2: Add test — headers present when within limit (not over limit)** + +```go +func TestServiceWithRequestHeadersWithinLimit(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 8}, + }) + + 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: 8}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "RateLimit-Limit", Value: "10"}, + {Key: "RateLimit-Remaining", Value: "8"}, + {Key: "RateLimit-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} +``` + +- [ ] **Step 3: Add test — both request and response headers enabled simultaneously with different names** + +```go +func TestServiceWithBothRequestAndResponseHeaders(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + os.Setenv("LIMIT_REQUEST_LIMIT_HEADER", "X-Upstream-Limit") + os.Setenv("LIMIT_REQUEST_REMAINING_HEADER", "X-Upstream-Remaining") + os.Setenv("LIMIT_REQUEST_RESET_HEADER", "X-Upstream-Reset") + os.Setenv("LIMIT_RESPONSE_HEADERS_ENABLED", "true") + os.Setenv("LIMIT_LIMIT_HEADER", "X-Downstream-Limit") + os.Setenv("LIMIT_REMAINING_HEADER", "X-Downstream-Remaining") + os.Setenv("LIMIT_RESET_HEADER", "X-Downstream-Reset") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + os.Unsetenv("LIMIT_REQUEST_LIMIT_HEADER") + os.Unsetenv("LIMIT_REQUEST_REMAINING_HEADER") + os.Unsetenv("LIMIT_REQUEST_RESET_HEADER") + os.Unsetenv("LIMIT_RESPONSE_HEADERS_ENABLED") + os.Unsetenv("LIMIT_LIMIT_HEADER") + os.Unsetenv("LIMIT_REMAINING_HEADER") + os.Unsetenv("LIMIT_RESET_HEADER") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, + }) + + 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}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "X-Upstream-Limit", Value: "10"}, + {Key: "X-Upstream-Remaining", Value: "0"}, + {Key: "X-Upstream-Reset", Value: "58"}, + }, + ResponseHeadersToAdd: []*core.HeaderValue{ + {Key: "X-Downstream-Limit", Value: "10"}, + {Key: "X-Downstream-Remaining", Value: "0"}, + {Key: "X-Downstream-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} +``` + +- [ ] **Step 4: Run all four new tests** + +```bash +go test ./test/service/... -run "TestServiceWithDefaultRequestHeaders|TestServiceWithCustomRequestHeaders|TestServiceWithRequestHeadersWithinLimit|TestServiceWithBothRequestAndResponseHeaders" -v +``` + +Expected: all 4 PASS. + +- [ ] **Step 5: Run the full test suite** + +```bash +go test ./... 2>&1 | tail -20 +``` + +Expected: all PASS. + +- [ ] **Step 6: Commit** + +```bash +git add test/service/ratelimit_test.go +git commit -m "test: add RequestHeadersToAdd test cases" +``` + +--- + +## Task 5: Update README + +**Files:** +- Modify: `README.md:1406-1416` + +- [ ] **Step 1: Insert the new section after the existing `Custom headers` section** + +In `README.md`, after line 1415 (`1. \`LIMIT_RESET_HEADER\` ...`) and before line 1417 (`# Tracing`), insert: + +```markdown + +The following environment variables control the custom request header feature (headers injected into the forwarded request sent to upstream services by Envoy): + +1. `LIMIT_REQUEST_HEADERS_ENABLED` - Enables the custom request headers. When enabled, Envoy injects these headers into the forwarded request, allowing upstream services to inspect rate limit state and make per-request routing decisions (e.g. skip a hot path when quota is exhausted) without the request being blocked. +1. `LIMIT_REQUEST_LIMIT_HEADER` - The default value is "RateLimit-Limit", setting the environment variable will specify an alternative header name +1. `LIMIT_REQUEST_REMAINING_HEADER` - The default value is "RateLimit-Remaining", setting the environment variable will specify an alternative header name +1. `LIMIT_REQUEST_RESET_HEADER` - The default value is "RateLimit-Reset", setting the environment variable will specify an alternative header name +``` + +- [ ] **Step 2: Verify the section reads correctly** + +```bash +grep -A 10 "LIMIT_REQUEST_HEADERS_ENABLED" README.md +``` + +Expected: the 4 env vars appear with their descriptions. + +- [ ] **Step 3: Commit** + +```bash +git add README.md +git commit -m "docs: add RequestHeadersToAdd documentation to README" +``` + +--- + +## Task 6: Final verification + +- [ ] **Step 1: Run the full test suite one more time** + +```bash +go test ./... 2>&1 | grep -E "FAIL|ok" +``` + +Expected: all lines show `ok`, no `FAIL`. + +- [ ] **Step 2: Check the build** + +```bash +go build ./... +``` + +Expected: exits with code 0, no output. + +- [ ] **Step 3: Review the diff** + +```bash +git diff origin/main..HEAD --stat +``` + +Confirm only these files changed: `src/settings/settings.go`, `src/service/ratelimit.go`, `test/service/ratelimit_test.go`, `README.md`, `docs/`. diff --git a/docs/superpowers/specs/2026-06-04-request-headers-to-add-design.md b/docs/superpowers/specs/2026-06-04-request-headers-to-add-design.md new file mode 100644 index 00000000..57e49a5a --- /dev/null +++ b/docs/superpowers/specs/2026-06-04-request-headers-to-add-design.md @@ -0,0 +1,90 @@ +# Design: `RequestHeadersToAdd` Support + +**Date:** 2026-06-04 +**Repo:** github.com/envoyproxy/ratelimit +**Scope:** Open source contribution + +## Problem + +The ratelimit service currently supports injecting rate limit headers into the **downstream response** via `ResponseHeadersToAdd` (enabled with `LIMIT_RESPONSE_HEADERS_ENABLED`). These headers go back to the caller, not to the upstream service. + +The Envoy `RateLimitResponse` proto also has a `RequestHeadersToAdd` field — headers Envoy injects into the **forwarded request** before sending it upstream. The ratelimit service never populates this field today. + +Upstream services (e.g. API gateways, sidecars) that sit behind Envoy have no per-request signal from the rate limiter. They cannot make routing decisions based on quota state without this. + +## Goal + +Add a `RequestHeadersToAdd` feature that mirrors `ResponseHeadersToAdd` exactly: same 3 standard headers (`RateLimit-Limit`, `RateLimit-Remaining`, `RateLimit-Reset`), same `minimumDescriptor` logic, independently enabled via env vars. + +## Non-Goals + +- Per-domain or per-descriptor scoping (would require XDS proto changes; deferred) +- Arbitrary custom header injection (out of scope for this change) +- Filtering headers to OVER_LIMIT-only responses (operators can inspect `RateLimit-Remaining: 0`) + +## Design + +### Settings + +Four new env vars added to `src/settings/settings.go`, parallel to the existing response header settings: + +| Env var | Default | Purpose | +|---|---|---| +| `LIMIT_REQUEST_HEADERS_ENABLED` | `false` | Master on/off switch | +| `LIMIT_REQUEST_LIMIT_HEADER` | `RateLimit-Limit` | Header name for the limit value | +| `LIMIT_REQUEST_REMAINING_HEADER` | `RateLimit-Remaining` | Header name for remaining count | +| `LIMIT_REQUEST_RESET_HEADER` | `RateLimit-Reset` | Header name for reset seconds | + +Default names match the response header defaults. Operators can set different names to distinguish upstream vs downstream headers (e.g. `x-grls-ratelimit-remaining` upstream vs `RateLimit-Remaining` downstream). + +The two features (`LIMIT_RESPONSE_HEADERS_ENABLED` and `LIMIT_REQUEST_HEADERS_ENABLED`) are independently controlled — operators can enable one, both, or neither. + +### Service struct (`src/service/ratelimit.go`) + +Four new fields added to the `service` struct, parallel to `customHeaders*`: + +```go +requestHeadersEnabled bool +requestHeaderLimitHeader string +requestHeaderRemainingHeader string +requestHeaderResetHeader string +``` + +Wired in `SetConfig()` alongside the existing response header fields (currently lines 94–101). + +### Response construction (`shouldRateLimitWorker`) + +After the existing `ResponseHeadersToAdd` block (currently lines 258–264), add: + +```go +if this.requestHeadersEnabled && minimumDescriptor != nil { + response.RequestHeadersToAdd = []*core.HeaderValue{ + {Key: this.requestHeaderLimitHeader, Value: }, + {Key: this.requestHeaderRemainingHeader, Value: }, + {Key: this.requestHeaderResetHeader, Value: }, + } +} +``` + +The `minimumDescriptor` (the descriptor closest to its limit) is reused as-is. Headers are always populated when enabled, regardless of whether the overall response is `OK` or `OVER_LIMIT` — mirroring response header behavior exactly. + +### Tests (`test/service/ratelimit_test.go`) + +Parallel test cases for `RequestHeadersToAdd` alongside the existing response header tests: + +1. **Headers present when enabled + within limit** — `RequestHeadersToAdd` populated correctly +2. **Headers present when enabled + over limit** — `RequestHeadersToAdd` populated correctly +3. **Headers absent when disabled** — `RequestHeadersToAdd` is nil +4. **Both enabled simultaneously with different header names** — `RequestHeadersToAdd` and `ResponseHeadersToAdd` independently populated with their respective configured names; no interference + +### README + +Add a `Global Rate Limit Request Headers` section immediately after the existing `Global Rate Limit Response Headers` section, documenting the 4 env vars and the use case: upstream services behind Envoy can read these headers to make per-request routing decisions (e.g. skip hot path on quota exhaustion) without the rate limit service blocking the request. + +## Backwards Compatibility + +Fully additive. `LIMIT_REQUEST_HEADERS_ENABLED` defaults to `false`. No YAML schema changes. No XDS proto changes. Existing deployments see no behavior change. + +## Future Work + +Per-domain scoping would require adding a field to `RateLimitConfig` in the XDS proto (`api/ratelimit/config/ratelimit/v3/rls_conf.proto`) and a coordinated change to `go-control-plane`. This is a natural follow-on once the base feature is established. diff --git a/examples/envoy/mock.yaml b/examples/envoy/mock.yaml index 226906ab..dc254147 100644 --- a/examples/envoy/mock.yaml +++ b/examples/envoy/mock.yaml @@ -4,6 +4,45 @@ static_resources: socket_address: address: 0.0.0.0 port_value: 9999 + filter_chains: + - filters: + - name: envoy.filters.network.http_connection_manager + typed_config: + "@type": type.googleapis.com/envoy.extensions.filters.network.http_connection_manager.v3.HttpConnectionManager + codec_type: AUTO + stat_prefix: ingress + # Access log for integration testing: logs the default RateLimit-* request headers + # injected by Envoy when LIMIT_REQUEST_HEADERS_ENABLED=true. + # If you configure custom header names via LIMIT_REQUEST_*_HEADER env vars, + # update the %REQ(...)% references below accordingly. + access_log: + - name: envoy.access_loggers.stdout + typed_config: + "@type": type.googleapis.com/envoy.extensions.access_loggers.stream.v3.StdoutAccessLog + log_format: + text_format_source: + inline_string: "[mock-upstream] method=%REQ(:METHOD)% path=%REQ(:PATH)% RateLimit-Limit=%REQ(RateLimit-Limit)% RateLimit-Remaining=%REQ(RateLimit-Remaining)% RateLimit-Reset=%REQ(RateLimit-Reset)%\n" + route_config: + name: ingress + virtual_hosts: + - name: backend + domains: + - "*" + routes: + - match: + prefix: "/" + direct_response: + status: "200" + body: + inline_string: "Hello World from Service 1" + http_filters: + - name: envoy.filters.http.router + typed_config: + "@type": type.googleapis.com/envoy.extensions.filters.http.router.v3.Router + - address: + socket_address: + address: 0.0.0.0 + port_value: 10001 filter_chains: - filters: - name: envoy.filters.network.http_connection_manager diff --git a/integration-test/docker-compose-integration-test.yml b/integration-test/docker-compose-integration-test.yml index 46d16a64..775b8e42 100644 --- a/integration-test/docker-compose-integration-test.yml +++ b/integration-test/docker-compose-integration-test.yml @@ -23,7 +23,7 @@ services: networks: - ratelimit-network volumes: - - ./examples/prom-statsd-exporter/conf.yaml:/etc/statsd-exporter/conf.yaml + - ${PWD}/examples/prom-statsd-exporter/conf.yaml:/etc/statsd-exporter/conf.yaml ratelimit: build: @@ -40,7 +40,7 @@ services: networks: - ratelimit-network volumes: - - ./examples/ratelimit/config:/data/ratelimit/config + - ${PWD}/examples/ratelimit/config:/data/ratelimit/config environment: - USE_STATSD=true - STATSD_HOST=statsd @@ -52,6 +52,7 @@ services: - RUNTIME_SUBDIRECTORY=ratelimit - RUNTIME_WATCH_ROOT=false - RESPONSE_DYNAMIC_METADATA=true + - LIMIT_REQUEST_HEADERS_ENABLED=true envoy-proxy: image: envoyproxy/envoy:dev @@ -66,7 +67,7 @@ services: depends_on: - ratelimit volumes: - - ./examples/envoy/proxy.yaml:/etc/envoy/envoy.yaml + - ${PWD}/examples/envoy/proxy.yaml:/etc/envoy/envoy.yaml networks: - ratelimit-network expose: @@ -87,7 +88,7 @@ services: - "--mode serve" - "--log-level info" volumes: - - ./examples/envoy/mock.yaml:/etc/envoy/envoy.yaml + - ${PWD}/examples/envoy/mock.yaml:/etc/envoy/envoy.yaml networks: - ratelimit-network expose: diff --git a/src/service/ratelimit.go b/src/service/ratelimit.go index 9bd9241d..dacec818 100644 --- a/src/service/ratelimit.go +++ b/src/service/ratelimit.go @@ -51,6 +51,10 @@ type service struct { customHeaderRemainingHeader string customHeaderResetHeader string customHeaderClock utils.TimeSource + requestHeadersEnabled bool + requestHeaderLimitHeader string + requestHeaderRemainingHeader string + requestHeaderResetHeader string globalShadowMode bool globalQuotaMode bool responseDynamicMetadataEnabled bool @@ -91,15 +95,21 @@ func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWi this.globalQuotaMode = rlSettings.GlobalQuotaMode this.responseDynamicMetadataEnabled = rlSettings.ResponseDynamicMetadata + this.customHeadersEnabled = rlSettings.RateLimitResponseHeadersEnabled if rlSettings.RateLimitResponseHeadersEnabled { - this.customHeadersEnabled = true - this.customHeaderLimitHeader = rlSettings.HeaderRatelimitLimit this.customHeaderRemainingHeader = rlSettings.HeaderRatelimitRemaining this.customHeaderResetHeader = rlSettings.HeaderRatelimitReset } + + this.requestHeadersEnabled = rlSettings.RateLimitRequestHeadersEnabled + if rlSettings.RateLimitRequestHeadersEnabled { + this.requestHeaderLimitHeader = rlSettings.HeaderRequestRatelimitLimit + this.requestHeaderRemainingHeader = rlSettings.HeaderRequestRatelimitRemaining + this.requestHeaderResetHeader = rlSettings.HeaderRequestRatelimitReset + } this.configLock.Unlock() logger.Info("Successfully loaded new configuration") } @@ -214,7 +224,7 @@ func (this *service) shouldRateLimitWorker( for i, descriptorStatus := range responseDescriptorStatuses { // Keep track of the descriptor closest to hit the ratelimit - if this.customHeadersEnabled && + if (this.customHeadersEnabled || this.requestHeadersEnabled) && descriptorStatus.CurrentLimit != nil && descriptorStatus.LimitRemaining < minLimitRemaining { minimumDescriptor = descriptorStatus @@ -263,6 +273,15 @@ func (this *service) shouldRateLimitWorker( } } + // Add request headers if requested + if this.requestHeadersEnabled && minimumDescriptor != nil { + response.RequestHeadersToAdd = []*core.HeaderValue{ + this.rateLimitRequestLimitHeader(minimumDescriptor), + this.rateLimitRequestRemainingHeader(minimumDescriptor), + this.rateLimitRequestResetHeader(minimumDescriptor), + } + } + // If there is a global shadow_mode, it should always return OK if finalCode == pb.RateLimitResponse_OVER_LIMIT && globalShadowMode { finalCode = pb.RateLimitResponse_OK @@ -397,6 +416,27 @@ func (this *service) rateLimitResetHeader( } } +func (this *service) rateLimitRequestLimitHeader(descriptor *pb.RateLimitResponse_DescriptorStatus) *core.HeaderValue { + return &core.HeaderValue{ + Key: this.requestHeaderLimitHeader, + Value: strconv.FormatUint(uint64(descriptor.CurrentLimit.RequestsPerUnit), 10), + } +} + +func (this *service) rateLimitRequestRemainingHeader(descriptor *pb.RateLimitResponse_DescriptorStatus) *core.HeaderValue { + return &core.HeaderValue{ + Key: this.requestHeaderRemainingHeader, + Value: strconv.FormatUint(uint64(descriptor.LimitRemaining), 10), + } +} + +func (this *service) rateLimitRequestResetHeader(descriptor *pb.RateLimitResponse_DescriptorStatus) *core.HeaderValue { + return &core.HeaderValue{ + Key: this.requestHeaderResetHeader, + Value: strconv.FormatInt(utils.CalculateReset(&descriptor.CurrentLimit.Unit, this.customHeaderClock).GetSeconds(), 10), + } +} + func (this *service) ShouldRateLimit( ctx context.Context, request *pb.RateLimitRequest, diff --git a/src/service/ratelimit_test.go b/src/service/ratelimit_test.go index bb0234d7..7d94580b 100644 --- a/src/service/ratelimit_test.go +++ b/src/service/ratelimit_test.go @@ -1,17 +1,23 @@ package ratelimit import ( + "context" + "os" + "sync" "testing" ratelimitv3 "github.com/envoyproxy/go-control-plane/envoy/extensions/common/ratelimit/v3" pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" "github.com/google/go-cmp/cmp" + gostats "github.com/lyft/gostats" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/testing/protocmp" "google.golang.org/protobuf/types/known/structpb" "github.com/envoyproxy/ratelimit/src/config" + mock_stats "github.com/envoyproxy/ratelimit/test/mocks/stats" ) func TestRatelimitToMetadata(t *testing.T) { @@ -221,3 +227,78 @@ func TestRatelimitToMetadata(t *testing.T) { }) } } + +// stubConfigEvent is a minimal ConfigUpdateEvent that always returns a no-op config +// with no error, allowing SetConfig to proceed to the settings-reading section. +type stubConfigEvent struct{} + +func (s stubConfigEvent) GetConfig() (config.RateLimitConfig, any) { + return &stubRLConfig{}, nil +} + +// stubRLConfig satisfies config.RateLimitConfig with no-op methods. +type stubRLConfig struct{} + +func (s *stubRLConfig) Dump() string { return "" } +func (s *stubRLConfig) GetLimit(_ context.Context, _ string, _ *ratelimitv3.RateLimitDescriptor) *config.RateLimit { + return nil +} +func (s *stubRLConfig) IsEmptyDomains() bool { return false } + +func newMinimalService(t *testing.T) *service { + t.Helper() + store := gostats.NewStore(gostats.NewNullSink(), false) + mgr := mock_stats.NewMockStatManager(store) + return &service{ + configLock: sync.RWMutex{}, + stats: mgr.NewServiceStats(), + } +} + +// TestResponseHeadersEnabledResetToFalseOnHotReload verifies that customHeadersEnabled +// is unconditionally assigned on every SetConfig call, not only set to true. +// Without the fix, the flag is only ever set to true and never cleared, so a +// hot-reload that turns off LIMIT_RESPONSE_HEADERS_ENABLED has no effect. +func TestResponseHeadersEnabledResetToFalseOnHotReload(t *testing.T) { + // Start with the feature enabled. + os.Setenv("LIMIT_RESPONSE_HEADERS_ENABLED", "true") + defer os.Unsetenv("LIMIT_RESPONSE_HEADERS_ENABLED") + + svc := newMinimalService(t) + + // First load: env var ON → flag must become true. + svc.SetConfig(stubConfigEvent{}, false) + assert.True(t, svc.customHeadersEnabled, "customHeadersEnabled should be true after initial load with env var on") + + // Operator turns off the env var; next hot-reload must clear the flag. + os.Unsetenv("LIMIT_RESPONSE_HEADERS_ENABLED") + svc.SetConfig(stubConfigEvent{}, false) + + // Bug: without the fix, this assertion fails because the if-only guard never + // resets the field to false. + assert.False(t, svc.customHeadersEnabled, "customHeadersEnabled should be false after reload with env var off") +} + +// TestRequestHeadersEnabledResetToFalseOnHotReload verifies that requestHeadersEnabled +// is unconditionally assigned on every SetConfig call, not only set to true. +// Without the fix, the flag is only ever set to true and never cleared, so a +// hot-reload that turns off LIMIT_REQUEST_HEADERS_ENABLED has no effect. +func TestRequestHeadersEnabledResetToFalseOnHotReload(t *testing.T) { + // Start with the feature enabled. + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + defer os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + + svc := newMinimalService(t) + + // First load: env var ON → flag must become true. + svc.SetConfig(stubConfigEvent{}, false) + assert.True(t, svc.requestHeadersEnabled, "requestHeadersEnabled should be true after initial load with env var on") + + // Operator turns off the env var; next hot-reload must clear the flag. + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + svc.SetConfig(stubConfigEvent{}, false) + + // Bug: without the fix, this assertion fails because the if-only guard never + // resets the field to false. + assert.False(t, svc.requestHeadersEnabled, "requestHeadersEnabled should be false after reload with env var off") +} diff --git a/src/settings/settings.go b/src/settings/settings.go index 2129cf80..ed28fd2a 100644 --- a/src/settings/settings.go +++ b/src/settings/settings.go @@ -120,6 +120,15 @@ type Settings struct { // value: remaining seconds HeaderRatelimitReset string `envconfig:"LIMIT_RESET_HEADER" default:"RateLimit-Reset"` + // Settings for optional injection of rate limit headers into the upstream request (request_headers_to_add) + RateLimitRequestHeadersEnabled bool `envconfig:"LIMIT_REQUEST_HEADERS_ENABLED" default:"false"` + // header name for the current limit value injected into the upstream request + HeaderRequestRatelimitLimit string `envconfig:"LIMIT_REQUEST_LIMIT_HEADER" default:"RateLimit-Limit"` + // header name for the remaining count injected into the upstream request + HeaderRequestRatelimitRemaining string `envconfig:"LIMIT_REQUEST_REMAINING_HEADER" default:"RateLimit-Remaining"` + // header name for the reset seconds injected into the upstream request + HeaderRequestRatelimitReset string `envconfig:"LIMIT_REQUEST_RESET_HEADER" default:"RateLimit-Reset"` + // 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..7cbc76c1 100644 --- a/test/service/ratelimit_test.go +++ b/test/service/ratelimit_test.go @@ -30,6 +30,7 @@ import ( "github.com/envoyproxy/ratelimit/src/redis" server "github.com/envoyproxy/ratelimit/src/server" ratelimit "github.com/envoyproxy/ratelimit/src/service" + "github.com/envoyproxy/ratelimit/src/settings" "github.com/envoyproxy/ratelimit/test/common" mock_config "github.com/envoyproxy/ratelimit/test/mocks/config" mock_limiter "github.com/envoyproxy/ratelimit/test/mocks/limiter" @@ -348,6 +349,81 @@ func TestMixedRuleShadowMode(test *testing.T) { t.assert.EqualValues(0, t.statStore.NewCounter("global_shadow_mode").Value()) } +func TestRequestHeadersSettingsDefaults(test *testing.T) { + s := settings.NewSettings() + assert.False(test, s.RateLimitRequestHeadersEnabled) + assert.Equal(test, "RateLimit-Limit", s.HeaderRequestRatelimitLimit) + assert.Equal(test, "RateLimit-Remaining", s.HeaderRequestRatelimitRemaining) + assert.Equal(test, "RateLimit-Reset", s.HeaderRequestRatelimitReset) +} + +func TestRequestHeadersDisabledByDefault(test *testing.T) { + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, + }) + + response, err := service.ShouldRateLimit(context.Background(), request) + t.assert.Nil(err) + t.assert.Nil(response.RequestHeadersToAdd) +} + +// TestRequestHeadersDisabledOnHotReload verifies that requestHeadersEnabled is reset to +// false on config hot-reload when LIMIT_REQUEST_HEADERS_ENABLED is turned off. Without +// the fix, the flag is only ever set to true and never cleared, so headers persist after +// the env var is removed. +func TestRequestHeadersDisabledOnHotReload(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() // service starts with requestHeadersEnabled = true + + // Now turn off the env var and trigger a config reload. + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + // After the reload with LIMIT_REQUEST_HEADERS_ENABLED unset, request headers must be absent. + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, + }) + + response, err := service.ShouldRateLimit(context.Background(), request) + t.assert.Nil(err) + t.assert.Nil(response.RequestHeadersToAdd) +} + func TestServiceWithCustomRatelimitHeaders(test *testing.T) { os.Setenv("LIMIT_RESPONSE_HEADERS_ENABLED", "true") os.Setenv("LIMIT_LIMIT_HEADER", "A-Ratelimit-Limit") @@ -460,6 +536,227 @@ func TestServiceWithDefaultRatelimitHeaders(test *testing.T) { t.assert.Nil(err) } +func TestServiceWithDefaultRequestHeaders(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}, {{"hello", "world"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + nil, + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.config.EXPECT().GetLimit(context.Background(), "different-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: nil, LimitRemaining: 0}, + }) + + 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: nil, LimitRemaining: 0}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "RateLimit-Limit", Value: "10"}, + {Key: "RateLimit-Remaining", Value: "0"}, + {Key: "RateLimit-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} + +func TestServiceWithCustomRequestHeaders(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + os.Setenv("LIMIT_REQUEST_LIMIT_HEADER", "X-RateLimit-Limit") + os.Setenv("LIMIT_REQUEST_REMAINING_HEADER", "X-RateLimit-Remaining") + os.Setenv("LIMIT_REQUEST_RESET_HEADER", "X-RateLimit-Reset") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + os.Unsetenv("LIMIT_REQUEST_LIMIT_HEADER") + os.Unsetenv("LIMIT_REQUEST_REMAINING_HEADER") + os.Unsetenv("LIMIT_REQUEST_RESET_HEADER") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}, {{"hello", "world"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + nil, + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.config.EXPECT().GetLimit(context.Background(), "different-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: nil, LimitRemaining: 0}, + }) + + 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: nil, LimitRemaining: 0}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "X-RateLimit-Limit", Value: "10"}, + {Key: "X-RateLimit-Remaining", Value: "0"}, + {Key: "X-RateLimit-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} + +func TestServiceWithRequestHeadersWithinLimit(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 8}, + }) + + 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: 8}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "RateLimit-Limit", Value: "10"}, + {Key: "RateLimit-Remaining", Value: "8"}, + {Key: "RateLimit-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} + +func TestServiceWithBothRequestAndResponseHeaders(test *testing.T) { + os.Setenv("LIMIT_REQUEST_HEADERS_ENABLED", "true") + os.Setenv("LIMIT_REQUEST_LIMIT_HEADER", "X-Upstream-Limit") + os.Setenv("LIMIT_REQUEST_REMAINING_HEADER", "X-Upstream-Remaining") + os.Setenv("LIMIT_REQUEST_RESET_HEADER", "X-Upstream-Reset") + os.Setenv("LIMIT_RESPONSE_HEADERS_ENABLED", "true") + os.Setenv("LIMIT_LIMIT_HEADER", "X-Downstream-Limit") + os.Setenv("LIMIT_REMAINING_HEADER", "X-Downstream-Remaining") + os.Setenv("LIMIT_RESET_HEADER", "X-Downstream-Reset") + defer func() { + os.Unsetenv("LIMIT_REQUEST_HEADERS_ENABLED") + os.Unsetenv("LIMIT_REQUEST_LIMIT_HEADER") + os.Unsetenv("LIMIT_REQUEST_REMAINING_HEADER") + os.Unsetenv("LIMIT_REQUEST_RESET_HEADER") + os.Unsetenv("LIMIT_RESPONSE_HEADERS_ENABLED") + os.Unsetenv("LIMIT_LIMIT_HEADER") + os.Unsetenv("LIMIT_REMAINING_HEADER") + os.Unsetenv("LIMIT_RESET_HEADER") + }() + + t := commonSetup(test) + defer t.controller.Finish() + service := t.setupBasicService() + + barrier := newBarrier() + t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) { + barrier.signal() + return t.config, nil + }) + t.configUpdateEventChan <- t.configUpdateEvent + barrier.wait() + + request := common.NewRateLimitRequest( + "different-domain", [][][2]string{{{"foo", "bar"}}}, 1) + limits := []*config.RateLimit{ + config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false), + } + t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0]) + t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return( + []*pb.RateLimitResponse_DescriptorStatus{ + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, + }) + + 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}, + }, + RequestHeadersToAdd: []*core.HeaderValue{ + {Key: "X-Upstream-Limit", Value: "10"}, + {Key: "X-Upstream-Remaining", Value: "0"}, + {Key: "X-Upstream-Reset", Value: "58"}, + }, + ResponseHeadersToAdd: []*core.HeaderValue{ + {Key: "X-Downstream-Limit", Value: "10"}, + {Key: "X-Downstream-Remaining", Value: "0"}, + {Key: "X-Downstream-Reset", Value: "58"}, + }, + }, + response) + t.assert.Nil(err) +} + func TestEmptyDomain(test *testing.T) { t := commonSetup(test) defer t.controller.Finish()