forked from iisjade/openbci-golang-server
-
Notifications
You must be signed in to change notification settings - Fork 0
/
websocket_connect.go
111 lines (98 loc) · 3 KB
/
websocket_connect.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
100
101
102
103
104
105
106
107
108
109
110
111
/* OpenBCI golang server allows users to control, visualize and store data
collected from the OpenBCI microcontroller.
Copyright (C) 2015 Kevin Schiesser
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package main
import (
"github.com/gorilla/websocket"
"log"
"net/http"
"time"
)
const (
// Time allowed to write a message to the peer.
writeWait = 1 * time.Second
// Time allowed to read the next pong message from the peer.
pongWait = 60 * time.Second
// Send pings to peer with this period. Must be less than pongWait.
pingPeriod = (pongWait * 9) / 10
// Max message size allowed to be written to server over websocket.
maxMessageSize = 512
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 4096 * 16,
}
type WSConn struct {
send chan *message
wsConn *websocket.Conn
}
func NewWSConn(w http.ResponseWriter, r *http.Request) (*WSConn, error) {
if r.Method != "GET" {
http.Error(w, "Method not allowed", 405)
}
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println("Error upgrading ws connection.", err)
}
return &WSConn{
wsConn: conn,
send: make(chan *message, 32),
}, err
}
func (ws *WSConn) write(mt int, payload []byte) error {
ws.wsConn.SetWriteDeadline(time.Now().Add(writeWait))
return ws.wsConn.WriteMessage(mt, payload)
}
func (ws *WSConn) writeJSON(payload *message) error {
ws.wsConn.SetWriteDeadline(time.Now().Add(writeWait))
return ws.wsConn.WriteJSON(payload)
}
func (ws *WSConn) WritePump(h *hub) {
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
h.unregister <- ws
ws.wsConn.Close()
}()
for {
select {
case message := <-ws.send:
if err := ws.writeJSON(message); err != nil {
return
}
case <-ticker.C:
if err := ws.write(websocket.PingMessage, []byte{}); err != nil {
return
}
}
}
}
//ReadPump listens for messages from the browser. In this case,
//the brower is responsible for closing the websocket connection.
//To do so it will write a close conn message to the ReadPump.
func (ws *WSConn) ReadPump(h *hub) {
defer func() {
h.unregister <- ws
ws.wsConn.Close()
}()
ws.wsConn.SetReadLimit(maxMessageSize)
ws.wsConn.SetReadDeadline(time.Now().Add(pongWait))
ws.wsConn.SetPongHandler(func(string) error { ws.wsConn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
for {
_, _, err := ws.wsConn.ReadMessage()
if err != nil {
break
}
}
}