Merge pull request #13 from functionalfoundry/jannis/auth-token-and-user

Add support for auth tokens and users associated with connections

Author: Zerim

connections.go

@@ -3,6 +3,7 @@ package graphqlws
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/google/uuid"
@@ -30,6 +31,12 @@ const (
 	writeTimeout = 10 * time.Second
 )
 
+// InitMessagePayload defines the parameters of a connection
+// init message.
+type InitMessagePayload struct {
+	AuthToken string `json:"authToken"`
+}
+
 // StartMessagePayload defines the parameters of an operation that
 // a client requests to be started.
 type StartMessagePayload struct {
@@ -59,6 +66,10 @@ func (msg OperationMessage) String() string {
 	return "<invalid>"
 }
 
+// UserFromAuthTokenFunc is a function that resolves an auth token
+// into a user (or returns an error if that isn't possible).
+type UserFromAuthTokenFunc func(token string) (interface{}, error)
+
 // ConnectionEventHandlers define the event handlers for a connection.
 // Event handlers allow other system components to react to events such
 // as the connection closing or an operation being started or stopped.
@@ -81,12 +92,22 @@ type ConnectionEventHandlers struct {
 	StopOperation func(Connection, string)
 }
 
+// ConnectionConfig defines the configuration parameters of a
+// GraphQL WebSocket connection.
+type ConnectionConfig struct {
+	UserFromAuthToken UserFromAuthTokenFunc
+	EventHandlers     ConnectionEventHandlers
+}
+
 // Connection is an interface to represent GraphQL WebSocket connections.
 // Each connection is associated with an ID that is unique to the server.
 type Connection interface {
 	// ID returns the unique ID of the connection.
 	ID() string
 
+	// User returns the user associated with the connection (or nil).
+	User() interface{}
+
 	// SendData sends results of executing an operation (typically a
 	// subscription) to the client.
 	SendData(string, *DataMessagePayload)
@@ -100,11 +121,12 @@ type Connection interface {
  */
 
 type connection struct {
-	id            string
-	ws            *websocket.Conn
-	eventHandlers *ConnectionEventHandlers
-	logger        *log.Entry
-	outgoing      chan OperationMessage
+	id       string
+	ws       *websocket.Conn
+	config   ConnectionConfig
+	logger   *log.Entry
+	outgoing chan OperationMessage
+	user     interface{}
 }
 
 func operationMessageForType(messageType string) OperationMessage {
@@ -116,11 +138,11 @@ func operationMessageForType(messageType string) OperationMessage {
 // NewConnection establishes a GraphQL WebSocket connection. It implements
 // the GraphQL WebSocket protocol by managing its internal state and handling
 // the client-server communication.
-func NewConnection(ws *websocket.Conn, eventHandlers *ConnectionEventHandlers) Connection {
+func NewConnection(ws *websocket.Conn, config ConnectionConfig) Connection {
 	conn := new(connection)
 	conn.id = uuid.New().String()
 	conn.ws = ws
-	conn.eventHandlers = eventHandlers
+	conn.config = config
 	conn.logger = NewLogger("connection/" + conn.id)
 
 	conn.outgoing = make(chan OperationMessage)
@@ -137,6 +159,10 @@ func (conn *connection) ID() string {
 	return conn.id
 }
 
+func (conn *connection) User() interface{} {
+	return conn.user
+}
+
 func (conn *connection) SendData(opID string, data *DataMessagePayload) {
 	msg := operationMessageForType(gqlData)
 	msg.ID = opID
@@ -162,8 +188,8 @@ func (conn *connection) close() {
 	close(conn.outgoing)
 
 	// Notify event handlers
-	if conn.eventHandlers != nil {
-		conn.eventHandlers.Close(conn)
+	if conn.config.EventHandlers.Close != nil {
+		conn.config.EventHandlers.Close(conn)
 	}
 
 	conn.logger.Info("Closed connection")
@@ -238,16 +264,31 @@ func (conn *connection) readLoop() {
 
 		// When the GraphQL WS connection is initiated, send an ACK back
 		case gqlConnectionInit:
-			conn.outgoing <- operationMessageForType(gqlConnectionAck)
+			data := InitMessagePayload{}
+			if err := json.Unmarshal(rawPayload, &data); err != nil {
+				conn.SendError(errors.New("Invalid GQL_CONNECTION_INIT payload"))
+			} else {
+				if conn.config.UserFromAuthToken != nil {
+					user, err := conn.config.UserFromAuthToken(data.AuthToken)
+					if err != nil {
+						conn.SendError(fmt.Errorf("Failed to authenticate user: %v", err))
+					} else {
+						conn.user = user
+						conn.outgoing <- operationMessageForType(gqlConnectionAck)
+					}
+				} else {
+					conn.outgoing <- operationMessageForType(gqlConnectionAck)
+				}
+			}
 
 		// Let event handlers deal with starting operations
 		case gqlStart:
-			if conn.eventHandlers != nil {
+			if conn.config.EventHandlers.StartOperation != nil {
 				data := StartMessagePayload{}
 				if err := json.Unmarshal(rawPayload, &data); err != nil {
 					conn.SendError(errors.New("Invalid GQL_START payload"))
 				} else {
-					errs := conn.eventHandlers.StartOperation(conn, msg.ID, &data)
+					errs := conn.config.EventHandlers.StartOperation(conn, msg.ID, &data)
 					if errs != nil {
 						conn.sendOperationErrors(msg.ID, errs)
 					}
@@ -256,8 +297,8 @@ func (conn *connection) readLoop() {
 
 		// Let event handlers deal with stopping operations
 		case gqlStop:
-			if conn.eventHandlers != nil {
-				conn.eventHandlers.StopOperation(conn, msg.ID)
+			if conn.config.EventHandlers.StopOperation != nil {
+				conn.config.EventHandlers.StopOperation(conn, msg.ID)
 			}
 
 		// When the GraphQL WS connection is terminated by the client,

examples/simple-server/main.go

@@ -31,7 +31,12 @@ func main() {
 
 	// Create subscription manager and GraphQL WS handler
 	subscriptionManager := graphqlws.NewSubscriptionManager(&schema)
-	websocketHandler := graphqlws.NewHandler(subscriptionManager)
+	websocketHandler := graphqlws.NewHandler(graphqlws.HandlerConfig{
+		SubscriptionManager: subscriptionManager,
+		UserFromAuthToken: func(token string) (interface{}, error) {
+			return "Default user", nil
+		},
+	})
 
 	// Serve the GraphQL WS endpoint
 	http.Handle("/subscriptions", websocketHandler)

handler.go

@@ -7,10 +7,16 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+// HandlerConfig stores the configuration of a GraphQL WebSocket handler.
+type HandlerConfig struct {
+	SubscriptionManager SubscriptionManager
+	UserFromAuthToken   UserFromAuthTokenFunc
+}
+
 // NewHandler creates a WebSocket handler for GraphQL WebSocket connections.
 // This handler takes a SubscriptionManager and adds/removes subscriptions
 // as they are started/stopped by the client.
-func NewHandler(subscriptionManager SubscriptionManager) http.Handler {
+func NewHandler(config HandlerConfig) http.Handler {
 	// Create a WebSocket upgrader that requires clients to implement
 	// the "graphql-ws" protocol
 	var upgrader = websocket.Upgrader{
@@ -19,6 +25,7 @@ func NewHandler(subscriptionManager SubscriptionManager) http.Handler {
 	}
 
 	logger := NewLogger("handler")
+	subscriptionManager := config.SubscriptionManager
 
 	// Create a map (used like a set) to manage client connections
 	var connections = make(map[Connection]bool)
@@ -42,40 +49,46 @@ func NewHandler(subscriptionManager SubscriptionManager) http.Handler {
 			}
 
 			// Establish a GraphQL WebSocket connection
-			conn := NewConnection(ws, &ConnectionEventHandlers{
-				Close: func(conn Connection) {
-					logger.WithFields(log.Fields{
-						"conn": conn.ID(),
-					}).Debug("Closing connection")
+			conn := NewConnection(ws, ConnectionConfig{
+				UserFromAuthToken: config.UserFromAuthToken,
+				EventHandlers: ConnectionEventHandlers{
+					Close: func(conn Connection) {
+						logger.WithFields(log.Fields{
+							"conn": conn.ID(),
+							"user": conn.User(),
+						}).Debug("Closing connection")
 
-					subscriptionManager.RemoveSubscriptions(conn)
+						subscriptionManager.RemoveSubscriptions(conn)
 
-					delete(connections, conn)
-				},
-				StartOperation: func(
-					conn Connection,
-					opID string,
-					data *StartMessagePayload,
-				) []error {
-					logger.WithFields(log.Fields{
-						"conn": conn.ID(),
-						"op":   opID,
-					}).Debug("Start operation")
+						delete(connections, conn)
+					},
+					StartOperation: func(
+						conn Connection,
+						opID string,
+						data *StartMessagePayload,
+					) []error {
+						logger.WithFields(log.Fields{
+							"conn": conn.ID(),
+							"op":   opID,
+							"user": conn.User(),
+						}).Debug("Start operation")
 
-					return subscriptionManager.AddSubscription(conn, &Subscription{
-						ID:            opID,
-						Query:         data.Query,
-						Variables:     data.Variables,
-						OperationName: data.OperationName,
-						SendData: func(subscription *Subscription, data *DataMessagePayload) {
-							conn.SendData(opID, data)
-						},
-					})
-				},
-				StopOperation: func(conn Connection, opID string) {
-					subscriptionManager.RemoveSubscription(conn, &Subscription{
-						ID: opID,
-					})
+						return subscriptionManager.AddSubscription(conn, &Subscription{
+							ID:            opID,
+							Query:         data.Query,
+							Variables:     data.Variables,
+							OperationName: data.OperationName,
+							Connection:    conn,
+							SendData: func(subscription *Subscription, data *DataMessagePayload) {
+								conn.SendData(opID, data)
+							},
+						})
+					},
+					StopOperation: func(conn Connection, opID string) {
+						subscriptionManager.RemoveSubscription(conn, &Subscription{
+							ID: opID,
+						})
+					},
 				},
 			})
 			connections[conn] = true