82 lines
1.6 KiB
Go
82 lines
1.6 KiB
Go
|
package websockets
|
||
|
|
|
||
|
|
import "github.com/gorilla/websocket"
|
||
|
|
import (
|
||
|
|
"time"
|
||
|
|
|
||
|
|
log "github.com/sirupsen/logrus"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Socket struct {
|
||
|
|
collection *Collection
|
||
|
|
|
||
|
|
conn *websocket.Conn
|
||
|
|
|
||
|
|
send chan []byte
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Socket) readPump() {
|
||
|
|
defer func() {
|
||
|
|
s.collection.unregister <- s
|
||
|
|
s.conn.Close()
|
||
|
|
}()
|
||
|
|
s.conn.SetReadLimit(maxMessageSize)
|
||
|
|
s.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||
|
|
s.conn.SetPongHandler(func(string) error {
|
||
|
|
s.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||
|
|
return nil
|
||
|
|
})
|
||
|
|
for {
|
||
|
|
t, m, err := s.conn.ReadMessage()
|
||
|
|
if err != nil {
|
||
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
|
||
|
|
log.WithError(err).Debug("Websocket closed unexpectedly.")
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
// Handle websocket responses somehow
|
||
|
|
if t == websocket.TextMessage {
|
||
|
|
log.Debug("Received websocket message: " + string(m))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Socket) writePump() {
|
||
|
|
ticker := time.NewTicker(pingPeriod)
|
||
|
|
defer func() {
|
||
|
|
ticker.Stop()
|
||
|
|
s.conn.Close()
|
||
|
|
}()
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case m, ok := <-s.send:
|
||
|
|
s.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||
|
|
if !ok {
|
||
|
|
// The collection closed the channel
|
||
|
|
s.conn.WriteMessage(websocket.CloseMessage, []byte{})
|
||
|
|
return
|
||
|
|
}
|
||
|
|
w, err := s.conn.NextWriter(websocket.TextMessage)
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
w.Write([]byte{'['})
|
||
|
|
w.Write(m)
|
||
|
|
n := len(s.send) - 1
|
||
|
|
for i := 0; i < n; i++ {
|
||
|
|
w.Write([]byte{','})
|
||
|
|
w.Write(<-s.send)
|
||
|
|
}
|
||
|
|
w.Write([]byte{']'})
|
||
|
|
if err := w.Close(); err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
case <-ticker.C:
|
||
|
|
s.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||
|
|
if err := s.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|