Skip to content

feat: rate limit config per key [WIP] #2089

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

Closed
SkArchon wants to merge 1 commit into from
Closed

feat: rate limit config per key [WIP] #2089

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
8 changes: 5 additions & 3 deletions router/core/graph_server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1315,10 +1315,12 @@ func (s *graphServer) buildGraphMux(
if s.redisClient != nil {
handlerOpts.RateLimitConfig = s.rateLimit
handlerOpts.RateLimiter, err = NewCosmoRateLimiter(&CosmoRateLimiterOptions{
RedisClient: s.redisClient,
Debug: s.rateLimit.Debug,
RejectStatusCode: s.rateLimit.SimpleStrategy.RejectStatusCode,
RedisClient: s.redisClient,
Debug: s.rateLimit.Debug,
//RejectStatusCode: s.rateLimit.SimpleStrategy.RejectStatusCode,
KeySuffixExpression: s.rateLimit.KeySuffixExpression,
RateLimitConfig: s.rateLimit.SimpleStrategy,
BaseRateLimitKey: s.rateLimit.Storage.KeyPrefix,
ExprManager: exprManager,
})
if err != nil {
Expand Down
10 changes: 3 additions & 7 deletions router/core/graphql_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,12 +269,8 @@ func (h *GraphQLHandler) configureRateLimiting(ctx *resolve.Context) *resolve.Co
ctx.SetRateLimiter(h.rateLimiter)
ctx.RateLimitOptions = resolve.RateLimitOptions{
Enable: true,
IncludeStatsInResponseExtension: !h.rateLimitConfig.SimpleStrategy.HideStatsFromResponseExtension,
Rate: h.rateLimitConfig.SimpleStrategy.Rate,
Burst: h.rateLimitConfig.SimpleStrategy.Burst,
Period: h.rateLimitConfig.SimpleStrategy.Period,
RateLimitKey: h.rateLimitConfig.Storage.KeyPrefix,
RejectExceedingRequests: h.rateLimitConfig.SimpleStrategy.RejectExceedingRequests,
IncludeStatsInResponseExtension: !h.rateLimitConfig.HideStatsFromResponseExtension,
//RateLimitKey: h.rateLimitConfig.Storage.KeyPrefix,
ErrorExtensionCode: resolve.RateLimitErrorExtensionCode{
Enabled: h.rateLimitConfig.ErrorExtensionCode.Enabled,
Code: h.rateLimitConfig.ErrorExtensionCode.Code,
Expand Down Expand Up @@ -317,7 +313,7 @@ func (h *GraphQLHandler) WriteError(ctx *resolve.Context, err error, res *resolv
Code: h.rateLimitConfig.ErrorExtensionCode.Code,
}
}
if !h.rateLimitConfig.SimpleStrategy.HideStatsFromResponseExtension {
if !h.rateLimitConfig.HideStatsFromResponseExtension {
buf := bytes.NewBuffer(make([]byte, 0, 1024))
err = h.rateLimiter.RenderResponseExtension(ctx, buf)
if err != nil {
Expand Down
62 changes: 44 additions & 18 deletions router/core/ratelimiter.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"errors"
"fmt"
rd "github.com/wundergraph/cosmo/router/internal/persistedoperation/operationstorage/redis"
"github.com/wundergraph/cosmo/router/pkg/config"
"io"
"reflect"
"sync"
Expand All @@ -25,10 +26,12 @@ type CosmoRateLimiterOptions struct {
RedisClient rd.RDCloser
Debug bool

RejectStatusCode int

KeySuffixExpression string
ExprManager *expr.Manager

BaseRateLimitKey string

RateLimitConfig config.RateLimitSimpleStrategy
}

func NewCosmoRateLimiter(opts *CosmoRateLimiterOptions) (rl *CosmoRateLimiter, err error) {
Expand All @@ -37,7 +40,8 @@ func NewCosmoRateLimiter(opts *CosmoRateLimiterOptions) (rl *CosmoRateLimiter, e
client: opts.RedisClient,
limiter: limiter,
debug: opts.Debug,
rejectStatusCode: opts.RejectStatusCode,
baseRateLimitKey: opts.BaseRateLimitKey,
rateLimitConfig: opts.RateLimitConfig,
}
if rl.rejectStatusCode == 0 {
rl.rejectStatusCode = 200
Expand All @@ -59,22 +63,41 @@ type CosmoRateLimiter struct {
rejectStatusCode int

keySuffixProgram *vm.Program

// TODO: To decouple from the config
rateLimitConfig config.RateLimitSimpleStrategy
baseRateLimitKey string
}

func (c *CosmoRateLimiter) RateLimitPreFetch(ctx *resolve.Context, info *resolve.FetchInfo, input json.RawMessage) (result *resolve.RateLimitDeny, err error) {
if c.isIntrospectionQuery(info.RootFields) {
return nil, nil
}
requestRate := c.calculateRate()
limit := redis_rate.Limit{
Rate: ctx.RateLimitOptions.Rate,
Burst: ctx.RateLimitOptions.Burst,
Period: ctx.RateLimitOptions.Period,
}
key, err := c.generateKey(ctx)
rawExprKey, key, err := c.generateKey(ctx, info)
if err != nil {
return nil, err
}

limitDetails, found := c.rateLimitConfig.KeyMapping[rawExprKey]
if !found {
if !c.rateLimitConfig.Enabled {
return nil, nil
} else {
limitDetails = c.rateLimitConfig.RateLimitSimpleStrategyEntry
}
} else {
if !limitDetails.Enabled {
return nil, nil
}
}

limit := redis_rate.Limit{
Rate: limitDetails.Rate,
Burst: limitDetails.Burst,
Period: limitDetails.Period,
}

allow, err := c.limiter.AllowN(ctx.Context(), key, limit, requestRate)
if err != nil {
return nil, err
Expand All @@ -83,29 +106,32 @@ func (c *CosmoRateLimiter) RateLimitPreFetch(ctx *resolve.Context, info *resolve
if allow.Allowed >= requestRate {
return nil, nil
}
if ctx.RateLimitOptions.RejectExceedingRequests {
if limitDetails.RejectExceedingRequests {
return nil, ErrRateLimitExceeded
}
return &resolve.RateLimitDeny{}, nil
}

func (c *CosmoRateLimiter) generateKey(ctx *resolve.Context) (string, error) {
func (c *CosmoRateLimiter) generateKey(ctx *resolve.Context, info *resolve.FetchInfo) (string, string, error) {
if c.keySuffixProgram == nil {
return ctx.RateLimitOptions.RateLimitKey, nil
return "", c.baseRateLimitKey, nil
}
rc := getRequestContext(ctx.Context())
if rc == nil {
return "", errors.New("no request context")
return "", "", errors.New("no request context")
}
str, err := expr.ResolveStringExpression(c.keySuffixProgram, rc.expressionContext)
clonedEc := rc.expressionContext.Clone()
clonedEc.Subgraph.Id = info.DataSourceID
clonedEc.Subgraph.Name = info.DataSourceName
str, err := expr.ResolveStringExpression(c.keySuffixProgram, *clonedEc)
if err != nil {
return "", fmt.Errorf("failed to resolve key suffix expression: %w", err)
return "", "", fmt.Errorf("failed to resolve key suffix expression: %w", err)
}
buf := bytes.NewBuffer(make([]byte, 0, len(ctx.RateLimitOptions.RateLimitKey)+len(str)+1))
_, _ = buf.WriteString(ctx.RateLimitOptions.RateLimitKey)
buf := bytes.NewBuffer(make([]byte, 0, len(c.baseRateLimitKey)+len(str)+1))
_, _ = buf.WriteString(c.baseRateLimitKey)
_ = buf.WriteByte(':')
_, _ = buf.WriteString(str)
return buf.String(), nil
return str, buf.String(), nil
}

func (c *CosmoRateLimiter) RejectStatusCode() int {
Expand Down
30 changes: 8 additions & 22 deletions router/core/ratelimiter_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
t.Parallel()
rl, err := NewCosmoRateLimiter(&CosmoRateLimiterOptions{})
assert.NoError(t, err)
key, err := rl.generateKey(expressionResolveContext(t, nil, nil))
key, err := rl.generateKey(expressionResolveContext(t, nil, nil), nil)
assert.NoError(t, err)
assert.Equal(t, "test", key)
})
Expand All @@ -51,9 +51,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
ExprManager: expr.CreateNewExprManager(),
})
require.NoError(t, err)
key, err := rl.generateKey(
expressionResolveContext(t, http.Header{"Authorization": []string{"token"}}, nil),
)
key, err := rl.generateKey(expressionResolveContext(t, http.Header{"Authorization": []string{"token"}}, nil), nil)
assert.NoError(t, err)
assert.Equal(t, "test:token", key)
})
Expand All @@ -64,9 +62,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
ExprManager: expr.CreateNewExprManager(),
})
assert.NoError(t, err)
key, err := rl.generateKey(
expressionResolveContext(t, http.Header{"Authorization": []string{"123"}}, nil),
)
key, err := rl.generateKey(expressionResolveContext(t, http.Header{"Authorization": []string{"123"}}, nil), nil)
assert.NoError(t, err)
assert.Equal(t, "test:123", key)
})
Expand All @@ -77,9 +73,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
ExprManager: expr.CreateNewExprManager(),
})
assert.NoError(t, err)
key, err := rl.generateKey(
expressionResolveContext(t, http.Header{"Authorization": []string{" token "}}, nil),
)
key, err := rl.generateKey(expressionResolveContext(t, http.Header{"Authorization": []string{" token "}}, nil), nil)
assert.NoError(t, err)
assert.Equal(t, "test:token", key)
})
Expand All @@ -90,9 +84,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
ExprManager: expr.CreateNewExprManager(),
})
assert.NoError(t, err)
key, err := rl.generateKey(
expressionResolveContext(t, nil, map[string]any{"sub": "token"}),
)
key, err := rl.generateKey(expressionResolveContext(t, nil, map[string]any{"sub": "token"}), nil)
assert.NoError(t, err)
assert.Equal(t, "test:token", key)
})
Expand All @@ -103,9 +95,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
ExprManager: expr.CreateNewExprManager(),
})
assert.NoError(t, err)
key, err := rl.generateKey(
expressionResolveContext(t, nil, map[string]any{"sub": 123}),
)
key, err := rl.generateKey(expressionResolveContext(t, nil, map[string]any{"sub": 123}), nil)
assert.Error(t, err)
assert.Empty(t, key)
})
Expand All @@ -116,9 +106,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
ExprManager: expr.CreateNewExprManager(),
})
assert.NoError(t, err)
key, err := rl.generateKey(
expressionResolveContext(t, http.Header{"X-Forwarded-For": []string{"192.168.0.1"}}, map[string]any{"sub": "token"}),
)
key, err := rl.generateKey(expressionResolveContext(t, http.Header{"X-Forwarded-For": []string{"192.168.0.1"}}, map[string]any{"sub": "token"}), nil)
assert.NoError(t, err)
assert.Equal(t, "test:token", key)
})
Expand All @@ -129,9 +117,7 @@ func TestRateLimiterGenerateKey(t *testing.T) {
ExprManager: expr.CreateNewExprManager(),
})
assert.NoError(t, err)
key, err := rl.generateKey(
expressionResolveContext(t, http.Header{"X-Forwarded-For": []string{"192.168.0.1"}}, nil),
)
key, err := rl.generateKey(expressionResolveContext(t, http.Header{"X-Forwarded-For": []string{"192.168.0.1"}}, nil), nil)
assert.NoError(t, err)
assert.Equal(t, "test:192.168.0.1", key)
})
Expand Down
2 changes: 1 addition & 1 deletion router/pkg/authentication/jwks_token_decoder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ func NewJwksTokenDecoder(ctx context.Context, logger *zap.Logger, configs []JWKS
c.URL: store,
},
PrioritizeHTTP: true,
RefreshUnknownKID: rate.NewLimiter(rate.Every(5*time.Minute), 1),
RefreshUnknownKID: rate.NewLimiter(rate.Every(2*time.Second), 1),
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot Jul 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Verify the dramatic increase in JWKS refresh frequency.

