websocketserver/handlers.go

package websocketserver

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/gorilla/websocket"
	"gitlab.com/gitlab-org/webide-file-sync/common"
)

type handlers map[string]handlerWithPayload
type connectFunc handlerWithoutPayload
type disconnectFunc handlerWithoutPayload
type handlerWithoutPayload func(*websocket.Conn) *common.Response
type handlerWithPayload func(*websocket.Conn, json.RawMessage) *common.Response

func NoPayload(f handlerWithoutPayload) handlerWithPayload {
	return func(c *websocket.Conn, mustBeEmpty json.RawMessage) *common.Response {
		var jsonNull = json.RawMessage("null")
		// Check that mustBeNil is nil, send an error response if not
		if mustBeEmpty != nil && len(mustBeEmpty) > 0 && !bytes.Equal(mustBeEmpty, jsonNull) {
			return &common.Response{StatusCode: http.StatusBadRequest, ErrorMessage: "protocol error"}
		}

		return f(c)
	}
}

func (s *Server) closeHandler(code int, text string) error {
	for _, n := range s.namespaces {
		if n.disconnectHandler != nil {
			n.disconnectHandler(s.conn)
		}
	}
	// setting the error to nil in order to let the websocket.CloseError bubble up
	return nil
}

func (s *Server) pongHandler(string) error {
	s.setReadDeadline(s.pongWait)
	return nil
}

// OnConnect registers a callback function in a namespace on a connection event
func (s *Server) OnConnect(namespace string, fn connectFunc) error {
	if err := s.initNamespace(namespace); err != nil {
		return err
	}

	s.namespaces[namespace].connectHandler = fn
	return nil
}

// OnDisconnect registers a callback function in a namespace on a disconnection event
func (s *Server) OnDisconnect(namespaceName string, fn disconnectFunc) error {
	if err := s.initNamespace(namespaceName); err != nil {
		return err
	}

	s.namespaces[namespaceName].disconnectHandler = fn
	return nil
}

// OnEvent registers a function handler per event on a given namespace
func (s *Server) OnEvent(namespaceName string, event string, f handlerWithPayload) error {
	if event == "" {
		return fmt.Errorf("event name can not be empty")
	}

	if err := s.initNamespace(namespaceName); err != nil {
		return err
	}

	s.namespaces[namespaceName].eventHandlers[event] = f
	return nil
}

func (s *Server) initNamespace(namespaceName string) error {
	if namespaceName == "" {
		return fmt.Errorf("namespace cannot be empty")
	}

	if s.namespaces[namespaceName] != nil {
		return nil
	}

	s.namespaces[namespaceName] = &namespace{eventHandlers: make(handlers)}
	return nil
}

func (s *Server) callEventHandler(event string, n *namespace, payload json.RawMessage) (*common.Response, error) {
	if event == "" {
		return nil, fmt.Errorf("event cannot be empty")
	}

	// Check if the handler for that event exists
	eventFn := n.eventHandlers[event]
	if eventFn == nil {
		return nil, fmt.Errorf("no callback associated with event %s", event)
	}

	return eventFn(s.conn, payload), nil
}