Flexible Authentication Hook

graphqlws/http.go

@@ -1,9 +1,8 @@
 package graphqlws
 
 import (
-	"net/http"
-
 	"github.com/gorilla/websocket"
+	"net/http"
 
 	"github.com/graph-gophers/graphql-transport-ws/graphqlws/internal/connection"
 )
@@ -17,6 +16,11 @@ var upgrader = websocket.Upgrader{
 
 // NewHandlerFunc returns an http.HandlerFunc that supports GraphQL over websockets
 func NewHandlerFunc(svc connection.GraphQLService, httpHandler http.Handler) http.HandlerFunc {
+	return NewHandlerFuncWithAuth(svc, httpHandler, nil)
+}
+
+// NewHandlerFunc returns an http.HandlerFunc that supports GraphQL over websockets
+func NewHandlerFuncWithAuth(svc connection.GraphQLService, httpHandler http.Handler, authFunc connection.AuthenticateFunc) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		for _, subprotocol := range websocket.Subprotocols(r) {
 			if subprotocol == "graphql-ws" {
@@ -30,7 +34,7 @@ func NewHandlerFunc(svc connection.GraphQLService, httpHandler http.Handler) htt
 					return
 				}
 
-				go connection.Connect(ws, svc)
+				go connection.Connect(ws, svc, connection.Authentication(r, authFunc))
 				return
 			}
 		}

graphqlws/internal/connection/connection.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net/http"
 	"time"
 )
 
@@ -47,18 +48,28 @@ type startMessagePayload struct {
 	Variables     map[string]interface{} `json:"variables"`
 }
 
-type initMessagePayload struct{}
-
 // GraphQLService interface
 type GraphQLService interface {
 	Subscribe(ctx context.Context, document string, operationName string, variableValues map[string]interface{}) (payloads <-chan interface{}, err error)
 }
 
+type AuthenticateFunc func(ctx context.Context, r *http.Request, payload map[string]interface{}) error
+
 type connection struct {
-	cancel       func()
-	service      GraphQLService
-	writeTimeout time.Duration
-	ws           wsConnection
+	cancel           func()
+	service          GraphQLService
+	writeTimeout     time.Duration
+	ws               wsConnection
+	authenticated    bool
+	authenticateFunc AuthenticateFunc
+	request          *http.Request
+}
+
+func Authentication(r *http.Request, f AuthenticateFunc) func(conn *connection) {
+	return func(conn *connection) {
+		conn.authenticateFunc = f
+		conn.request = r
+	}
 }
 
 // ReadLimit limits the maximum size of incoming messages
@@ -159,15 +170,31 @@ func (conn *connection) readLoop(ctx context.Context, send sendFunc) {
 
 		switch msg.Type {
 		case typeConnectionInit:
-			var initMsg initMessagePayload
+
+			var initMsg map[string]interface{}
 			if err := json.Unmarshal(msg.Payload, &initMsg); err != nil {
 				ep := errPayload(fmt.Errorf("invalid payload for type: %s", msg.Type))
 				send("", typeConnectionError, ep)
 				continue
 			}
+
+			if conn.authenticateFunc != nil {
+				if err := conn.authenticateFunc(ctx, conn.request, initMsg); err != nil {
+					send("", typeConnectionError, errPayload(err))
+					continue
+				}
+			}
+			conn.authenticated = true
 			send("", typeConnectionAck, nil)
 
 		case typeStart:
+
+			if !conn.authenticated && conn.authenticateFunc != nil {
+				ep := errPayload(errors.New("authentication required."))
+				send("", typeConnectionError, ep)
+				continue
+			}
+
 			// TODO: check an operation with the same ID hasn't been started already
 			if msg.ID == "" {
 				ep := errPayload(errors.New("missing ID for start operation"))