Replace gorilla with melody for websocket API (#1205)

vendor/gopkg.in/olahol/melody.v1/melody.go

@@ -0,0 +1,313 @@
+package melody
+
+import (
+	"errors"
+	"net/http"
+	"sync"
+
+	"github.com/gorilla/websocket"
+)
+
+// Close codes defined in RFC 6455, section 11.7.
+// Duplicate of codes from gorilla/websocket for convenience.
+const (
+	CloseNormalClosure           = 1000
+	CloseGoingAway               = 1001
+	CloseProtocolError           = 1002
+	CloseUnsupportedData         = 1003
+	CloseNoStatusReceived        = 1005
+	CloseAbnormalClosure         = 1006
+	CloseInvalidFramePayloadData = 1007
+	ClosePolicyViolation         = 1008
+	CloseMessageTooBig           = 1009
+	CloseMandatoryExtension      = 1010
+	CloseInternalServerErr       = 1011
+	CloseServiceRestart          = 1012
+	CloseTryAgainLater           = 1013
+	CloseTLSHandshake            = 1015
+)
+
+// Duplicate of codes from gorilla/websocket for convenience.
+var validReceivedCloseCodes = map[int]bool{
+	// see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
+
+	CloseNormalClosure:           true,
+	CloseGoingAway:               true,
+	CloseProtocolError:           true,
+	CloseUnsupportedData:         true,
+	CloseNoStatusReceived:        false,
+	CloseAbnormalClosure:         false,
+	CloseInvalidFramePayloadData: true,
+	ClosePolicyViolation:         true,
+	CloseMessageTooBig:           true,
+	CloseMandatoryExtension:      true,
+	CloseInternalServerErr:       true,
+	CloseServiceRestart:          true,
+	CloseTryAgainLater:           true,
+	CloseTLSHandshake:            false,
+}
+
+type handleMessageFunc func(*Session, []byte)
+type handleErrorFunc func(*Session, error)
+type handleCloseFunc func(*Session, int, string) error
+type handleSessionFunc func(*Session)
+type filterFunc func(*Session) bool
+
+// Melody implements a websocket manager.
+type Melody struct {
+	Config                   *Config
+	Upgrader                 *websocket.Upgrader
+	messageHandler           handleMessageFunc
+	messageHandlerBinary     handleMessageFunc
+	messageSentHandler       handleMessageFunc
+	messageSentHandlerBinary handleMessageFunc
+	errorHandler             handleErrorFunc
+	closeHandler             handleCloseFunc
+	connectHandler           handleSessionFunc
+	disconnectHandler        handleSessionFunc
+	pongHandler              handleSessionFunc
+	hub                      *hub
+}
+
+// New creates a new melody instance with default Upgrader and Config.
+func New() *Melody {
+	upgrader := &websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin:     func(r *http.Request) bool { return true },
+	}
+
+	hub := newHub()
+
+	go hub.run()
+
+	return &Melody{
+		Config:                   newConfig(),
+		Upgrader:                 upgrader,
+		messageHandler:           func(*Session, []byte) {},
+		messageHandlerBinary:     func(*Session, []byte) {},
+		messageSentHandler:       func(*Session, []byte) {},
+		messageSentHandlerBinary: func(*Session, []byte) {},
+		errorHandler:             func(*Session, error) {},
+		closeHandler:             nil,
+		connectHandler:           func(*Session) {},
+		disconnectHandler:        func(*Session) {},
+		pongHandler:              func(*Session) {},
+		hub:                      hub,
+	}
+}
+
+// HandleConnect fires fn when a session connects.
+func (m *Melody) HandleConnect(fn func(*Session)) {
+	m.connectHandler = fn
+}
+
+// HandleDisconnect fires fn when a session disconnects.
+func (m *Melody) HandleDisconnect(fn func(*Session)) {
+	m.disconnectHandler = fn
+}
+
+// HandlePong fires fn when a pong is received from a session.
+func (m *Melody) HandlePong(fn func(*Session)) {
+	m.pongHandler = fn
+}
+
+// HandleMessage fires fn when a text message comes in.
+func (m *Melody) HandleMessage(fn func(*Session, []byte)) {
+	m.messageHandler = fn
+}
+
+// HandleMessageBinary fires fn when a binary message comes in.
+func (m *Melody) HandleMessageBinary(fn func(*Session, []byte)) {
+	m.messageHandlerBinary = fn
+}
+
+// HandleSentMessage fires fn when a text message is successfully sent.
+func (m *Melody) HandleSentMessage(fn func(*Session, []byte)) {
+	m.messageSentHandler = fn
+}
+
+// HandleSentMessageBinary fires fn when a binary message is successfully sent.
+func (m *Melody) HandleSentMessageBinary(fn func(*Session, []byte)) {
+	m.messageSentHandlerBinary = fn
+}
+
+// HandleError fires fn when a session has an error.
+func (m *Melody) HandleError(fn func(*Session, error)) {
+	m.errorHandler = fn
+}
+
+// HandleClose sets the handler for close messages received from the session.
+// The code argument to h is the received close code or CloseNoStatusReceived
+// if the close message is empty. The default close handler sends a close frame
+// back to the session.
+//
+// The application must read the connection to process close messages as
+// described in the section on Control Frames above.
+//
+// The connection read methods return a CloseError when a close frame is
+// received. Most applications should handle close messages as part of their
+// normal error handling. Applications should only set a close handler when the
+// application must perform some action before sending a close frame back to
+// the session.
+func (m *Melody) HandleClose(fn func(*Session, int, string) error) {
+	if fn != nil {
+		m.closeHandler = fn
+	}
+}
+
+// HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance.
+func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error {
+	return m.HandleRequestWithKeys(w, r, nil)
+}
+
+// HandleRequestWithKeys does the same as HandleRequest but populates session.Keys with keys.
+func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]interface{}) error {
+	if m.hub.closed() {
+		return errors.New("melody instance is closed")
+	}
+
+	conn, err := m.Upgrader.Upgrade(w, r, nil)
+
+	if err != nil {
+		return err
+	}
+
+	session := &Session{
+		Request: r,
+		Keys:    keys,
+		conn:    conn,
+		output:  make(chan *envelope, m.Config.MessageBufferSize),
+		melody:  m,
+		open:    true,
+		rwmutex: &sync.RWMutex{},
+	}
+
+	m.hub.register <- session
+
+	m.connectHandler(session)
+
+	go session.writePump()
+
+	session.readPump()
+
+	if !m.hub.closed() {
+		m.hub.unregister <- session
+	}
+
+	session.close()
+
+	m.disconnectHandler(session)
+
+	return nil
+}
+
+// Broadcast broadcasts a text message to all sessions.
+func (m *Melody) Broadcast(msg []byte) error {
+	if m.hub.closed() {
+		return errors.New("melody instance is closed")
+	}
+
+	message := &envelope{t: websocket.TextMessage, msg: msg}
+	m.hub.broadcast <- message
+
+	return nil
+}
+
+// BroadcastFilter broadcasts a text message to all sessions that fn returns true for.
+func (m *Melody) BroadcastFilter(msg []byte, fn func(*Session) bool) error {
+	if m.hub.closed() {
+		return errors.New("melody instance is closed")
+	}
+
+	message := &envelope{t: websocket.TextMessage, msg: msg, filter: fn}
+	m.hub.broadcast <- message
+
+	return nil
+}
+
+// BroadcastOthers broadcasts a text message to all sessions except session s.
+func (m *Melody) BroadcastOthers(msg []byte, s *Session) error {
+	return m.BroadcastFilter(msg, func(q *Session) bool {
+		return s != q
+	})
+}
+
+// BroadcastMultiple broadcasts a text message to multiple sessions given in the sessions slice.
+func (m *Melody) BroadcastMultiple(msg []byte, sessions []*Session) error {
+	for _, sess := range sessions {
+		if writeErr := sess.Write(msg); writeErr != nil {
+			return writeErr
+		}
+	}
+	return nil
+}
+
+// BroadcastBinary broadcasts a binary message to all sessions.
+func (m *Melody) BroadcastBinary(msg []byte) error {
+	if m.hub.closed() {
+		return errors.New("melody instance is closed")
+	}
+
+	message := &envelope{t: websocket.BinaryMessage, msg: msg}
+	m.hub.broadcast <- message
+
+	return nil
+}
+
+// BroadcastBinaryFilter broadcasts a binary message to all sessions that fn returns true for.
+func (m *Melody) BroadcastBinaryFilter(msg []byte, fn func(*Session) bool) error {
+	if m.hub.closed() {
+		return errors.New("melody instance is closed")
+	}
+
+	message := &envelope{t: websocket.BinaryMessage, msg: msg, filter: fn}
+	m.hub.broadcast <- message
+
+	return nil
+}
+
+// BroadcastBinaryOthers broadcasts a binary message to all sessions except session s.
+func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) error {
+	return m.BroadcastBinaryFilter(msg, func(q *Session) bool {
+		return s != q
+	})
+}
+
+// Close closes the melody instance and all connected sessions.
+func (m *Melody) Close() error {
+	if m.hub.closed() {
+		return errors.New("melody instance is already closed")
+	}
+
+	m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: []byte{}}
+
+	return nil
+}
+
+// CloseWithMsg closes the melody instance with the given close payload and all connected sessions.
+// Use the FormatCloseMessage function to format a proper close message payload.
+func (m *Melody) CloseWithMsg(msg []byte) error {
+	if m.hub.closed() {
+		return errors.New("melody instance is already closed")
+	}
+
+	m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: msg}
+
+	return nil
+}
+
+// Len return the number of connected sessions.
+func (m *Melody) Len() int {
+	return m.hub.len()
+}
+
+// IsClosed returns the status of the melody instance.
+func (m *Melody) IsClosed() bool {
+	return m.hub.closed()
+}
+
+// FormatCloseMessage formats closeCode and text as a WebSocket close message.
+func FormatCloseMessage(closeCode int, text string) []byte {
+	return websocket.FormatCloseMessage(closeCode, text)
+}