[websocket] Add a config to reject initial connection

handler/graphql.go

@@ -40,6 +40,7 @@ type Config struct {
 	tracer                          graphql.Tracer
 	complexityLimit                 int
 	complexityLimitFunc             graphql.ComplexityLimitFunc
+	websocketOnInitFunc             func(ctx context.Context, initPayload InitPayload) bool
 	disableIntrospection            bool
 	connectionKeepAlivePingInterval time.Duration
 	uploadMaxMemory                 int64
@@ -250,6 +251,14 @@ func (tw *tracerWrapper) EndOperationExecution(ctx context.Context) {
 	tw.tracer1.EndOperationExecution(ctx)
 }
 
+// WebsocketOnInitFunc is called when the server receives connection init message from the client.
+// This can be used to check initial payload to see whether to accept the websocket connection.
+func WebsocketOnInitFunc(websocketOnInitFunc func(ctx context.Context, initPayload InitPayload) bool) Option {
+	return func(cfg *Config) {
+		cfg.websocketOnInitFunc = websocketOnInitFunc
+	}
+}
+
 // CacheSize sets the maximum size of the query cache.
 // If size is less than or equal to 0, the cache is disabled.
 func CacheSize(size int) Option {

handler/websocket.go

@@ -12,7 +12,7 @@ import (
 
 	"github.com/99designs/gqlgen/graphql"
 	"github.com/gorilla/websocket"
-	"github.com/hashicorp/golang-lru"
+	lru "github.com/hashicorp/golang-lru"
 	"github.com/vektah/gqlparser"
 	"github.com/vektah/gqlparser/ast"
 	"github.com/vektah/gqlparser/gqlerror"
@@ -94,6 +94,15 @@ func (c *wsConnection) init() bool {
 			}
 		}
 
+		if c.cfg.websocketOnInitFunc != nil {
+			isValidPayload := c.cfg.websocketOnInitFunc(c.ctx, c.initPayload)
+			if !isValidPayload {
+				c.sendConnectionError("invalid init payload")
+				c.close(websocket.CloseNormalClosure, "terminated")
+				return false
+			}
+		}
+
 		c.write(&operationMessage{Type: connectionAckMsg})
 	case connectionTerminateMsg:
 		c.close(websocket.CloseNormalClosure, "terminated")

handler/websocket_test.go

@@ -1,6 +1,7 @@
 package handler
 
 import (
+	"context"
 	"encoding/json"
 	"net/http/httptest"
 	"strings"
@@ -158,6 +159,55 @@ func TestWebsocketWithKeepAlive(t *testing.T) {
 	})
 }
 
+func TestWebsocketOnInitFunc(t *testing.T) {
+	next := make(chan struct{})
+
+	t.Run("accept connection if WebsocketOnInitFunc is NOT provided", func(t *testing.T) {
+		h := GraphQL(&executableSchemaStub{next})
+		srv := httptest.NewServer(h)
+		defer srv.Close()
+
+		c := wsConnect(srv.URL)
+		defer c.Close()
+
+		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
+
+		require.Equal(t, connectionAckMsg, readOp(c).Type)
+	})
+
+	t.Run("accept connection if WebsocketOnInitFunc is provided and is accepting connection", func(t *testing.T) {
+		h := GraphQL(&executableSchemaStub{next}, WebsocketOnInitFunc(func(ctx context.Context, initPayload InitPayload) bool {
+			return true
+		}))
+		srv := httptest.NewServer(h)
+		defer srv.Close()
+
+		c := wsConnect(srv.URL)
+		defer c.Close()
+
+		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
+
+		require.Equal(t, connectionAckMsg, readOp(c).Type)
+	})
+
+	t.Run("reject connection if WebsocketOnInitFunc is provided and is accepting connection", func(t *testing.T) {
+		h := GraphQL(&executableSchemaStub{next}, WebsocketOnInitFunc(func(ctx context.Context, initPayload InitPayload) bool {
+			return false
+		}))
+		srv := httptest.NewServer(h)
+		defer srv.Close()
+
+		c := wsConnect(srv.URL)
+		defer c.Close()
+
+		require.NoError(t, c.WriteJSON(&operationMessage{Type: connectionInitMsg}))
+
+		msg := readOp(c)
+		require.Equal(t, connectionErrorMsg, msg.Type)
+		require.Equal(t, `{"message":"invalid init payload"}`, string(msg.Payload))
+	})
+}
+
 func wsConnect(url string) *websocket.Conn {
 	c, _, err := websocket.DefaultDialer.Dial(strings.Replace(url, "http://", "ws://", -1), nil)
 	if err != nil {