From b274787d6ddc1bcc83db1b2613b0432750bfa2ce Mon Sep 17 00:00:00 2001 From: Dominik Korittki <23359034+dkorittki@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:24:42 +0200 Subject: [PATCH 1/2] fix: ignore header rules for pubsub trigger sources --- router-tests/events/trigger_test.go | 37 ++++-- router/core/context.go | 10 ++ .../header_rule_engine_buildheader_test.go | 105 ++++++++++++++++++ 3 files changed, 144 insertions(+), 8 deletions(-) diff --git a/router-tests/events/trigger_test.go b/router-tests/events/trigger_test.go index 21df4cc290..e1fb06916e 100644 --- a/router-tests/events/trigger_test.go +++ b/router-tests/events/trigger_test.go @@ -1,12 +1,15 @@ package events_test import ( + "net/http" "sync" "testing" "time" "github.com/stretchr/testify/require" "github.com/wundergraph/cosmo/router-tests/testenv" + "github.com/wundergraph/cosmo/router/core" + "github.com/wundergraph/cosmo/router/pkg/config" ) // TestEDFSTriggerDeduplication verifies trigger ID generation for Cosmo Streams subscriptions. @@ -33,8 +36,9 @@ func TestEDFSTriggerDeduplication(t *testing.T) { // provider ID, so they must resolve to the same trigger ID — wait for exactly // one trigger to be initialized before asserting. xEnv.WaitForTriggerCount(1, time.Second*10) - xEnv.RequireTriggerCount(1) xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename":"Employee"}`), 2, time.Second*10) + // Assert the exact count only after both subscriptions have been served + xEnv.RequireTriggerCount(1) }() // Subscription 1: selects only id. @@ -104,14 +108,23 @@ func TestEDFSTriggerDeduplication(t *testing.T) { }) }) - // Two subscriptions with the same query but different initial_payload (headers) should + // Two subscriptions with the same query but different headers should // still share a single NATS trigger because the trigger ID is based on the NATS subject, // not on connection-level metadata like headers. - t.Run("same subject different initial payload shares one trigger", func(t *testing.T) { + t.Run("same subject different headers shares one trigger", func(t *testing.T) { t.Parallel() testenv.Run(t, &testenv.Config{ RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate, EnableNats: true, + RouterOptions: []core.Option{ + core.WithHeaderRules(config.HeaderRules{ + All: &config.GlobalHeaderRule{ + Request: []*config.RequestHeaderRule{ + {Operation: config.HeaderRuleOperationPropagate, Named: "Authorization"}, + }, + }, + }), + }, }, func(t *testing.T, xEnv *testenv.Environment) { var done sync.WaitGroup done.Add(2) @@ -119,14 +132,19 @@ func TestEDFSTriggerDeduplication(t *testing.T) { go func() { xEnv.WaitForSubscriptionCount(2, time.Second*10) xEnv.WaitForTriggerCount(1, time.Second*10) - xEnv.RequireTriggerCount(1) xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename":"Employee"}`), 2, time.Second*10) + // Asserted after both subscriptions have been served, see the comment in the + // "different selected fields" subtest above. + xEnv.RequireTriggerCount(1) }() - // Subscription 1: sends Authorization header A in connection_init. + // Subscription 1: sends Authorization header A go func() { defer done.Done() - conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, []byte(`{"headers":{"Authorization":"Bearer token-a"}}`)) + conn := xEnv.InitGraphQLWebSocketConnection( + http.Header{"Authorization": []string{"Bearer token-a"}}, nil, + []byte(`{"headers":{"Authorization":"Bearer token-a"}}`), + ) defer conn.Close() err := testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{ @@ -154,10 +172,13 @@ func TestEDFSTriggerDeduplication(t *testing.T) { require.Equal(t, "1", complete.ID) }() - // Subscription 2: sends a different Authorization header in connection_init. + // Subscription 2: sends Authorization header B go func() { defer done.Done() - conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, []byte(`{"headers":{"Authorization":"Bearer token-b"}}`)) + conn := xEnv.InitGraphQLWebSocketConnection( + http.Header{"Authorization": []string{"Bearer token-b"}}, nil, + []byte(`{"headers":{"Authorization":"Bearer token-b"}}`), + ) defer conn.Close() err := testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{ diff --git a/router/core/context.go b/router/core/context.go index bda68f50c1..14000b26c5 100644 --- a/router/core/context.go +++ b/router/core/context.go @@ -24,6 +24,7 @@ import ( "github.com/wundergraph/cosmo/router/pkg/authentication" "github.com/wundergraph/cosmo/router/pkg/config" "github.com/wundergraph/cosmo/router/pkg/graphqlschemausage" + pubsub "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" ctrace "github.com/wundergraph/cosmo/router/pkg/trace" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" @@ -334,6 +335,15 @@ func SubgraphHeadersBuilder(ctx *requestContext, headerPropagation *HeaderPropag headers := make(map[string]*HeaderWithHash, len(p.Response.Response.DataSources)+1) makeHeaders(headers, p.Response.Response.DataSources) + // avoid adding header rules for pubsub triggers sources, as no headers are passed to them. + _, isPubSub := p.Response.Trigger.Source.(pubsub.SubscriptionDataSource) + if isPubSub { + return &headerBuilder{ + headers: headers, + allHash: keyGen.Sum64(), + } + } + h, hh := headerPropagation.BuildRequestHeaderForSubgraph(p.Response.Trigger.SourceName, ctx) headers[p.Response.Trigger.SourceName] = &HeaderWithHash{ Header: h, diff --git a/router/core/header_rule_engine_buildheader_test.go b/router/core/header_rule_engine_buildheader_test.go index 24f4e66f8a..2a12f604fe 100644 --- a/router/core/header_rule_engine_buildheader_test.go +++ b/router/core/header_rule_engine_buildheader_test.go @@ -10,6 +10,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/wundergraph/cosmo/router/pkg/config" + pubsub "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/graphql_datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" "go.uber.org/zap" @@ -324,6 +326,109 @@ func TestSubgraphHeadersBuilder_SubscriptionPlan_IncludesTriggerAndResponse(t *t assert.Equal(t, hashTrig, hashTrig2) } +func TestSubgraphHeadersBuilder_SubscriptionPlan_SkipsPubSubTriggerSource(t *testing.T) { + ht, err := NewHeaderPropagation(t.Context(), zap.NewNop(), &config.HeaderRules{ + All: &config.GlobalHeaderRule{ + Request: []*config.RequestHeaderRule{ + {Operation: "propagate", Named: "X-A"}, + {Operation: "set", Name: "X-Static", Value: "static"}, + }, + }, + }, nil) + require.NoError(t, err) + + newCtx := func(headerValue string) *requestContext { + clientReq := httptest.NewRequest("POST", "http://localhost", nil) + clientReq.Header.Set("X-A", headerValue) + + return &requestContext{ + logger: zap.NewNop(), + responseWriter: httptest.NewRecorder(), + request: clientReq, + operation: &operationContext{}, + subgraphResolver: NewSubgraphResolver(nil), + } + } + + newPlan := func(source resolve.SubscriptionDataSource) *plan.SubscriptionResponsePlan { + return &plan.SubscriptionResponsePlan{ + Response: &resolve.GraphQLSubscription{ + Response: &resolve.GraphQLResponse{ + DataSources: []resolve.DataSourceInfo{{Name: "sg-resp"}}, + }, + Trigger: resolve.GraphQLSubscriptionTrigger{ + SourceName: "sg-trigger", + Source: source, + }, + }, + } + } + + t.Run("pubsub trigger source gets no headers", func(t *testing.T) { + ctx := newCtx("va") + pubSubDS := &pubsub.PubSubSubscriptionDataSource[pubsub.SubscriptionEventConfiguration]{} + + hb := SubgraphHeadersBuilder(ctx, ht, newPlan(pubSubDS)) + require.NotNil(t, hb) + + // The trigger source is a pubsub data source, so no header rules are built for it. + hTrig, hashTrig := hb.HeadersForSubgraph("sg-trigger") + assert.Nil(t, hTrig) + assert.Zero(t, hashTrig) + + // Regular response data sources are unaffected. + hResp, hashResp := hb.HeadersForSubgraph("sg-resp") + require.NotNil(t, hResp) + assert.Equal(t, "va", hResp.Get("X-A")) + assert.Equal(t, "static", hResp.Get("X-Static")) + require.NotZero(t, hashResp) + }) + + t.Run("pubsub trigger source is excluded from the all hash", func(t *testing.T) { + pubSubDS := &pubsub.PubSubSubscriptionDataSource[pubsub.SubscriptionEventConfiguration]{} + nonPubSubDS := &graphql_datasource.SubscriptionSource{} + + pubSubHash := SubgraphHeadersBuilder(newCtx("va"), ht, + newPlan(pubSubDS)).HashAll() + + // A plan containing only the response data source must produce the same hash, proving + // the trigger source no longer contributes to it. + responseOnlyHash := SubgraphHeadersBuilder(newCtx("va"), ht, &plan.SynchronousResponsePlan{ + Response: &resolve.GraphQLResponse{ + DataSources: []resolve.DataSourceInfo{{Name: "sg-resp"}}, + }, + }).HashAll() + assert.Equal(t, responseOnlyHash, pubSubHash) + + // A non-pubsub trigger source still contributes to the hash. + regularHash := SubgraphHeadersBuilder(newCtx("va"), ht, + newPlan(nonPubSubDS)).HashAll() + assert.NotEqual(t, regularHash, pubSubHash) + }) + + t.Run("pubsub all hash is independent of propagated client headers", func(t *testing.T) { + pubSubDS := &pubsub.PubSubSubscriptionDataSource[pubsub.SubscriptionEventConfiguration]{} + pubSubOnlyPlan := func() *plan.SubscriptionResponsePlan { + return &plan.SubscriptionResponsePlan{ + Response: &resolve.GraphQLSubscription{ + Response: &resolve.GraphQLResponse{}, + Trigger: resolve.GraphQLSubscriptionTrigger{ + SourceName: "sg-trigger", + Source: pubSubDS, + }, + }, + } + } + + // Two clients sending different values for a propagated header must produce the same + // hash, so they share a single trigger instead of each creating their own. + assert.Equal(t, + SubgraphHeadersBuilder(newCtx("va"), ht, pubSubOnlyPlan()).HashAll(), + SubgraphHeadersBuilder(newCtx("vb"), ht, pubSubOnlyPlan()).HashAll(), + ) + }) +} + func TestSubgraphHeadersBuilder_MissingPrePopulatedCache(t *testing.T) { ht, err := NewHeaderPropagation(t.Context(), zap.NewNop(), &config.HeaderRules{ All: &config.GlobalHeaderRule{ From b3ef0c0763157e69476e0bde587e8b4268c2036b Mon Sep 17 00:00:00 2001 From: Dominik Korittki <23359034+dkorittki@users.noreply.github.com> Date: Tue, 4 Aug 2026 10:11:40 +0200 Subject: [PATCH 2/2] fix: fix test + add new one Fixes the single flight test assumption that edfs triggers are not deduplicated on different headers. In contrast to normal subscriptions they now will. The test now expects one trigger will still ensuring deduplication to subgraphs during resolving still works. Also added another websocket subscription test to make sure that for non-edfs and edfs triggers the deduplication based on headers work as expected. --- router-tests/operations/singleflight_test.go | 60 ++++++-------- router-tests/subscriptions/websocket_test.go | 84 ++++++++++++++++++++ 2 files changed, 109 insertions(+), 35 deletions(-) diff --git a/router-tests/operations/singleflight_test.go b/router-tests/operations/singleflight_test.go index e30ce6ab55..f9eb8e62e9 100644 --- a/router-tests/operations/singleflight_test.go +++ b/router-tests/operations/singleflight_test.go @@ -1,7 +1,6 @@ package integration import ( - "context" "fmt" "net/http" "sync" @@ -521,7 +520,11 @@ func TestSingleFlight(t *testing.T) { require.Less(t, actualSubgraphRequests, numOfOperations) }) }) - t.Run("subscription deduplication with multiple subgraphs - different headers", func(t *testing.T) { + // All subscriptions share a single EDFS trigger regardless of their headers, because the + // trigger ID of a pubsub source is derived from the subject and provider only. The nested + // fetches each subscription performs to resolve its response are still built from the + // per-request propagated headers, so those must not be de-duplicated by single flight. + t.Run("subscription with different headers does not deduplicate subgraph fetches", func(t *testing.T) { t.Parallel() testenv.Run(t, &testenv.Config{ RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate, @@ -556,24 +559,14 @@ func TestSingleFlight(t *testing.T) { ) done.Add(int(numOfOperations)) - // Continuously publish until all consumers have received their message. - // NATSPublishUntilMinMessagesSent is insufficient here because cumulative - // MessagesSent can reach 10 before all 10 consumers are served (retries - // deliver to already-served consumers, inflating the count). - publishCtx, publishCancel := context.WithCancel(xEnv.Context) + // Wait for all subscriptions to be established before triggering. The differing + // Authorization headers do not split the trigger, so a single message fans out + // to all subscriptions. go func() { xEnv.WaitForSubscriptionCount(uint64(numOfOperations), time.Second*15) - xEnv.WaitForTriggerCount(uint64(numOfOperations), time.Second*15) - for { - select { - case <-publishCtx.Done(): - return - default: - } - _ = xEnv.NatsConnectionDefault.Publish(xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename": "Employee"}`)) - _ = xEnv.NatsConnectionDefault.Flush() - time.Sleep(500 * time.Millisecond) - } + xEnv.WaitForTriggerCount(1, time.Second*15) + // Trigger the subscription via NATS to get updates for all subscriptions + xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename": "Employee"}`), 1, time.Second*15) }() for i := int64(0); i < numOfOperations; i++ { @@ -606,29 +599,26 @@ func TestSingleFlight(t *testing.T) { }) require.NoError(t, err) - // Read messages until we get "complete", draining any extra - // "next" messages that may arrive from publish retries - for { - var reply testenv.WebSocketMessage - err = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - require.NoError(t, err) - err = testenv.WSReadJSON(t, conn, &reply) - require.NoError(t, err) - if reply.Type == "complete" { - require.Equal(t, "1", reply.ID) - break - } - } + // Read the complete message + var complete testenv.WebSocketMessage + err = conn.SetReadDeadline(time.Now().Add(1 * time.Second)) + require.NoError(t, err) + err = testenv.WSReadJSON(t, conn, &complete) + require.NoError(t, err) + require.Equal(t, "complete", complete.Type) + require.Equal(t, "1", complete.ID) }(i) } done.Wait() - publishCancel() xEnv.WaitForSubscriptionCount(0, time.Second*5) - // We expect no request de-duplication because different headers must not be de-duplicated - // Publish retries may increase the count, so check >= not == + // We expect no request de-duplication because the fetches carry different headers. + // The NATS event itself supplies __typename and id — the only fields the pubsub + // data source owns — so resolving details.forename and details.surname costs one + // entity fetch to the employees subgraph per subscription: 10 in total. actualSubgraphRequests := xEnv.SubgraphRequestCount.Global.Load() - require.GreaterOrEqual(t, actualSubgraphRequests, numOfOperations) + require.Equal(t, numOfOperations, actualSubgraphRequests) + require.Equal(t, numOfOperations, xEnv.SubgraphRequestCount.Employees.Load()) }) }) t.Run("mutation with multiple subgraphs deduplication", func(t *testing.T) { diff --git a/router-tests/subscriptions/websocket_test.go b/router-tests/subscriptions/websocket_test.go index c221a6ecd6..be9c26bac4 100644 --- a/router-tests/subscriptions/websocket_test.go +++ b/router-tests/subscriptions/websocket_test.go @@ -741,6 +741,90 @@ func TestWebSockets(t *testing.T) { xEnv.WaitForSubscriptionCount(0, time.Second*5) }) }) + // Trigger IDs for regular (non-EDFS) subgraph subscriptions include the hash of the + // headers built by the SubgraphHeadersBuilder, so clients whose propagated headers + // differ must never share an upstream trigger. Upgrade header, query param and initial + // payload forwarding are disabled in both subtests so that the trigger input is identical + // for both clients and header propagation is the only discriminator. + t.Run("subscription trigger deduplication with header propagation", func(t *testing.T) { + t.Parallel() + + headerRules := config.HeaderRules{ + All: &config.GlobalHeaderRule{ + Request: []*config.RequestHeaderRule{ + { + Operation: config.HeaderRuleOperationPropagate, + Named: "Authorization", + }, + }, + }, + } + + websocketConfig := func(cfg *config.WebSocketConfiguration) { + cfg.ForwardUpgradeHeaders.Enabled = false + cfg.ForwardUpgradeQueryParams.Enabled = false + cfg.ForwardInitialPayload = false + } + + subscribeCurrentTime := func(t *testing.T, xEnv *testenv.Environment, authorization string) *websocket.Conn { + t.Helper() + + conn := xEnv.InitGraphQLWebSocketConnection(http.Header{ + "Authorization": []string{authorization}, + }, nil, nil) + + err := testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{ + ID: "1", + Type: "subscribe", + Payload: []byte(`{"query":"subscription { currentTime { unixTime timeStamp }}"}`), + }) + require.NoError(t, err) + + return conn + } + + t.Run("different headers use separate triggers", func(t *testing.T) { + t.Parallel() + + testenv.Run(t, &testenv.Config{ + ModifyWebsocketConfiguration: websocketConfig, + RouterOptions: []core.Option{ + core.WithHeaderRules(headerRules), + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + connA := subscribeCurrentTime(t, xEnv, "Bearer token-a") + defer connA.Close() + + connB := subscribeCurrentTime(t, xEnv, "Bearer token-b") + defer connB.Close() + + xEnv.WaitForSubscriptionCount(2, time.Second*15) + xEnv.WaitForTriggerCount(2, time.Second*15) + xEnv.RequireTriggerCount(2) + }) + }) + + t.Run("same headers share one trigger", func(t *testing.T) { + t.Parallel() + + testenv.Run(t, &testenv.Config{ + ModifyWebsocketConfiguration: websocketConfig, + RouterOptions: []core.Option{ + core.WithHeaderRules(headerRules), + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + connA := subscribeCurrentTime(t, xEnv, "Bearer token-a") + defer connA.Close() + + connB := subscribeCurrentTime(t, xEnv, "Bearer token-a") + defer connB.Close() + + xEnv.WaitForSubscriptionCount(2, time.Second*15) + xEnv.WaitForTriggerCount(1, time.Second*15) + xEnv.RequireTriggerCount(1) + }) + }) + }) t.Run("empty allow lists should allow all headers and query args", func(t *testing.T) { t.Parallel()