/
websocket_server.go
executable file
·72 lines (63 loc) · 1.98 KB
/
websocket_server.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
package server
import (
"net/http"
"strconv"
"github.com/gearghost/go-protoo/logger"
"github.com/gearghost/go-protoo/transport"
"github.com/gorilla/websocket"
)
type WebSocketServerConfig struct {
Host string
Port int
CertFile string
KeyFile string
HTMLRoot string
WebSocketPath string
}
func DefaultConfig() WebSocketServerConfig {
return WebSocketServerConfig{
Host: "0.0.0.0",
Port: 8443,
HTMLRoot: ".",
WebSocketPath: "/ws",
}
}
type WebSocketServer struct {
handleWebSocket func(ws *transport.WebSocketTransport, request *http.Request)
// Websocket upgrader
upgrader websocket.Upgrader
}
func NewWebSocketServer(handler func(ws *transport.WebSocketTransport, request *http.Request)) *WebSocketServer {
var server = &WebSocketServer{
handleWebSocket: handler,
}
server.upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
return server
}
func (server *WebSocketServer) handleWebSocketRequest(writer http.ResponseWriter, request *http.Request) {
responseHeader := http.Header{}
responseHeader.Add("Sec-WebSocket-Protocol", "protoo")
socket, err := server.upgrader.Upgrade(writer, request, responseHeader)
if err != nil {
panic(err)
}
wsTransport := transport.NewWebSocketTransport(socket)
server.handleWebSocket(wsTransport, request)
wsTransport.ReadMessage()
}
func (server *WebSocketServer) Bind(cfg WebSocketServerConfig) {
// Websocket handle func
http.HandleFunc(cfg.WebSocketPath, server.handleWebSocketRequest)
http.Handle("/", http.FileServer(http.Dir(cfg.HTMLRoot)))
if cfg.CertFile == "" || cfg.KeyFile == "" {
logger.Infof("non-TLS WebSocketServer listening on: %s:%d", cfg.Host, cfg.Port)
panic(http.ListenAndServe(cfg.Host+":"+strconv.Itoa(cfg.Port), nil))
} else {
logger.Infof("TLS WebSocketServer listening on: %s:%d", cfg.Host, cfg.Port)
panic(http.ListenAndServeTLS(cfg.Host+":"+strconv.Itoa(cfg.Port), cfg.CertFile, cfg.KeyFile, nil))
}
}