-
Notifications
You must be signed in to change notification settings - Fork 0
/
wsConnClient.go
134 lines (103 loc) · 2.86 KB
/
wsConnClient.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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package connection
import (
"fmt"
"sync"
"github.com/Dharitri-org/drtg-communication/websocket/data"
logger "github.com/Dharitri-org/drtg-logger"
"github.com/gorilla/websocket"
)
var log = logger.GetOrCreate("connection")
type wsConnClient struct {
mut sync.RWMutex
conn *websocket.Conn
clientID string
}
// NewWSConnClient creates a new wrapper over a websocket connection
func NewWSConnClient() *wsConnClient {
return &wsConnClient{}
}
// NewWSConnClientWithConn creates a new wrapper over a provided websocket connection
func NewWSConnClientWithConn(conn *websocket.Conn) *wsConnClient {
wsc := &wsConnClient{
conn: conn,
}
wsc.clientID = fmt.Sprintf("%p", wsc)
return wsc
}
// OpenConnection will open a new client with a background context
func (wsc *wsConnClient) OpenConnection(url string) error {
wsc.mut.Lock()
defer wsc.mut.Unlock()
if wsc.conn != nil {
return data.ErrConnectionAlreadyOpen
}
var err error
wsc.conn, _, err = websocket.DefaultDialer.Dial(url, nil)
if err != nil {
return err
}
return nil
}
// ReadMessage calls the underlying reading message ws connection func
func (wsc *wsConnClient) ReadMessage() (messageType int, p []byte, err error) {
conn, err := wsc.getConn()
if err != nil {
return 0, nil, err
}
return conn.ReadMessage()
}
// WriteMessage calls the underlying write message ws connection func
func (wsc *wsConnClient) WriteMessage(messageType int, payload []byte) error {
wsc.mut.Lock()
defer wsc.mut.Unlock()
if wsc.conn == nil {
return data.ErrConnectionNotOpen
}
return wsc.conn.WriteMessage(messageType, payload)
}
// IsOpen will return true if the connection is open, false otherwise
func (wsc *wsConnClient) IsOpen() bool {
wsc.mut.RLock()
defer wsc.mut.RUnlock()
return wsc.conn != nil
}
func (wsc *wsConnClient) getConn() (*websocket.Conn, error) {
wsc.mut.RLock()
defer wsc.mut.RUnlock()
if wsc.conn == nil {
return nil, data.ErrConnectionNotOpen
}
conn := wsc.conn
return conn, nil
}
// GetID will return the unique id of the client
func (wsc *wsConnClient) GetID() string {
return wsc.clientID
}
// Close will try to cleanly close the connection, if possible
func (wsc *wsConnClient) Close() error {
// critical section
wsc.mut.Lock()
defer wsc.mut.Unlock()
if wsc.conn == nil {
return data.ErrConnectionNotOpen
}
log.Debug("closing ws connection...")
//Cleanly close the connection by sending a close message and then
//waiting (with timeout) for the server to close the connection.
err := wsc.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
if err != nil {
log.Trace("cannot send close message", "error", err)
}
wsc.conn.CloseHandler()
err = wsc.conn.Close()
if err != nil {
return err
}
wsc.conn = nil
return nil
}
// IsInterfaceNil -
func (wsc *wsConnClient) IsInterfaceNil() bool {
return wsc == nil
}