github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/model/websocket_client.go (about) 1 // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved. 2 // See LICENSE.txt for license information. 3 4 package model 5 6 import ( 7 "bytes" 8 "encoding/json" 9 "net/http" 10 "sync/atomic" 11 "time" 12 13 "github.com/gorilla/websocket" 14 ) 15 16 const ( 17 SOCKET_MAX_MESSAGE_SIZE_KB = 8 * 1024 // 8KB 18 PING_TIMEOUT_BUFFER_SECONDS = 5 19 ) 20 21 type msgType int 22 23 const ( 24 msgTypeJSON msgType = iota + 1 25 msgTypePong 26 ) 27 28 type writeMessage struct { 29 msgType msgType 30 data interface{} 31 } 32 33 const avgReadMsgSizeBytes = 1024 34 35 // WebSocketClient stores the necessary information required to 36 // communicate with a WebSocket endpoint. 37 // A client must read from PingTimeoutChannel, EventChannel and ResponseChannel to prevent 38 // deadlocks from occurring in the program. 39 type WebSocketClient struct { 40 Url string // The location of the server like "ws://localhost:8065" 41 ApiUrl string // The API location of the server like "ws://localhost:8065/api/v3" 42 ConnectUrl string // The WebSocket URL to connect to like "ws://localhost:8065/api/v3/path/to/websocket" 43 Conn *websocket.Conn // The WebSocket connection 44 AuthToken string // The token used to open the WebSocket connection 45 Sequence int64 // The ever-incrementing sequence attached to each WebSocket action 46 PingTimeoutChannel chan bool // The channel used to signal ping timeouts 47 EventChannel chan *WebSocketEvent // The channel used to receive various events pushed from the server. For example: typing, posted 48 ResponseChannel chan *WebSocketResponse // The channel used to receive responses for requests made to the server 49 ListenError *AppError // A field that is set if there was an abnormal closure of the WebSocket connection 50 writeChan chan writeMessage 51 52 pingTimeoutTimer *time.Timer 53 quitPingWatchdog chan struct{} 54 55 quitWriterChan chan struct{} 56 resetTimerChan chan struct{} 57 closed int32 58 } 59 60 // NewWebSocketClient constructs a new WebSocket client with convenience 61 // methods for talking to the server. 62 func NewWebSocketClient(url, authToken string) (*WebSocketClient, *AppError) { 63 return NewWebSocketClientWithDialer(websocket.DefaultDialer, url, authToken) 64 } 65 66 // NewWebSocketClientWithDialer constructs a new WebSocket client with convenience 67 // methods for talking to the server using a custom dialer. 68 func NewWebSocketClientWithDialer(dialer *websocket.Dialer, url, authToken string) (*WebSocketClient, *AppError) { 69 conn, _, err := dialer.Dial(url+API_URL_SUFFIX+"/websocket", nil) 70 if err != nil { 71 return nil, NewAppError("NewWebSocketClient", "model.websocket_client.connect_fail.app_error", nil, err.Error(), http.StatusInternalServerError) 72 } 73 74 client := &WebSocketClient{ 75 Url: url, 76 ApiUrl: url + API_URL_SUFFIX, 77 ConnectUrl: url + API_URL_SUFFIX + "/websocket", 78 Conn: conn, 79 AuthToken: authToken, 80 Sequence: 1, 81 PingTimeoutChannel: make(chan bool, 1), 82 EventChannel: make(chan *WebSocketEvent, 100), 83 ResponseChannel: make(chan *WebSocketResponse, 100), 84 writeChan: make(chan writeMessage), 85 quitPingWatchdog: make(chan struct{}), 86 quitWriterChan: make(chan struct{}), 87 resetTimerChan: make(chan struct{}), 88 } 89 90 client.configurePingHandling() 91 go client.writer() 92 93 client.SendMessage(WEBSOCKET_AUTHENTICATION_CHALLENGE, map[string]interface{}{"token": authToken}) 94 95 return client, nil 96 } 97 98 // NewWebSocketClient4 constructs a new WebSocket client with convenience 99 // methods for talking to the server. Uses the v4 endpoint. 100 func NewWebSocketClient4(url, authToken string) (*WebSocketClient, *AppError) { 101 return NewWebSocketClient4WithDialer(websocket.DefaultDialer, url, authToken) 102 } 103 104 // NewWebSocketClient4WithDialer constructs a new WebSocket client with convenience 105 // methods for talking to the server using a custom dialer. Uses the v4 endpoint. 106 func NewWebSocketClient4WithDialer(dialer *websocket.Dialer, url, authToken string) (*WebSocketClient, *AppError) { 107 return NewWebSocketClientWithDialer(dialer, url, authToken) 108 } 109 110 // Connect creates a websocket connection with the given ConnectUrl. 111 // This is racy and error-prone should not be used. Use any of the New* functions to create a websocket. 112 func (wsc *WebSocketClient) Connect() *AppError { 113 return wsc.ConnectWithDialer(websocket.DefaultDialer) 114 } 115 116 // ConnectWithDialer creates a websocket connection with the given ConnectUrl using the dialer. 117 // This is racy and error-prone and should not be used. Use any of the New* functions to create a websocket. 118 func (wsc *WebSocketClient) ConnectWithDialer(dialer *websocket.Dialer) *AppError { 119 var err error 120 wsc.Conn, _, err = dialer.Dial(wsc.ConnectUrl, nil) 121 if err != nil { 122 return NewAppError("Connect", "model.websocket_client.connect_fail.app_error", nil, err.Error(), http.StatusInternalServerError) 123 } 124 // Super racy and should not be done anyways. 125 // All of this needs to be redesigned for v6. 126 wsc.configurePingHandling() 127 // If it has been closed before, we just restart the writer. 128 if atomic.CompareAndSwapInt32(&wsc.closed, 1, 0) { 129 wsc.writeChan = make(chan writeMessage) 130 wsc.quitWriterChan = make(chan struct{}) 131 go wsc.writer() 132 wsc.resetTimerChan = make(chan struct{}) 133 wsc.quitPingWatchdog = make(chan struct{}) 134 } 135 136 wsc.EventChannel = make(chan *WebSocketEvent, 100) 137 wsc.ResponseChannel = make(chan *WebSocketResponse, 100) 138 139 wsc.SendMessage(WEBSOCKET_AUTHENTICATION_CHALLENGE, map[string]interface{}{"token": wsc.AuthToken}) 140 141 return nil 142 } 143 144 // Close closes the websocket client. It is recommended that a closed client should not be 145 // reused again. Rather a new client should be created anew. 146 func (wsc *WebSocketClient) Close() { 147 // CAS to 1 and proceed. Return if already 1. 148 if !atomic.CompareAndSwapInt32(&wsc.closed, 0, 1) { 149 return 150 } 151 wsc.quitWriterChan <- struct{}{} 152 close(wsc.writeChan) 153 // We close the connection, which breaks the reader loop. 154 // Then we let the defer block in the reader do further cleanup. 155 wsc.Conn.Close() 156 } 157 158 // TODO: un-export the Conn so that Write methods go through the writer 159 func (wsc *WebSocketClient) writer() { 160 for { 161 select { 162 case msg := <-wsc.writeChan: 163 switch msg.msgType { 164 case msgTypeJSON: 165 wsc.Conn.WriteJSON(msg.data) 166 case msgTypePong: 167 wsc.Conn.WriteMessage(websocket.PongMessage, []byte{}) 168 } 169 case <-wsc.quitWriterChan: 170 return 171 } 172 } 173 } 174 175 // Listen starts the read loop of the websocket client. 176 func (wsc *WebSocketClient) Listen() { 177 // This loop can exit in 2 conditions: 178 // 1. Either the connection breaks naturally. 179 // 2. Close was explicitly called, which closes the connection manually. 180 // 181 // Due to the way the API is written, there is a requirement that a client may NOT 182 // call Listen at all and can still call Close and Connect. 183 // Therefore, we let the cleanup of the reader stuff rely on closing the connection 184 // and then we do the cleanup in the defer block. 185 // 186 // First, we close some channels and then CAS to 1 and proceed to close the writer chan also. 187 // This is needed because then the defer clause does not double-close the writer when (2) happens. 188 // But if (1) happens, we set the closed bit, and close the rest of the stuff. 189 go func() { 190 defer func() { 191 close(wsc.EventChannel) 192 close(wsc.ResponseChannel) 193 close(wsc.quitPingWatchdog) 194 close(wsc.resetTimerChan) 195 // We CAS to 1 and proceed. 196 if !atomic.CompareAndSwapInt32(&wsc.closed, 0, 1) { 197 return 198 } 199 wsc.quitWriterChan <- struct{}{} 200 close(wsc.writeChan) 201 wsc.Conn.Close() // This can most likely be removed. Needs to be checked. 202 }() 203 204 var buf bytes.Buffer 205 buf.Grow(avgReadMsgSizeBytes) 206 207 for { 208 // Reset buffer. 209 buf.Reset() 210 _, r, err := wsc.Conn.NextReader() 211 if err != nil { 212 if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseNoStatusReceived) { 213 wsc.ListenError = NewAppError("NewWebSocketClient", "model.websocket_client.connect_fail.app_error", nil, err.Error(), http.StatusInternalServerError) 214 } 215 return 216 } 217 // Use pre-allocated buffer. 218 _, err = buf.ReadFrom(r) 219 if err != nil { 220 // This should use a different error ID, but en.json is not imported anyways. 221 // It's a different bug altogether but we let it be for now. 222 // See MM-24520. 223 wsc.ListenError = NewAppError("NewWebSocketClient", "model.websocket_client.connect_fail.app_error", nil, err.Error(), http.StatusInternalServerError) 224 return 225 } 226 227 event := WebSocketEventFromJson(bytes.NewReader(buf.Bytes())) 228 if event == nil { 229 continue 230 } 231 if event.IsValid() { 232 wsc.EventChannel <- event 233 continue 234 } 235 236 var response WebSocketResponse 237 if err := json.Unmarshal(buf.Bytes(), &response); err == nil && response.IsValid() { 238 wsc.ResponseChannel <- &response 239 continue 240 } 241 } 242 }() 243 } 244 245 func (wsc *WebSocketClient) SendMessage(action string, data map[string]interface{}) { 246 req := &WebSocketRequest{} 247 req.Seq = wsc.Sequence 248 req.Action = action 249 req.Data = data 250 251 wsc.Sequence++ 252 wsc.writeChan <- writeMessage{ 253 msgType: msgTypeJSON, 254 data: req, 255 } 256 } 257 258 // UserTyping will push a user_typing event out to all connected users 259 // who are in the specified channel 260 func (wsc *WebSocketClient) UserTyping(channelId, parentId string) { 261 data := map[string]interface{}{ 262 "channel_id": channelId, 263 "parent_id": parentId, 264 } 265 266 wsc.SendMessage("user_typing", data) 267 } 268 269 // GetStatuses will return a map of string statuses using user id as the key 270 func (wsc *WebSocketClient) GetStatuses() { 271 wsc.SendMessage("get_statuses", nil) 272 } 273 274 // GetStatusesByIds will fetch certain user statuses based on ids and return 275 // a map of string statuses using user id as the key 276 func (wsc *WebSocketClient) GetStatusesByIds(userIds []string) { 277 data := map[string]interface{}{ 278 "user_ids": userIds, 279 } 280 wsc.SendMessage("get_statuses_by_ids", data) 281 } 282 283 func (wsc *WebSocketClient) configurePingHandling() { 284 wsc.Conn.SetPingHandler(wsc.pingHandler) 285 wsc.pingTimeoutTimer = time.NewTimer(time.Second * (60 + PING_TIMEOUT_BUFFER_SECONDS)) 286 go wsc.pingWatchdog() 287 } 288 289 func (wsc *WebSocketClient) pingHandler(appData string) error { 290 if atomic.LoadInt32(&wsc.closed) == 1 { 291 return nil 292 } 293 wsc.resetTimerChan <- struct{}{} 294 wsc.writeChan <- writeMessage{ 295 msgType: msgTypePong, 296 } 297 return nil 298 } 299 300 // pingWatchdog is used to send values to the PingTimeoutChannel whenever a timeout occurs. 301 // We use the resetTimerChan from the pingHandler to pass the signal, and then reset the timer 302 // after draining it. And if the timer naturally expires, we also extend it to prevent it from 303 // being deadlocked when the resetTimerChan case runs. Because timer.Stop would return false, 304 // and the code would be forever stuck trying to read from C. 305 func (wsc *WebSocketClient) pingWatchdog() { 306 for { 307 select { 308 case <-wsc.resetTimerChan: 309 if !wsc.pingTimeoutTimer.Stop() { 310 <-wsc.pingTimeoutTimer.C 311 } 312 wsc.pingTimeoutTimer.Reset(time.Second * (60 + PING_TIMEOUT_BUFFER_SECONDS)) 313 314 case <-wsc.pingTimeoutTimer.C: 315 wsc.PingTimeoutChannel <- true 316 wsc.pingTimeoutTimer.Reset(time.Second * (60 + PING_TIMEOUT_BUFFER_SECONDS)) 317 case <-wsc.quitPingWatchdog: 318 return 319 } 320 } 321 }