Repository: Khan/genqlient
Pull request #360 (2024-11-14 16:55:58)

Subscription connection parameters support
Adds support for subscription (websocket) connection parameters support.

We are having an issue with authorization when using subscriptions in genqlient. Our mobile app team is using Apollo for GraphQL requests, and Apollo seems to sends its auth in a way that is different from how genqlient does it:

https://www.apollographql.com/docs/react/data/subscriptions#5-authenticate-over-websocket-optional

This is causing a problems, because we use genqlient for testing the backend. It seems that the way Apollo is implemented, authentication (and other headers) are not actually sent as HTTP headers, but as connection parameters instead.

I have:
- [x] Written a clear PR title and description (above)
- [x] Signed the [Khan Academy CLA](https://www.khanacademy.org/r/cla)
- [x] Added tests covering my changes, if applicable
- [x] Included a link to the issue fixed, if applicable
- [x] Included documentation, for new features
- [x] Added an entry to the changelog

HaraldNordgren: A big gotcha regarding HTTP headers and connection parameters, that I realized recently, is that websocket requests in 99desgins gqlgen follow a different path, going through here:

```go
	gqlgenServer.AddTransport(transport.Websocket{
		InitFunc: func(ctx context.Context, initPayload transport.InitPayload) (context.Context, *transport.InitPayload, error) {
			// initPayload contains connection parameters
		},
	})
```
HaraldNordgren: Ping @benjaminjkraft and @matthieu4294967296moineau, do you think it would be possible to get a quick look here? 🤗
HaraldNordgren: Ping @benjaminjkraft 🤗
jbaldwinroberts: this would be useful for us too - thanks
HaraldNordgren: Hi @benjaminjkraft, I know you are likely very busy. But it would be much appreciated to get some review here 🤗
HaraldNordgren: @benjaminjkraft Hope everything is under control and not too stressful 🤗

Thanks for getting a review here, much appreciated!

I think it's a great point regarding options. I have pushed a new version with this added. Let me know what you think!
HaraldNordgren: I couldn't figure out a way to generate the 99designs/gqlgen code for countAuthorized. Running 'go generate ./...' like I would do for my regular repositories fails.

So I created countAuthorized infrastructure myself. I hope  didn't introduce any issue. All the tests are passing at least 😃
HaraldNordgren: @benjaminjkraft I think it does make sense with the headers.

But it feels a bit tricky to get it right with the important "Sec-WebSocket-Protocol" key, which should always be present in the output. Maybe better to do as a follow-up?
benjaminjkraft: Oh, I see what you mean. Sure, I'll merge this and if we sort it out before we cut a release great, and if not oh well.

Patch:
diff --git a/graphql/client.go b/graphql/client.go
index d81d1ce..334f1e6 100644
--- a/graphql/client.go
+++ b/graphql/client.go
@@ -126,25 +126,41 @@ func NewClientUsingGet(endpoint string, httpClient Doer) Client {
 	return newClient(endpoint, httpClient, http.MethodGet)
 }
 
+type WebSocketOption func(*webSocketClient)
+
 // NewClientUsingWebSocket returns a [WebSocketClient] which makes subscription requests
 // to the given endpoint using webSocket.
 //
 // The client does not support queries nor mutations, and will return an error
 // if passed a request that attempts one.
-func NewClientUsingWebSocket(endpoint string, wsDialer Dialer, headers http.Header) WebSocketClient {
+func NewClientUsingWebSocket(endpoint string, wsDialer Dialer, headers http.Header, opts ...WebSocketOption) WebSocketClient {
 	if headers == nil {
 		headers = http.Header{}
 	}
 	if headers.Get("Sec-WebSocket-Protocol") == "" {
 		headers.Add("Sec-WebSocket-Protocol", "graphql-transport-ws")
 	}
-	return &webSocketClient{
+	client := &webSocketClient{
 		Dialer:        wsDialer,
 		Header:        headers,
 		errChan:       make(chan error),
 		endpoint:      endpoint,
 		subscriptions: subscriptionMap{map_: make(map[string]subscription)},
 	}
+
+	for _, opt := range opts {
+		opt(client)
+	}
+
+	return client
+}
+
+// WithConnectionParams sets up connection params to be sent to the server
+// during the initial connection handshake.
+func WithConnectionParams(connParams map[string]interface{}) WebSocketOption {
+	return func(ws *webSocketClient) {
+		ws.connParams = connParams
+	}
 }
 
 func newClient(endpoint string, httpClient Doer, method string) Client {
diff --git a/graphql/websocket.go b/graphql/websocket.go
index 10b77da..5fb97ce 100644
--- a/graphql/websocket.go
+++ b/graphql/websocket.go
@@ -48,12 +48,18 @@ type webSocketClient struct {
 	Header        http.Header
 	endpoint      string
 	conn          WSConn
+	connParams    map[string]interface{}
 	errChan       chan error
 	subscriptions subscriptionMap
 	isClosing     bool
 	sync.Mutex
 }
 
+type webSocketInitMessage struct {
+	Payload map[string]interface{} `json:"payload"`
+	Type    string                 `json:"type"`
+}
+
 type webSocketSendMessage struct {
 	Payload *Request `json:"payload"`
 	Type    string   `json:"type"`
@@ -67,8 +73,9 @@ type webSocketReceiveMessage struct {
 }
 
 func (w *webSocketClient) sendInit() error {
-	connInitMsg := webSocketSendMessage{
-		Type: webSocketTypeConnInit,
+	connInitMsg := webSocketInitMessage{
+		Type:    webSocketTypeConnInit,
+		Payload: w.connParams,
 	}
 	return w.sendStructAsJSON(connInitMsg)
 }
diff --git a/internal/integration/generated.go b/internal/integration/generated.go
index 9bb972d..446f5d2 100644
--- a/internal/integration/generated.go
+++ b/internal/integration/generated.go
@@ -1311,6 +1311,14 @@ type __queryWithVariablesInput struct {
 // GetId returns __queryWithVariablesInput.Id, and is useful for accessing the field via an interface.
 func (v *__queryWithVariablesInput) GetId() string { return v.Id }
 
+// countAuthorizedResponse is returned by countAuthorized on success.
+type countAuthorizedResponse struct {
+	CountAuthorized int `json:"countAuthorized"`
+}
+
+// GetCountAuthorized returns countAuthorizedResponse.CountAuthorized, and is useful for accessing the field via an interface.
+func (v *countAuthorizedResponse) GetCountAuthorized() int { return v.CountAuthorized }
+
 // countResponse is returned by count on success.
 type countResponse struct {
 	Count int `json:"count"`
@@ -3144,6 +3152,54 @@ func countForwardData(interfaceChan interface{}, jsonRawMsg json.RawMessage) err
 	return nil
 }
 
+// The subscription executed by countAuthorized.
+const countAuthorized_Operation = `
+subscription countAuthorized {
+	countAuthorized
+}
+`
+
+// To unsubscribe, use [graphql.WebSocketClient.Unsubscribe]
+func countAuthorized(
+	ctx_ context.Context,
+	client_ graphql.WebSocketClient,
+) (dataChan_ chan countAuthorizedWsResponse, subscriptionID_ string, err_ error) {
+	req_ := &graphql.Request{
+		OpName: "countAuthorized",
+		Query:  countAuthorized_Operation,
+	}
+
+	dataChan_ = make(chan countAuthorizedWsResponse)
+	subscriptionID_, err_ = client_.Subscribe(req_, dataChan_, countAuthorizedForwardData)
+
+	return dataChan_, subscriptionID_, err_
+}
+
+type countAuthorizedWsResponse graphql.BaseResponse[*countAuthorizedResponse]
+
+func countAuthorizedForwardData(interfaceChan interface{}, jsonRawMsg json.RawMessage) error {
+	var gqlResp graphql.Response
+	var wsResp countAuthorizedWsResponse
+	err := json.Unmarshal(jsonRawMsg, &gqlResp)
+	if err != nil {
+		return err
+	}
+	if len(gqlResp.Errors) == 0 {
+		err = json.Unmarshal(jsonRawMsg, &wsResp)
+		if err != nil {
+			return err
+		}
+	} else {
+		wsResp.Errors = gqlResp.Errors
+	}
+	dataChan_, ok := interfaceChan.(chan countAuthorizedWsResponse)
+	if !ok {
+		return errors.New("failed to cast interface into 'chan countAuthorizedWsResponse'")
+	}
+	dataChan_ <- wsResp
+	return nil
+}
+
 // The mutation executed by createUser.
 const createUser_Operation = `
 mutation createUser ($user: NewUser!) {
diff --git a/internal/integration/roundtrip.go b/internal/integration/roundtrip.go
index a020d7b..835fa69 100644
--- a/internal/integration/roundtrip.go
+++ b/internal/integration/roundtrip.go
@@ -156,14 +156,20 @@ func (md *MyDialer) DialContext(ctx context.Context, urlStr string, requestHeade
 	return graphql.WSConn(conn), err
 }
 
-func newRoundtripWebSocketClient(t *testing.T, endpoint string) graphql.WebSocketClient {
+func newRoundtripWebSocketClient(t *testing.T, endpoint string, opts ...graphql.WebSocketOption) graphql.WebSocketClient {
 	dialer := websocket.DefaultDialer
 	if !strings.HasPrefix(endpoint, "ws") {
 		_, address, _ := strings.Cut(endpoint, "://")
 		endpoint = "ws://" + address
 	}
+
 	return &roundtripClient{
-		wsWrapped: graphql.NewClientUsingWebSocket(endpoint, &MyDialer{Dialer: dialer}, nil),
-		t:         t,
+		wsWrapped: graphql.NewClientUsingWebSocket(
+			endpoint,
+			&MyDialer{Dialer: dialer},
+			nil,
+			opts...,
+		),
+		t: t,
 	}
 }
diff --git a/internal/integration/schema.graphql b/internal/integration/schema.graphql
index 9f148cc..16e13ca 100644
--- a/internal/integration/schema.graphql
+++ b/internal/integration/schema.graphql
@@ -19,6 +19,7 @@ type Mutation {
 
 type Subscription {
   count: Int!
+  countAuthorized: Int!
 }
 
 type User implements Being & Lucky {
diff --git a/internal/integration/server/gqlgen_exec.go b/internal/integration/server/gqlgen_exec.go
index ccf5d2a..c3f9557 100644
--- a/internal/integration/server/gqlgen_exec.go
+++ b/internal/integration/server/gqlgen_exec.go
@@ -78,7 +78,8 @@ type ComplexityRoot struct {
 	}
 
 	Subscription struct {
-		Count func(childComplexity int) int
+		Count           func(childComplexity int) int
+		CountAuthorized func(childComplexity int) int
 	}
 
 	User struct {
@@ -108,6 +109,7 @@ type QueryResolver interface {
 }
 type SubscriptionResolver interface {
 	Count(ctx context.Context) (<-chan int, error)
+	CountAuthorized(ctx context.Context) (<-chan int, error)
 }
 
 type executableSchema struct {
@@ -291,6 +293,13 @@ func (e *executableSchema) Complexity(typeName, field string, childComplexity in
 
 		return e.complexity.Subscription.Count(childComplexity), true
 
+	case "Subscription.countAuthorized":
+		if e.complexity.Subscription.CountAuthorized == nil {
+			break
+		}
+
+		return e.complexity.Subscription.CountAuthorized(childComplexity), true
+
 	case "User.birthdate":
 		if e.complexity.User.Birthdate == nil {
 			break
@@ -484,6 +493,7 @@ type Mutation {
 
 type Subscription {
   count: Int!
+  countAuthorized: Int!
 }
 
 type User implements Being & Lucky {
@@ -1811,6 +1821,64 @@ func (ec *executionContext) fieldContext_Subscription_count(ctx context.Context,
 	return fc, nil
 }
 
+func (ec *executionContext) _Subscription_countAuthorized(ctx context.Context, field graphql.CollectedField) (ret func(ctx context.Context) graphql.Marshaler) {
+	fc, err := ec.fieldContext_Subscription_countAuthorized(ctx, field)
+	if err != nil {
+		return nil
+	}
+	ctx = graphql.WithFieldContext(ctx, fc)
+	defer func() {
+		if r := recover(); r != nil {
+			ec.Error(ctx, ec.Recover(ctx, r))
+			ret = nil
+		}
+	}()
+	resTmp, err := ec.ResolverMiddleware(ctx, func(rctx context.Context) (interface{}, error) {
+		ctx = rctx // use context from middleware stack in children
+		return ec.resolvers.Subscription().CountAuthorized(rctx)
+	})
+	if err != nil {
+		ec.Error(ctx, err)
+		return nil
+	}
+	if resTmp == nil {
+		if !graphql.HasFieldError(ctx, fc) {
+			ec.Errorf(ctx, "must not be null")
+		}
+		return nil
+	}
+	return func(ctx context.Context) graphql.Marshaler {
+		select {
+		case res, ok := <-resTmp.(<-chan int):
+			if !ok {
+				return nil
+			}
+			return graphql.WriterFunc(func(w io.Writer) {
+				w.Write([]byte{'{'})
+				graphql.MarshalString(field.Alias).MarshalGQL(w)
+				w.Write([]byte{':'})
+				ec.marshalNInt2int(ctx, field.Selections, res).MarshalGQL(w)
+				w.Write([]byte{'}'})
+			})
+		case <-ctx.Done():
+			return nil
+		}
+	}
+}
+
+func (ec *executionContext) fieldContext_Subscription_countAuthorized(ctx context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
+	fc = &graphql.FieldContext{
+		Object:     "Subscription",
+		Field:      field,
+		IsMethod:   true,
+		IsResolver: true,
+		Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
+			return nil, errors.New("field of type Int does not have child fields")
+		},
+	}
+	return fc, nil
+}
+
 func (ec *executionContext) _User_id(ctx context.Context, field graphql.CollectedField, obj *User) (ret graphql.Marshaler) {
 	fc, err := ec.fieldContext_User_id(ctx, field)
 	if err != nil {
@@ -4398,6 +4466,8 @@ func (ec *executionContext) _Subscription(ctx context.Context, sel ast.Selection
 	switch fields[0].Name {
 	case "count":
 		return ec._Subscription_count(ctx, fields[0])
+	case "countAuthorized":
+		return ec._Subscription_countAuthorized(ctx, fields[0])
 	default:
 		panic("unknown field " + strconv.Quote(fields[0].Name))
 	}
diff --git a/internal/integration/server/server.go b/internal/integration/server/server.go
index 2c79909..e626e34 100644
--- a/internal/integration/server/server.go
+++ b/internal/integration/server/server.go
@@ -173,11 +173,45 @@ func (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
 	return respChan, nil
 }
 
+func (s *subscriptionResolver) CountAuthorized(ctx context.Context) (<-chan int, error) {
+	if getAuthToken(ctx) != "authorized-user-token" {
+		return nil, fmt.Errorf("unauthorized")
+	}
+
+	return s.Count(ctx)
+}
+
+const AuthKey = "authToken"
+
+type (
+	authTokenCtxKey struct{}
+)
+
+func withAuthToken(ctx context.Context, token string) context.Context {
+	return context.WithValue(ctx, authTokenCtxKey{}, token)
+}
+
+func getAuthToken(ctx context.Context) string {
+	if tkn, ok := ctx.Value(authTokenCtxKey{}).(string); ok {
+		return tkn
+	}
+	return ""
+}
+
 func RunServer() *httptest.Server {
 	gqlgenServer := handler.New(NewExecutableSchema(Config{Resolvers: &resolver{}}))
 	gqlgenServer.AddTransport(transport.POST{})
 	gqlgenServer.AddTransport(transport.GET{})
-	gqlgenServer.AddTransport(transport.Websocket{})
+
+	gqlgenServer.AddTransport(transport.Websocket{
+		InitFunc: func(ctx context.Context, initPayload transport.InitPayload) (context.Context, *transport.InitPayload, error) {
+			if authToken, ok := initPayload[AuthKey].(string); ok && authToken != "" {
+				ctx = withAuthToken(ctx, authToken)
+			}
+			return ctx, &initPayload, nil
+		},
+	})
+
 	gqlgenServer.AroundResponses(func(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
 		graphql.RegisterExtension(ctx, "foobar", "test")
 		return next(ctx)


Test patch:
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index 70b0cb9..be4e168 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -98,6 +98,7 @@ func TestSubscription(t *testing.T) {
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
 			wsClient := newRoundtripWebSocketClient(t, server.URL)
+
 			errChan, err := wsClient.Start(ctx)
 			require.NoError(t, err)
 
@@ -146,6 +147,93 @@ func TestSubscription(t *testing.T) {
 	}
 }
 
+func TestSubscriptionConnectionParams(t *testing.T) {
+	_ = `# @genqlient
+	subscription countAuthorized { countAuthorized }`
+
+	authKey := server.AuthKey
+
+	ctx := context.Background()
+	server := server.RunServer()
+	defer server.Close()
+
+	cases := []struct {
+		connParams    map[string]interface{}
+		name          string
+		expectedError string
+		opts          []graphql.WebSocketOption
+	}{
+		{
+			name: "authorized_user_gets_counter",
+			opts: []graphql.WebSocketOption{
+				graphql.WithConnectionParams(map[string]interface{}{
+					authKey: "authorized-user-token",
+				}),
+			},
+		},
+		{
+			name:          "unauthorized_user_gets_error",
+			expectedError: "input: countAuthorized unauthorized\n",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			wsClient := newRoundtripWebSocketClient(
+				t,
+				server.URL,
+				tc.opts...,
+			)
+
+			errChan, err := wsClient.Start(ctx)
+			require.NoError(t, err)
+
+			dataChan, subscriptionID, err := countAuthorized(ctx, wsClient)
+			require.NoError(t, err)
+			defer wsClient.Close()
+
+			var (
+				counter = 0
+				start   = time.Now()
+			)
+
+			for loop := true; loop; {
+				select {
+				case resp, more := <-dataChan:
+					if !more {
+						loop = false
+						break
+					}
+
+					if tc.expectedError != "" {
+						require.Error(t, resp.Errors)
+						assert.Equal(t, tc.expectedError, resp.Errors.Error())
+						continue
+					}
+
+					require.NotNil(t, resp.Data)
+					assert.Equal(t, counter, resp.Data.CountAuthorized)
+					require.Nil(t, resp.Errors)
+
+					if time.Since(start) > 5*time.Second {
+						err := wsClient.Unsubscribe(subscriptionID)
+						require.NoError(t, err)
+						loop = false
+					}
+
+					counter++
+
+				case err := <-errChan:
+					require.NoError(t, err)
+
+				case <-time.After(10 * time.Second):
+					require.NoError(t, fmt.Errorf("subscription timed out"))
+				}
+			}
+		})
+	}
+}
+
 func TestServerError(t *testing.T) {
 	_ = `# @genqlient
 	query failingQuery { fail me { id } }`