feat: allow subscriptions to be excluded from rate limiting

demo/pkg/subgraphs/courses/generated/mapping.json

@@ -205,4 +205,4 @@
       ]
     }
   ]
-}
+}

docs-website/router/configuration.mdx

@@ -2221,6 +2221,7 @@ For mor information on how to use the expression language, please refer to the [
 | Rate Limiting Simple Strategy](/router/configuration#rate-limiting-simple-strategy)      | simple                |                                               |                                                                                                                                                                                                                                                                                        |               |
 |                                                                                          | storage               | <Icon icon="square" />                        | [Rate Limiting Redis Storage](/router/configuration#rate-limiting-redis-storage)                                                                                                                                                                                                       |               |
 | RATE_LIMIT_KEY_SUFFIX_EXPRESSION                                                         | key_suffix_expression | <Icon icon="square" />                        | The expression to define a key suffix for the rate limit, e.g. by using request headers, claims, or a combination of both with a fallback strategy. The expression is specified as a string and needs to evaluate to a string. Please see https://expr-lang.org/ for more information. |               |
+| RATE_LIMIT_EXCLUDE_SUBSCRIPTIONS                                                         | exclude_subscriptions | <Icon icon="square" />                        | Exclude subscription operations from rate limiting. Queries and mutations remain rate limited, including those sent over subscription transports like WebSockets.                                                                                                                     | false         |
 |                                                                                          | error_extension_code  | <Icon icon="square" />                        |                                                                                                                                                                                                                                                                                        |               |
 | [Rate Limit Error Extension Code](/router/configuration#rate-limit-error-extension-code) |                       |                                               |                                                                                                                                                                                                                                                                                        |               |
 
@@ -2279,11 +2280,16 @@ rate_limit:
         key_prefix: "cosmo_rate_limit"
     debug: false
     key_suffix_expression: "request.header.Get('X-Api-Key')"
+    exclude_subscriptions: false
     error_extension_code:
         enabled: true
         code: "RATE_LIMIT_EXCEEDED"
 ```
 
+<Warning>
+  Excluding subscriptions from rate limiting can expose the router and your subgraphs to a denial-of-service risk, as an unrestricted client can open subscriptions that trigger unbounded subgraph fetches. Prefer a `key_suffix_expression` based on request headers where possible.
+</Warning>
+
 ## Subgraph Data Propagation
 
 Controls what subgraphs can propagate back to the client. Errors and the response `extensions` object are governed by separate but related policies.

router-tests/security/ratelimit_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -1058,6 +1059,240 @@ func TestRateLimit(t *testing.T) {
 	})
 }
 
+func TestRateLimitExcludeSubscriptions(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode.")
+	}
+
+	t.Parallel()
+
+	const subscriptionQuery = `{"query":"subscription { employeeUpdated(employeeID: 3) { id details { forename surname } }}"}`
+	employeeUpdatedEvent := []byte(`{"id":3,"__typename": "Employee"}`)
+
+	newRateLimitConfig := func(key string, excludeSubscriptions, rejectExceedingRequests bool) *config.RateLimitConfiguration {
+		return &config.RateLimitConfiguration{
+			Enabled:  true,
+			Strategy: "simple",
+			SimpleStrategy: config.RateLimitSimpleStrategy{
+				Rate:                    1,
+				Burst:                   1,
+				Period:                  time.Second * 10,
+				RejectExceedingRequests: rejectExceedingRequests,
+			},
+			Storage: config.RedisConfiguration{
+				URLs:      []string{"redis://localhost:6379"},
+				KeyPrefix: key,
+			},
+			Debug:                true,
+			ExcludeSubscriptions: excludeSubscriptions,
+		}
+	}
+
+	cleanupKey := func(t *testing.T, key string) {
+		t.Cleanup(func() {
+			client := redis.NewClient(&redis.Options{Addr: "localhost:6379", Password: "test"})
+			del := client.Del(context.Background(), key)
+			require.NoError(t, del.Err())
+		})
+	}
+
+	t.Run("subscription events are rate limited when exclude_subscriptions is disabled", func(t *testing.T) {
+		t.Parallel()
+
+		key := uuid.New().String()
+		cleanupKey(t, key)
+
+		testenv.Run(t, &testenv.Config{
+			RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate,
+			EnableNats:               true,
+			RouterOptions: []core.Option{
+				core.WithRateLimitConfig(newRateLimitConfig(key, false, false)),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, nil)
+			err := testenv.WSWriteJSON(t, conn, testenv.WebSocketMessage{
+				ID:      "1",
+				Type:    "subscribe",
+				Payload: []byte(subscriptionQuery),
+			})
+			require.NoError(t, err)
+
+			xEnv.WaitForSubscriptionCount(1, time.Second*15)
+			xEnv.WaitForTriggerCount(1, time.Second*15)
+
+			subject := xEnv.GetPubSubName("employeeUpdated.3")
+			xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, subject, employeeUpdatedEvent, 1, time.Second*15)
+
+			// The first event consumes the whole budget with the entity fetch resolving the details.
+			var res testenv.WebSocketMessage
+			err = testenv.WSReadJSON(t, conn, &res)
+			require.NoError(t, err)
+			require.Equal(t, "next", res.Type)
+			require.Equal(t, "1", res.ID)
+			require.Equal(t, fmt.Sprintf(`{"data":{"employeeUpdated":{"id":3,"details":{"forename":"Stefan","surname":"Avram"}}},"extensions":{"rateLimit":{"key":"%s","requestRate":1,"remaining":0,"retryAfterMs":1234,"resetAfterMs":1234}}}`, key), string(res.Payload))
+
+			// The entity fetch of the second event must be denied.
+			err = xEnv.NatsConnectionDefault.Publish(subject, employeeUpdatedEvent)
+			require.NoError(t, err)
+			require.NoError(t, xEnv.NatsConnectionDefault.Flush())
+
+			err = testenv.WSReadJSON(t, conn, &res)
+			require.NoError(t, err)
+			require.Equal(t, "next", res.Type)
+			require.Equal(t, "1", res.ID)
+			require.Contains(t, string(res.Payload), "Rate limit exceeded")
+
+			require.NoError(t, conn.Close())
+			xEnv.WaitForSubscriptionCount(0, time.Second*15)
+		})
+	})
+	t.Run("subscription events over websocket are not rate limited when exclude_subscriptions is enabled", func(t *testing.T) {
+		t.Parallel()
+
+		key := uuid.New().String()
+		cleanupKey(t, key)
+
+		testenv.Run(t, &testenv.Config{
+			RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate,
+			EnableNats:               true,
+			RouterOptions: []core.Option{
+				core.WithRateLimitConfig(newRateLimitConfig(key, true, true)),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, nil)
+			err := testenv.WSWriteJSON(t, conn, testenv.WebSocketMessage{
+				ID:      "1",
+				Type:    "subscribe",
+				Payload: []byte(subscriptionQuery),
+			})
+			require.NoError(t, err)
+
+			xEnv.WaitForSubscriptionCount(1, time.Second*15)
+			xEnv.WaitForTriggerCount(1, time.Second*15)
+
+			subject := xEnv.GetPubSubName("employeeUpdated.3")
+
+			// With a budget of 1, any rate limiting of the entity fetches would deny the second and third event.
+			for range 3 {
+				xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, subject, employeeUpdatedEvent, 1, time.Second*15)
+
+				var res testenv.WebSocketMessage
+				err = testenv.WSReadJSON(t, conn, &res)
+				require.NoError(t, err)
+				require.Equal(t, "next", res.Type)
+				require.Equal(t, "1", res.ID)
+				require.Equal(t, `{"data":{"employeeUpdated":{"id":3,"details":{"forename":"Stefan","surname":"Avram"}}}}`, string(res.Payload))
+			}
+
+			// Queries remain rate limited.
+			res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query:     `query ($n:Int!) { employee(id:$n) { id details { forename surname } } }`,
+				Variables: json.RawMessage(`{"n":1}`),
+			})
+			require.Equal(t, fmt.Sprintf(`{"data":{"employee":{"id":1,"details":{"forename":"Jens","surname":"Neuse"}}},"extensions":{"rateLimit":{"key":"%s","requestRate":1,"remaining":0,"retryAfterMs":1234,"resetAfterMs":1234}}}`, key), res.Body)
+			res = xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+				Query:     `query ($n:Int!) { employee(id:$n) { id details { forename surname } } }`,
+				Variables: json.RawMessage(`{"n":1}`),
+			})
+			require.Equal(t, fmt.Sprintf(`{"errors":[{"message":"Rate limit exceeded"}],"data":null,"extensions":{"rateLimit":{"key":"%s","requestRate":1,"remaining":0,"retryAfterMs":1234,"resetAfterMs":1234}}}`, key), res.Body)
+
+			require.NoError(t, conn.Close())
+			xEnv.WaitForSubscriptionCount(0, time.Second*15)
+		})
+	})
+	t.Run("queries over websocket are still rate limited when exclude_subscriptions is enabled", func(t *testing.T) {
+		t.Parallel()
+
+		key := uuid.New().String()
+		cleanupKey(t, key)
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithRateLimitConfig(newRateLimitConfig(key, true, true)),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, nil)
+			err := testenv.WSWriteJSON(t, conn, testenv.WebSocketMessage{
+				ID:      "1",
+				Type:    "subscribe",
+				Payload: []byte(`{"query":"query { employee(id:1) { id } }"}`),
+			})
+			require.NoError(t, err)
+
+			var res testenv.WebSocketMessage
+			err = testenv.WSReadJSON(t, conn, &res)
+			require.NoError(t, err)
+			require.Equal(t, "next", res.Type)
+			require.Equal(t, "1", res.ID)
+			require.Equal(t, fmt.Sprintf(`{"data":{"employee":{"id":1}},"extensions":{"rateLimit":{"key":"%s","requestRate":1,"remaining":0,"retryAfterMs":1234,"resetAfterMs":1234}}}`, key), string(res.Payload))
+
+			err = testenv.WSReadJSON(t, conn, &res)
+			require.NoError(t, err)
+			require.Equal(t, "complete", res.Type)
+			require.Equal(t, "1", res.ID)
+
+			err = testenv.WSWriteJSON(t, conn, testenv.WebSocketMessage{
+				ID:      "2",
+				Type:    "subscribe",
+				Payload: []byte(`{"query":"query { employee(id:1) { id } }"}`),
+			})
+			require.NoError(t, err)
+
+			err = testenv.WSReadJSON(t, conn, &res)
+			require.NoError(t, err)
+			require.Equal(t, "2", res.ID)
+			require.Contains(t, string(res.Payload), "Rate limit exceeded")
+
+			require.NoError(t, conn.Close())
+		})
+	})
+	t.Run("subscription events over sse are not rate limited when exclude_subscriptions is enabled", func(t *testing.T) {
+		t.Parallel()
+
+		key := uuid.New().String()
+		cleanupKey(t, key)
+
+		testenv.Run(t, &testenv.Config{
+			RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate,
+			EnableNats:               true,
+			RouterOptions: []core.Option{
+				core.WithRateLimitConfig(newRateLimitConfig(key, true, true)),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+			defer cancel()
+
+			events := make(chan string, 8)
+			go xEnv.GraphQLSubscriptionOverSSE(ctx, testenv.GraphQLRequest{
+				Query: `subscription { employeeUpdated(employeeID: 3) { id details { forename surname } }}`,
+				Header: map[string][]string{
+					"Content-Type": {"application/json"},
+					"Accept":       {"text/event-stream"},
+				},
+			}, func(data string) {
+				events <- data
+			})
+
+			xEnv.WaitForSubscriptionCount(1, time.Second*15)
+			xEnv.WaitForTriggerCount(1, time.Second*15)
+
+			subject := xEnv.GetPubSubName("employeeUpdated.3")
+
+			// With a budget of 1, any rate limiting of the entity fetches would deny the second and third event.
+			for range 3 {
+				xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, subject, employeeUpdatedEvent, 1, time.Second*15)
+
+				select {
+				case data := <-events:
+					require.Equal(t, `{"data":{"employeeUpdated":{"id":3,"details":{"forename":"Stefan","surname":"Avram"}}}}`, strings.TrimSpace(data))
+				case <-ctx.Done():
+					t.Fatal("timed out waiting for subscription event")
+				}
+			}
+		})
+	})
+}
+
 const (
 	bigNestedQuery = `query Demo {
   products {

router/core/graphql_handler.go

@@ -187,7 +187,7 @@ func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if h.engineLoaderHooks != nil {
 		resolveCtx.SetEngineLoaderHooks(h.engineLoaderHooks)
 	}
-	resolveCtx = h.configureRateLimiting(resolveCtx)
+	resolveCtx = h.configureRateLimiting(resolveCtx, reqCtx.operation.opType)
 	if reqCtx.customFieldValueRenderer != nil {
 		resolveCtx.SetFieldValueRenderer(reqCtx.customFieldValueRenderer)
 	}
@@ -418,7 +418,7 @@ func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (h *GraphQLHandler) configureRateLimiting(ctx *resolve.Context) *resolve.Context {
+func (h *GraphQLHandler) configureRateLimiting(ctx *resolve.Context, opType OperationType) *resolve.Context {
 	if h.rateLimiter == nil {
 		return ctx
 	}
@@ -431,6 +431,9 @@ func (h *GraphQLHandler) configureRateLimiting(ctx *resolve.Context) *resolve.Co
 	if h.rateLimitConfig.Strategy != "simple" {
 		return ctx
 	}
+	if h.rateLimitConfig.ExcludeSubscriptions && opType == OperationTypeSubscription {
+		return ctx
+	}
 	ctx.SetRateLimiter(h.rateLimiter)
 	ctx.RateLimitOptions = resolve.RateLimitOptions{
 		Enable:                          true,

router/core/websocket.go

@@ -1159,7 +1159,7 @@ func (h *WebSocketConnectionHandler) executeSubscription(registration *Subscript
 			resolveCtx.SetPreFetchFieldAuthorizer(h.graphqlHandler.authorizer)
 		}
 	}
-	resolveCtx = h.graphqlHandler.configureRateLimiting(resolveCtx)
+	resolveCtx = h.graphqlHandler.configureRateLimiting(resolveCtx, operationCtx.opType)
 
 	// Put in a closure to evaluate err after defer
 	defer func() {

router/pkg/config/config.go

@@ -742,6 +742,8 @@ 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"`
+	// ExcludeSubscriptions disables rate limiting for subscription operations only.
+	ExcludeSubscriptions bool `yaml:"exclude_subscriptions" envDefault:"false" env:"RATE_LIMIT_EXCLUDE_SUBSCRIPTIONS"`
 }
 
 type RateLimitErrorExtensionCode struct {

router/pkg/config/config.schema.json

@@ -2604,6 +2604,11 @@
           "type": "string",
           "description": "The expression to define a key suffix for the rate limit, e.g. by using request headers, claims, or a combination of both with a fallback strategy. The expression is specified as a string and needs to evaluate to a string. Please see https://expr-lang.org/ for more information."
         },
+        "exclude_subscriptions": {
+          "type": "boolean",
+          "default": false,
+          "description": "Exclude subscription operations from rate limiting. Be aware that disabling rate limiting for subscriptions can expose the router and subgraphs to a denial-of-service risk from unrestricted clients."
+        },
         "error_extension_code": {
           "type": "object",
           "description": "If enabled, a code will be added to the extensions.code field of error objects related to rate limiting. This allows clients to easily identify if an error happened due to rate limiting.",

router/pkg/config/fixtures/full.yaml

@@ -466,6 +466,7 @@ engine:
 rate_limit:
   enabled: true
   strategy: 'simple'
+  exclude_subscriptions: true
   storage:
     cluster_enabled: true
     urls:

router/pkg/config/testdata/config_defaults.json

@@ -398,7 +398,8 @@
     "ErrorExtensionCode": {
       "Enabled": true,
       "Code": "RATE_LIMIT_EXCEEDED"
-    }
+    },
+    "ExcludeSubscriptions": false
   },
   "LocalhostFallbackInsideDocker": true,
   "CDN": {

router/pkg/config/testdata/config_full.json

@@ -773,7 +773,8 @@
     "ErrorExtensionCode": {
       "Enabled": true,
       "Code": "RATE_LIMIT_EXCEEDED"
-    }
+    },
+    "ExcludeSubscriptions": true
   },
   "LocalhostFallbackInsideDocker": true,
   "CDN": {