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-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() 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{