-
Notifications
You must be signed in to change notification settings - Fork 667
/
websocket.go
99 lines (88 loc) · 1.95 KB
/
websocket.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
package binance
import (
"net/http"
"time"
"github.com/gorilla/websocket"
)
// WsHandler handle raw websocket message
type WsHandler func(message []byte)
// ErrHandler handles errors
type ErrHandler func(err error)
// WsConfig webservice configuration
type WsConfig struct {
Endpoint string
}
func newWsConfig(endpoint string) *WsConfig {
return &WsConfig{
Endpoint: endpoint,
}
}
var wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (doneC, stopC chan struct{}, err error) {
Dialer := websocket.Dialer{
Proxy: http.ProxyFromEnvironment,
HandshakeTimeout: 45 * time.Second,
EnableCompression: false,
}
c, _, err := Dialer.Dial(cfg.Endpoint, nil)
if err != nil {
return nil, nil, err
}
c.SetReadLimit(655350)
doneC = make(chan struct{})
stopC = make(chan struct{})
go func() {
// This function will exit either on error from
// websocket.Conn.ReadMessage or when the stopC channel is
// closed by the client.
defer close(doneC)
if WebsocketKeepalive {
keepAlive(c, WebsocketTimeout)
}
// Wait for the stopC channel to be closed. We do that in a
// separate goroutine because ReadMessage is a blocking
// operation.
silent := false
go func() {
select {
case <-stopC:
silent = true
case <-doneC:
}
c.Close()
}()
for {
_, message, err := c.ReadMessage()
if err != nil {
if !silent {
errHandler(err)
}
return
}
handler(message)
}
}()
return
}
func keepAlive(c *websocket.Conn, timeout time.Duration) {
ticker := time.NewTicker(timeout)
lastResponse := time.Now()
c.SetPongHandler(func(msg string) error {
lastResponse = time.Now()
return nil
})
go func() {
defer ticker.Stop()
for {
deadline := time.Now().Add(10 * time.Second)
err := c.WriteControl(websocket.PingMessage, []byte{}, deadline)
if err != nil {
return
}
<-ticker.C
if time.Since(lastResponse) > timeout {
c.Close()
return
}
}
}()
}