-
Notifications
You must be signed in to change notification settings - Fork 702
/
websocket.go
72 lines (64 loc) · 1.78 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
package xgin
import (
"io"
"net"
"net/http"
"time"
"github.com/gorilla/websocket"
)
//WebSocketConn websocket conn, see websocket.Conn
type WebSocketConn interface {
Subprotocol() string
Close() error
LocalAddr() net.Addr
RemoteAddr() net.Addr
WriteControl(messageType int, data []byte, deadline time.Time) error
NextWriter(messageType int) (io.WriteCloser, error)
WritePreparedMessage(pm *websocket.PreparedMessage) error
WriteMessage(messageType int, data []byte) error
SetWriteDeadline(t time.Time) error
NextReader() (messageType int, r io.Reader, err error)
ReadMessage() (messageType int, p []byte, err error)
SetReadDeadline(t time.Time) error
SetReadLimit(limit int64)
CloseHandler() func(code int, text string) error
SetCloseHandler(h func(code int, text string) error)
PingHandler() func(appData string) error
SetPingHandler(h func(appData string) error)
PongHandler() func(appData string) error
SetPongHandler(h func(appData string) error)
UnderlyingConn() net.Conn
EnableWriteCompression(enable bool)
SetCompressionLevel(level int) error
}
//WebSocketFunc ..
type WebSocketFunc func(WebSocketConn, error)
//WebSocket ..
type WebSocket struct {
Pattern string
Handler WebSocketFunc
*websocket.Upgrader
Header http.Header
}
//Upgrade get upgrage request
func (ws *WebSocket) Upgrade(w http.ResponseWriter, r *http.Request) {
conn, err := ws.Upgrader.Upgrade(w, r, ws.Header)
if err == nil {
defer conn.Close()
}
ws.Handler(conn, err)
}
//WebSocketOption ..
type WebSocketOption func(*WebSocket)
//WebSocketOptions ..
func WebSocketOptions(pattern string, handler WebSocketFunc, opts ...WebSocketOption) *WebSocket {
ws := &WebSocket{
Pattern: pattern,
Handler: handler,
Upgrader: &websocket.Upgrader{},
}
for _, opt := range opts {
opt(ws)
}
return ws
}