75 lines
1.5 KiB
Go
75 lines
1.5 KiB
Go
package ws
|
|
|
|
import (
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
const (
|
|
writeWait = 10 * time.Second
|
|
pongWait = 60 * time.Second
|
|
pingPeriod = (pongWait * 9) / 10
|
|
)
|
|
|
|
type Client struct {
|
|
hub *Hub
|
|
conn *websocket.Conn
|
|
send chan []byte
|
|
}
|
|
|
|
func NewClient(hub *Hub, conn *websocket.Conn) *Client {
|
|
return &Client{
|
|
hub: hub,
|
|
conn: conn,
|
|
send: make(chan []byte, 256),
|
|
}
|
|
}
|
|
|
|
// ReadPump keeps the connection alive and handles pong frames.
|
|
// We don't expect real messages from the client (read-only WS).
|
|
func (c *Client) ReadPump() {
|
|
defer func() {
|
|
c.hub.Unregister(c)
|
|
c.conn.Close()
|
|
}()
|
|
c.conn.SetReadLimit(512)
|
|
c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
|
c.conn.SetPongHandler(func(string) error {
|
|
c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
|
return nil
|
|
})
|
|
for {
|
|
if _, _, err := c.conn.ReadMessage(); err != nil {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// WritePump sends messages from hub to the client and pings to keep alive.
|
|
func (c *Client) WritePump() {
|
|
ticker := time.NewTicker(pingPeriod)
|
|
defer func() {
|
|
ticker.Stop()
|
|
c.conn.Close()
|
|
}()
|
|
for {
|
|
select {
|
|
case msg, ok := <-c.send:
|
|
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
|
if !ok {
|
|
c.conn.WriteMessage(websocket.CloseMessage, nil)
|
|
return
|
|
}
|
|
if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
|
|
return
|
|
}
|
|
case <-ticker.C:
|
|
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
|
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|