Replace gorilla with melody for websocket API (#1205)

1 files changed, 219 insertions, 0 deletions

diff --git a/vendor/gopkg.in/olahol/melody.v1/session.go b/vendor/gopkg.in/olahol/melody.v1/session.go
new file mode 100644
index 00000000..3997cefe
--- /dev/null
+++ b/vendor/gopkg.in/olahol/melody.v1/session.go
@@ -0,0 +1,219 @@
+package melody
+
+import (
+	"errors"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// Session wrapper around websocket connections.
+type Session struct {
+	Request *http.Request
+	Keys    map[string]interface{}
+	conn    *websocket.Conn
+	output  chan *envelope
+	melody  *Melody
+	open    bool
+	rwmutex *sync.RWMutex
+}
+
+func (s *Session) writeMessage(message *envelope) {
+	if s.closed() {
+		s.melody.errorHandler(s, errors.New("tried to write to closed a session"))
+		return
+	}
+
+	select {
+	case s.output <- message:
+	default:
+		s.melody.errorHandler(s, errors.New("session message buffer is full"))
+	}
+}
+
+func (s *Session) writeRaw(message *envelope) error {
+	if s.closed() {
+		return errors.New("tried to write to a closed session")
+	}
+
+	s.conn.SetWriteDeadline(time.Now().Add(s.melody.Config.WriteWait))
+	err := s.conn.WriteMessage(message.t, message.msg)
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *Session) closed() bool {
+	s.rwmutex.RLock()
+	defer s.rwmutex.RUnlock()
+
+	return !s.open
+}
+
+func (s *Session) close() {
+	if !s.closed() {
+		s.rwmutex.Lock()
+		s.open = false
+		s.conn.Close()
+		close(s.output)
+		s.rwmutex.Unlock()
+	}
+}
+
+func (s *Session) ping() {
+	s.writeRaw(&envelope{t: websocket.PingMessage, msg: []byte{}})
+}
+
+func (s *Session) writePump() {
+	ticker := time.NewTicker(s.melody.Config.PingPeriod)
+	defer ticker.Stop()
+
+loop:
+	for {
+		select {
+		case msg, ok := <-s.output:
+			if !ok {
+				break loop
+			}
+
+			err := s.writeRaw(msg)
+
+			if err != nil {
+				s.melody.errorHandler(s, err)
+				break loop
+			}
+
+			if msg.t == websocket.CloseMessage {
+				break loop
+			}
+
+			if msg.t == websocket.TextMessage {
+				s.melody.messageSentHandler(s, msg.msg)
+			}
+
+			if msg.t == websocket.BinaryMessage {
+				s.melody.messageSentHandlerBinary(s, msg.msg)
+			}
+		case <-ticker.C:
+			s.ping()
+		}
+	}
+}
+
+func (s *Session) readPump() {
+	s.conn.SetReadLimit(s.melody.Config.MaxMessageSize)
+	s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
+
+	s.conn.SetPongHandler(func(string) error {
+		s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
+		s.melody.pongHandler(s)
+		return nil
+	})
+
+	if s.melody.closeHandler != nil {
+		s.conn.SetCloseHandler(func(code int, text string) error {
+			return s.melody.closeHandler(s, code, text)
+		})
+	}
+
+	for {
+		t, message, err := s.conn.ReadMessage()
+
+		if err != nil {
+			s.melody.errorHandler(s, err)
+			break
+		}
+
+		if t == websocket.TextMessage {
+			s.melody.messageHandler(s, message)
+		}
+
+		if t == websocket.BinaryMessage {
+			s.melody.messageHandlerBinary(s, message)
+		}
+	}
+}
+
+// Write writes message to session.
+func (s *Session) Write(msg []byte) error {
+	if s.closed() {
+		return errors.New("session is closed")
+	}
+
+	s.writeMessage(&envelope{t: websocket.TextMessage, msg: msg})
+
+	return nil
+}
+
+// WriteBinary writes a binary message to session.
+func (s *Session) WriteBinary(msg []byte) error {
+	if s.closed() {
+		return errors.New("session is closed")
+	}
+
+	s.writeMessage(&envelope{t: websocket.BinaryMessage, msg: msg})
+
+	return nil
+}
+
+// Close closes session.
+func (s *Session) Close() error {
+	if s.closed() {
+		return errors.New("session is already closed")
+	}
+
+	s.writeMessage(&envelope{t: websocket.CloseMessage, msg: []byte{}})
+
+	return nil
+}
+
+// CloseWithMsg closes the session with the provided payload.
+// Use the FormatCloseMessage function to format a proper close message payload.
+func (s *Session) CloseWithMsg(msg []byte) error {
+	if s.closed() {
+		return errors.New("session is already closed")
+	}
+
+	s.writeMessage(&envelope{t: websocket.CloseMessage, msg: msg})
+
+	return nil
+}
+
+// Set is used to store a new key/value pair exclusivelly for this session.
+// It also lazy initializes s.Keys if it was not used previously.
+func (s *Session) Set(key string, value interface{}) {
+	if s.Keys == nil {
+		s.Keys = make(map[string]interface{})
+	}
+
+	s.Keys[key] = value
+}
+
+// Get returns the value for the given key, ie: (value, true).
+// If the value does not exists it returns (nil, false)
+func (s *Session) Get(key string) (value interface{}, exists bool) {
+	if s.Keys != nil {
+		value, exists = s.Keys[key]
+	}
+
+	return
+}
+
+// MustGet returns the value for the given key if it exists, otherwise it panics.
+func (s *Session) MustGet(key string) interface{} {
+	if value, exists := s.Get(key); exists {
+		return value
+	}
+
+	panic("Key \"" + key + "\" does not exist")
+}
+
+// IsClosed returns the status of the connection.
+func (s *Session) IsClosed() bool {
+	return s.closed()
+}