Allow setting websocket connection parameters

graphql/client.go

@@ -132,6 +132,16 @@ func NewClientUsingGet(endpoint string, httpClient Doer) Client {
 // 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 {
+	return NewClientUsingWebSocketWithConnectionParams(endpoint, wsDialer, headers, nil)
+}
+
+// NewClientUsingWebSocketWithConnectionParams returns a [WebSocketClient] which makes subscription requests
+// to the given endpoint using webSocket. It allows to pass additional connection parameters
+// to the server during the initial connection handshake.
+//
+// connectionParams is a map of connection parameters to be sent to the server
+// during the initial connection handshake.
+func NewClientUsingWebSocketWithConnectionParams(endpoint string, wsDialer Dialer, headers http.Header, connParams map[string]any) WebSocketClient {
 	if headers == nil {
 		headers = http.Header{}
 	}
@@ -141,6 +151,7 @@ func NewClientUsingWebSocket(endpoint string, wsDialer Dialer, headers http.Head
 	return &webSocketClient{
 		Dialer:        wsDialer,
 		Header:        headers,
+		connParams:    connParams,
 		errChan:       make(chan error),
 		endpoint:      endpoint,
 		subscriptions: subscriptionMap{map_: make(map[string]subscription)},

graphql/websocket.go

@@ -48,16 +48,22 @@ type webSocketClient struct {
 	Header        http.Header
 	endpoint      string
 	conn          WSConn
+	connParams    map[string]any
 	errChan       chan error
 	subscriptions subscriptionMap
 	isClosing     bool
 	sync.Mutex
 }
 
+type webSocketInitMessage struct {
+	Type    string                 `json:"type"`
+	Payload map[string]interface{} `json:"payload"`
+}
+
 type webSocketSendMessage struct {
-	Payload *Request `json:"payload"`
 	Type    string   `json:"type"`
 	ID      string   `json:"id"`
+	Payload *Request `json:"payload"`
 }
 
 type webSocketReceiveMessage struct {
@@ -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)
 }

internal/integration/integration_test.go

@@ -66,13 +66,17 @@ func TestSubscription(t *testing.T) {
 	_ = `# @genqlient
 	subscription count { count }`
 
+	authKey := server.AuthKey
+
 	ctx := context.Background()
 	server := server.RunServer()
 	defer server.Close()
 
 	cases := []struct {
 		name           string
 		unsubThreshold time.Duration
+		connParams     map[string]interface{}
+		counterStart   int
 		expected       subscriptionResult
 	}{
 		{
@@ -83,6 +87,18 @@ func TestSubscription(t *testing.T) {
 				serverChannelClosed: true,
 			},
 		},
+		{
+			name:           "server_closed_authorized_user_gets_incremented_counter",
+			unsubThreshold: 5 * time.Second,
+			connParams: map[string]interface{}{
+				authKey: "authorized-user-token",
+			},
+			counterStart: 1000,
+			expected: subscriptionResult{
+				clientUnsubscribed:  false,
+				serverChannelClosed: true,
+			},
+		},
 		{
 			name:           "client_unsubscribed",
 			unsubThreshold: 300 * time.Millisecond,
@@ -95,7 +111,8 @@ func TestSubscription(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			wsClient := newRoundtripWebSocketClient(t, server.URL)
+			wsClient := newRoundtripWebSocketClient(t, server.URL, tc.connParams)
+
 			errChan, err := wsClient.Start(ctx)
 			require.NoError(t, err)
 
@@ -104,7 +121,7 @@ func TestSubscription(t *testing.T) {
 			defer wsClient.Close()
 
 			var (
-				counter = 0
+				counter = tc.counterStart
 				start   = time.Now()
 				result  = subscriptionResult{}
 			)

internal/integration/roundtrip.go

@@ -156,14 +156,14 @@ 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, connectionParams map[string]interface{}) 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),
+		wsWrapped: graphql.NewClientUsingWebSocketWithConnectionParams(endpoint, &MyDialer{Dialer: dialer}, nil, connectionParams),
 		t:         t,
 	}
 }

internal/integration/server/server.go

@@ -156,28 +156,60 @@ func (m mutationResolver) CreateUser(ctx context.Context, input NewUser) (*User,
 	return &newUser, nil
 }
 
+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 (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
+	respCounter := 0
+	if getAuthToken(ctx) == "authorized-user-token" {
+		respCounter = 1000
+	}
+
 	respChan := make(chan int, 1)
 	go func(respChan chan int) {
 		defer close(respChan)
-		counter := 0
+		closeCounter := 0
 		for {
-			if counter == 10 {
+			if closeCounter == 10 {
 				return
 			}
-			respChan <- counter
-			counter++
+			closeCounter++
+			respChan <- respCounter
+			respCounter++
 			time.Sleep(100 * time.Millisecond)
 		}
 	}(respChan)
 	return respChan, nil
 }
 
+type (
+	authTokenCtxKey struct{}
+)
+
+const AuthKey = "authToken"
+
 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)