The refresh rate for unknown Key IDs has been increased from once every 5 minutes to once every 2 seconds - a 150x increase in frequency. This could significantly impact performance and potentially overwhelm JWKS providers with requests.

Please verify this change is intentional and consider the following:

  • Will external JWKS providers handle this increased request frequency?
  • Has this been tested with production-like traffic patterns?
  • Is there documentation explaining why this aggressive refresh rate is necessary?

🏁 Script executed:

#!/bin/bash
# Description: Search for any configuration or documentation explaining the JWKS refresh rate change
# Expected: Find configuration options, comments, or documentation explaining the rationale

# Search for JWKS refresh rate configurations or comments
rg -A 5 -B 5 "RefreshUnknownKID|JWKS.*refresh|refresh.*JWKS" --type go

Length of output: 1639

Document and validate JWKS unknown‐KID refresh rate change

The hard-coded RefreshUnknownKID limiter was lowered from 5 minutes to 2 seconds (150× increase), but there’s no accompanying comment or configuration entry explaining why. This could lead to a surge of HTTP requests against JWKS endpoints.

Please address the following in router/pkg/authentication/jwks_token_decoder.go:99:

  • Confirm that external JWKS providers can sustain refresh attempts every 2 seconds under production load.
  • Add a code comment or expose this interval via configuration (e.g., in authentication.JWKSConfig) with a clear rationale.
  • Update any related documentation (README or design docs) to explain why such an aggressive refresh rate is necessary.
