diff --git a/v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go b/v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go index 4d6c99182d..b58bed6058 100644 --- a/v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go +++ b/v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go @@ -14,6 +14,7 @@ import ( "unicode" "github.com/buger/jsonparser" + "github.com/cespare/xxhash/v2" "github.com/jensneuse/abstractlogger" "github.com/pkg/errors" "github.com/tidwall/sjson" @@ -1997,6 +1998,11 @@ type SubscriptionSource struct { subscriptionOnStartFns []SubscriptionOnStartFn } +func (s *SubscriptionSource) HashTriggerInput(input []byte, xxh *xxhash.Digest) error { + _, err := xxh.Write(input) + return err +} + // Start the subscription. The updater is called on new events. Start needs to be called in a separate goroutine. func (s *SubscriptionSource) Start(ctx *resolve.Context, headers http.Header, input []byte, updater resolve.SubscriptionUpdater) error { var options GraphQLSubscriptionOptions diff --git a/v2/pkg/engine/resolve/datasource.go b/v2/pkg/engine/resolve/datasource.go index be9169c1d3..b06444e010 100644 --- a/v2/pkg/engine/resolve/datasource.go +++ b/v2/pkg/engine/resolve/datasource.go @@ -4,6 +4,8 @@ import ( "context" "net/http" + "github.com/cespare/xxhash/v2" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" ) @@ -16,6 +18,9 @@ type SubscriptionDataSource interface { // Start is called when a new subscription is created. It establishes the connection to the data source. // The updater is used to send updates to the client. Deduplication of the request must be done before calling this method. Start(ctx *Context, headers http.Header, input []byte, updater SubscriptionUpdater) error + // HashTriggerInput writes identity-relevant fields of the subscription into xxh. + // The resolver appends the subgraph headers hash afterward to produce the final trigger ID. + HashTriggerInput(input []byte, xxh *xxhash.Digest) error } // HookableSubscriptionDataSource is a hookable interface for subscription data sources. diff --git a/v2/pkg/engine/resolve/resolve.go b/v2/pkg/engine/resolve/resolve.go index a45f3487ec..4108bbb56d 100644 --- a/v2/pkg/engine/resolve/resolve.go +++ b/v2/pkg/engine/resolve/resolve.go @@ -1274,9 +1274,15 @@ func (r *Resolver) UnsubscribeClient(connectionID ConnectionID) error { // prepareTrigger safely gets the headers for the trigger Subgraph and computes the hash across headers and input // the generated hash is the unique triggerID // the headers must be forwarded to the DataSource to create the trigger -func (r *Resolver) prepareTrigger(ctx *Context, sourceName string, input []byte) (headers http.Header, triggerID uint64) { +func (r *Resolver) prepareTrigger(ctx *Context, sourceName string, input []byte, source SubscriptionDataSource) ( + headers http.Header, triggerID uint64, err error) { keyGen := pool.Hash64.Get() - _, _ = keyGen.Write(input) + defer pool.Hash64.Put(keyGen) + + if err = source.HashTriggerInput(input, keyGen); err != nil { + return nil, 0, err + } + if ctx.SubgraphHeadersBuilder != nil { var headersHash uint64 headers, headersHash = ctx.SubgraphHeadersBuilder.HeadersForSubgraph(sourceName) @@ -1286,9 +1292,10 @@ func (r *Resolver) prepareTrigger(ctx *Context, sourceName string, input []byte) _, _ = keyGen.Write(b[:]) } } + triggerID = keyGen.Sum64() - pool.Hash64.Put(keyGen) - return headers, triggerID + + return headers, triggerID, nil } func (r *Resolver) ResolveGraphQLSubscription(ctx *Context, subscription *GraphQLSubscription, writer SubscriptionResponseWriter) error { @@ -1328,7 +1335,11 @@ func (r *Resolver) ResolveGraphQLSubscription(ctx *Context, subscription *GraphQ return nil } - headers, triggerID := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input) + headers, triggerID, err := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input, subscription.Trigger.Source) + if err != nil { + msg := []byte(`{"errors":[{"message":"failed to prepare subscription trigger"}]}`) + return writeFlushComplete(writer, msg) + } id := SubscriptionIdentifier{ ConnectionID: NewConnectionID(), SubscriptionID: 0, @@ -1422,7 +1433,11 @@ func (r *Resolver) AsyncResolveGraphQLSubscription(ctx *Context, subscription *G return err } - headers, triggerID := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input) + headers, triggerID, err := r.prepareTrigger(ctx, subscription.Trigger.SourceName, input, subscription.Trigger.Source) + if err != nil { + msg := []byte(`{"errors":[{"message":"failed to prepare subscription trigger"}]}`) + return writeFlushComplete(writer, msg) + } return r.addSubscription(triggerID, &addSubscription{ ctx: ctx, diff --git a/v2/pkg/engine/resolve/resolve_test.go b/v2/pkg/engine/resolve/resolve_test.go index 0ee86948f1..0ba646b229 100644 --- a/v2/pkg/engine/resolve/resolve_test.go +++ b/v2/pkg/engine/resolve/resolve_test.go @@ -14,6 +14,7 @@ import ( "testing" "time" + "github.com/cespare/xxhash/v2" "github.com/golang/mock/gomock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -5766,6 +5767,11 @@ type _fakeStream struct { subscriptionOnStartFn func(ctx StartupHookContext, input []byte) (err error) } +func (f *_fakeStream) HashTriggerInput(input []byte, xxh *xxhash.Digest) error { + _, err := xxh.Write(input) + return err +} + func (f *_fakeStream) SubscriptionOnStart(ctx StartupHookContext, input []byte) (err error) { if f.subscriptionOnStartFn == nil { return nil @@ -8020,6 +8026,11 @@ type startFailStream struct { subBReady chan struct{} } +func (s *startFailStream) HashTriggerInput(input []byte, xxh *xxhash.Digest) error { + _, err := xxh.Write(input) + return err +} + func (s *startFailStream) Start(_ *Context, _ http.Header, _ []byte, _ SubscriptionUpdater) error { <-s.subBReady return errors.New("connection refused") @@ -8104,6 +8115,11 @@ type hookFailStream struct { sourceStarted atomic.Bool } +func (s *hookFailStream) HashTriggerInput(input []byte, xxh *xxhash.Digest) error { + _, err := xxh.Write(input) + return err +} + func (s *hookFailStream) Start(_ *Context, _ http.Header, _ []byte, _ SubscriptionUpdater) error { s.sourceStarted.Store(true) select {} diff --git a/v2/pkg/engine/resolve/resolver_subscription_test.go b/v2/pkg/engine/resolve/resolver_subscription_test.go index 1329dad030..ed2adeb57d 100644 --- a/v2/pkg/engine/resolve/resolver_subscription_test.go +++ b/v2/pkg/engine/resolve/resolver_subscription_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "github.com/cespare/xxhash/v2" "github.com/stretchr/testify/require" ) @@ -66,6 +67,11 @@ type FakeSource struct { interval time.Duration } +func (f *FakeSource) HashTriggerInput(input []byte, xxh *xxhash.Digest) error { + _, err := xxh.Write(input) + return err +} + func (f *FakeSource) Start(ctx *Context, headers http.Header, input []byte, updater SubscriptionUpdater) error { go func() { for i, u := range f.updates {