🤖 Prompt for AI Agents 
In router/pkg/authentication/jwks_token_decoder.go at line 99, the
RefreshUnknownKID rate limiter was changed from 5 minutes to 2 seconds without
explanation. Confirm with the team or documentation that external JWKS providers
can handle refreshes every 2 seconds in production. Add a clear code comment
near this line explaining the reason for this aggressive refresh rate.
Additionally, expose this interval as a configurable parameter in
authentication.JWKSConfig instead of hardcoding it. Finally, update any relevant
documentation or README files to describe the rationale behind this change and
its impact.

}

jwks, err := createKeyFunc(ctx, jwksetHTTPClientOptions)
Expand Down
27 changes: 21 additions & 6 deletions router/pkg/config/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -507,6 +507,10 @@ type RateLimitConfiguration struct {
Debug bool `yaml:"debug" envDefault:"false" env:"RATE_LIMIT_DEBUG"`
KeySuffixExpression string `yaml:"key_suffix_expression,omitempty" env:"RATE_LIMIT_KEY_SUFFIX_EXPRESSION"`
ErrorExtensionCode RateLimitErrorExtensionCode `yaml:"error_extension_code"`

// This makes more sense here since it's a single value for the whole rate limit configuration
// Since we cannot have this per subgraph
HideStatsFromResponseExtension bool `yaml:"hide_stats_from_response_extension" envDefault:"false" env:"RATE_LIMIT_SIMPLE_HIDE_STATS_FROM_RESPONSE_EXTENSION"`
}

type RateLimitErrorExtensionCode struct {
Expand All @@ -521,12 +525,23 @@ type RedisConfiguration struct {
}

type RateLimitSimpleStrategy struct {
Rate int `yaml:"rate" envDefault:"10" env:"RATE_LIMIT_SIMPLE_RATE"`
Burst int `yaml:"burst" envDefault:"10" env:"RATE_LIMIT_SIMPLE_BURST"`
Period time.Duration `yaml:"period" envDefault:"1s" env:"RATE_LIMIT_SIMPLE_PERIOD"`
RejectExceedingRequests bool `yaml:"reject_exceeding_requests" envDefault:"false" env:"RATE_LIMIT_SIMPLE_REJECT_EXCEEDING_REQUESTS"`
RejectStatusCode int `yaml:"reject_status_code" envDefault:"200" env:"RATE_LIMIT_SIMPLE_REJECT_STATUS_CODE"`
HideStatsFromResponseExtension bool `yaml:"hide_stats_from_response_extension" envDefault:"false" env:"RATE_LIMIT_SIMPLE_HIDE_STATS_FROM_RESPONSE_EXTENSION"`
RateLimitSimpleStrategyEntry
//Enabled bool `yaml:"enabled" envDefault:"true" env:"RATE_LIMIT_SIMPLE_ENABLED"`
//Rate int `yaml:"rate" envDefault:"10" env:"RATE_LIMIT_SIMPLE_RATE"`
//Burst int `yaml:"burst" envDefault:"10" env:"RATE_LIMIT_SIMPLE_BURST"`
//Period time.Duration `yaml:"period" envDefault:"1s" env:"RATE_LIMIT_SIMPLE_PERIOD"`
//RejectExceedingRequests bool `yaml:"reject_exceeding_requests" envDefault:"false" env:"RATE_LIMIT_SIMPLE_REJECT_EXCEEDING_REQUESTS"`
//RejectStatusCode int `yaml:"reject_status_code" envDefault:"200" env:"RATE_LIMIT_SIMPLE_REJECT_STATUS_CODE"`
KeyMapping map[string]RateLimitSimpleStrategyEntry `yaml:"key_mapping"`
}

type RateLimitSimpleStrategyEntry struct {
Enabled bool `yaml:"enabled" envDefault:"true" env:"RATE_LIMIT_SIMPLE_ENABLED"`
Rate int `yaml:"rate" envDefault:"10" env:"RATE_LIMIT_SIMPLE_RATE"`
Burst int `yaml:"burst" envDefault:"10" env:"RATE_LIMIT_SIMPLE_BURST"`
Period time.Duration `yaml:"period" envDefault:"1s" env:"RATE_LIMIT_SIMPLE_PERIOD"`
RejectExceedingRequests bool `yaml:"reject_exceeding_requests" envDefault:"false" env:"RATE_LIMIT_SIMPLE_REJECT_EXCEEDING_REQUESTS"`
RejectStatusCode int `yaml:"reject_status_code" envDefault:"200" env:"RATE_LIMIT_SIMPLE_REJECT_STATUS_CODE"`
}

type CDNConfiguration struct {
Expand Down
41 changes: 41 additions & 0 deletions router/pkg/config/config.schema.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1820,6 +1820,47 @@
"type": "boolean",
"default": false,
"description": "Hide the rate limit stats from the response extension. If the value is true, the rate limit stats are not included in the response extension."
},
"key_mapping": {
"type": "object",
"description": "The configuration per key entry.",
"additionalProperties": {
Comment on lines +1823 to +1827
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot Jul 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add basic constraints to key_mapping to avoid silent configuration errors

key_mapping is currently declared as a totally free-form object. A typo such as "raet": 50 or an empty object {} will happily pass validation but later explode at runtime.
Consider tightening validation with e.g. patternProperties, propertyNames, or at least a minProperties: 1 to ensure that every entry contains something meaningful.

 "key_mapping": {
   "type": "object",
   "description": "The configuration per key entry.",
+  "minProperties": 1,
   "additionalProperties": {
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
},
"key_mapping": {
"type": "object",
"description": "The configuration per key entry.",
"additionalProperties": {
},
"key_mapping": {
"type": "object",
"description": "The configuration per key entry.",
"minProperties": 1,
"additionalProperties": {
🤖 Prompt for AI Agents 
In router/pkg/config/config.schema.json around lines 1823 to 1827, the
key_mapping object lacks constraints, allowing empty or misspelled entries to
pass validation and cause runtime errors. Add validation rules such as
patternProperties or propertyNames to restrict allowed keys, and include
minProperties: 1 to ensure the object is not empty, thereby preventing silent
configuration errors.

"type": "object",
"additionalProperties": false,
"description": "The configuration for all subgraphs. The configuration is used to configure the traffic shaping for all subgraphs.",
"properties": {
"enabled": {
"type": "boolean",
"description": "The rate at which the requests are allowed. The rate is specified as a number of requests per second."
},
"rate": {
"type": "integer",
"description": "The rate at which the requests are allowed. The rate is specified as a number of requests per second.",
"minimum": 1
},
"burst": {
"type": "integer",
"description": "The maximum number of requests that are allowed to exceed the rate. The burst is specified as a number of requests.",
"minimum": 1
},
"period": {
"type": "string",
"description": "The period of time over which the rate limit is enforced. The period is specified as a string with a number and a unit, e.g. 10ms, 1s, 1m, 1h. The supported units are 'ms', 's', 'm', 'h'.",
"duration": {
"minimum": "1s"
}
},
"reject_exceeding_requests": {
"type": "boolean",
"description": "Reject the requests that exceed the rate limit. If the value is true, the requests that exceed the rate limit are rejected."
},
"reject_status_code": {
"type": "integer",
"description": "The status code to return when the request is rejected. The default value is 200 (OK) as we're returning a well formed GraphQL response.",
"default": 200
}
}
Comment on lines +1836 to +1862
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot Jul 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Require at least one rate-limiting parameter inside each mapping entry

All properties inside a single entry are optional, so a definition like:

rate_limit:
  simple_strategy:
    key_mapping:
      employees: {}

would pass schema validation yet leave the router with undefined values at runtime.
At minimum, enforce the same trio as the global defaults:

 "additionalProperties": false,
 ...
-"description": "The configuration for all subgraphs. The configuration is used to configure the traffic shaping for all subgraphs.",
+"description": "Per-key rate-limit settings. At least one of rate | burst | period must be supplied.",
 ...
-                "properties": {
+                "properties": {
 ...
                 },
+                "anyOf": [
+                  { "required": ["rate"] },
+                  { "required": ["burst"] },
+                  { "required": ["period"] }
+                ]

Alternatively, replicate the global required: ["rate","burst","period"] if you don’t want to inherit omitted values.

🤖 Prompt for AI Agents 
In router/pkg/config/config.schema.json around lines 1836 to 1862, the
rate-limiting entries currently have all properties optional, allowing empty
objects that cause undefined runtime values. To fix this, add a "required" field
to each rate-limiting entry schema that enforces at least the "rate", "burst",
and "period" properties to be present, matching the global defaults. This
ensures each entry has the necessary parameters defined and prevents invalid
configurations.

}
}
},
"required": ["rate", "burst", "period"]
Expand Down
Loading