From 2b1b001bccb3d637cd51f69daa678eb0f8ade508 Mon Sep 17 00:00:00 2001 From: Agniva De Sarker Date: Mon, 30 Mar 2020 21:09:52 +0530 Subject: [PATCH] MM-23503: Fix race in websocket_client writer (#14160) * MM-23503: Fix race in websocket_client writer We create a separate writer goroutine where all the writes happen. * Fixing the case of double close * Incorporate review comments * Use CAS * Fix incorrect comment * Check if client is closed in pingHandler too Co-authored-by: mattermod --- model/websocket_client.go | 99 +++++++++++++++++++++++++++------- model/websocket_client_test.go | 58 ++++++++++++++++++++ 2 files changed, 137 insertions(+), 20 deletions(-) create mode 100644 model/websocket_client_test.go diff --git a/model/websocket_client.go b/model/websocket_client.go index 619b469a4c..f81affb984 100644 --- a/model/websocket_client.go +++ b/model/websocket_client.go @@ -7,6 +7,7 @@ import ( "bytes" "encoding/json" "net/http" + "sync/atomic" "time" "github.com/gorilla/websocket" @@ -17,6 +18,18 @@ const ( PING_TIMEOUT_BUFFER_SECONDS = 5 ) +type msgType int + +const ( + msgTypeJSON msgType = iota + 1 + msgTypePong +) + +type writeMessage struct { + msgType msgType + data interface{} +} + // WebSocketClient stores the necessary information required to // communicate with a WebSocket endpoint. type WebSocketClient struct { @@ -30,8 +43,13 @@ type WebSocketClient struct { EventChannel chan *WebSocketEvent // The channel used to receive various events pushed from the server. For example: typing, posted ResponseChannel chan *WebSocketResponse // The channel used to receive responses for requests made to the server ListenError *AppError // A field that is set if there was an abnormal closure of the WebSocket connection - pingTimeoutTimer *time.Timer - closeChannel chan struct{} + writeChan chan writeMessage + + pingTimeoutTimer *time.Timer + quitPingWatchdog chan struct{} + + quitWriterChan chan struct{} + closed int32 } // NewWebSocketClient constructs a new WebSocket client with convenience @@ -49,21 +67,22 @@ func NewWebSocketClientWithDialer(dialer *websocket.Dialer, url, authToken strin } client := &WebSocketClient{ - url, - url + API_URL_SUFFIX, - url + API_URL_SUFFIX + "/websocket", - conn, - authToken, - 1, - make(chan bool, 1), - make(chan *WebSocketEvent, 100), - make(chan *WebSocketResponse, 100), - nil, - nil, - make(chan struct{}), + Url: url, + ApiUrl: url + API_URL_SUFFIX, + ConnectUrl: url + API_URL_SUFFIX + "/websocket", + Conn: conn, + AuthToken: authToken, + Sequence: 1, + PingTimeoutChannel: make(chan bool, 1), + EventChannel: make(chan *WebSocketEvent, 100), + ResponseChannel: make(chan *WebSocketResponse, 100), + writeChan: make(chan writeMessage), + quitPingWatchdog: make(chan struct{}), + quitWriterChan: make(chan struct{}), } client.configurePingHandling() + go client.writer() client.SendMessage(WEBSOCKET_AUTHENTICATION_CHALLENGE, map[string]interface{}{"token": authToken}) @@ -82,18 +101,29 @@ func NewWebSocketClient4WithDialer(dialer *websocket.Dialer, url, authToken stri return NewWebSocketClientWithDialer(dialer, url, authToken) } +// Connect creates a websocket connection with the given ConnectUrl. +// This is racy and error-prone should not be used. Use any of the New* functions to create a websocket. func (wsc *WebSocketClient) Connect() *AppError { return wsc.ConnectWithDialer(websocket.DefaultDialer) } +// ConnectWithDialer creates a websocket connection with the given ConnectUrl using the dialer. +// This is racy and error-prone and should not be used. Use any of the New* functions to create a websocket. func (wsc *WebSocketClient) ConnectWithDialer(dialer *websocket.Dialer) *AppError { var err error wsc.Conn, _, err = dialer.Dial(wsc.ConnectUrl, nil) if err != nil { return NewAppError("Connect", "model.websocket_client.connect_fail.app_error", nil, err.Error(), http.StatusInternalServerError) } - + // Super racy and should not be done anyways. + // All of this needs to be redesigned for v6. wsc.configurePingHandling() + // If it has been closed before, we just restart the writer. + if atomic.CompareAndSwapInt32(&wsc.closed, 1, 0) { + wsc.writeChan = make(chan writeMessage) + wsc.quitWriterChan = make(chan struct{}) + go wsc.writer() + } wsc.EventChannel = make(chan *WebSocketEvent, 100) wsc.ResponseChannel = make(chan *WebSocketResponse, 100) @@ -104,16 +134,38 @@ func (wsc *WebSocketClient) ConnectWithDialer(dialer *websocket.Dialer) *AppErro } func (wsc *WebSocketClient) Close() { + // CAS to 1 and proceed. Return if already 1. + if !atomic.CompareAndSwapInt32(&wsc.closed, 0, 1) { + return + } + wsc.quitWriterChan <- struct{}{} + close(wsc.writeChan) wsc.Conn.Close() } +func (wsc *WebSocketClient) writer() { + for { + select { + case msg := <-wsc.writeChan: + switch msg.msgType { + case msgTypeJSON: + wsc.Conn.WriteJSON(msg.data) + case msgTypePong: + wsc.Conn.WriteMessage(websocket.PongMessage, []byte{}) + } + case <-wsc.quitWriterChan: + return + } + } +} + func (wsc *WebSocketClient) Listen() { go func() { defer func() { wsc.Conn.Close() close(wsc.EventChannel) close(wsc.ResponseChannel) - close(wsc.closeChannel) + close(wsc.quitPingWatchdog) }() for { @@ -153,8 +205,10 @@ func (wsc *WebSocketClient) SendMessage(action string, data map[string]interface req.Data = data wsc.Sequence++ - - wsc.Conn.WriteJSON(req) + wsc.writeChan <- writeMessage{ + msgType: msgTypeJSON, + data: req, + } } // UserTyping will push a user_typing event out to all connected users @@ -189,12 +243,17 @@ func (wsc *WebSocketClient) configurePingHandling() { } func (wsc *WebSocketClient) pingHandler(appData string) error { + if atomic.LoadInt32(&wsc.closed) == 1 { + return nil + } if !wsc.pingTimeoutTimer.Stop() { <-wsc.pingTimeoutTimer.C } wsc.pingTimeoutTimer.Reset(time.Second * (60 + PING_TIMEOUT_BUFFER_SECONDS)) - wsc.Conn.WriteMessage(websocket.PongMessage, []byte{}) + wsc.writeChan <- writeMessage{ + msgType: msgTypePong, + } return nil } @@ -202,6 +261,6 @@ func (wsc *WebSocketClient) pingWatchdog() { select { case <-wsc.pingTimeoutTimer.C: wsc.PingTimeoutChannel <- true - case <-wsc.closeChannel: + case <-wsc.quitPingWatchdog: } } diff --git a/model/websocket_client_test.go b/model/websocket_client_test.go new file mode 100644 index 0000000000..4bdc011731 --- /dev/null +++ b/model/websocket_client_test.go @@ -0,0 +1,58 @@ +// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved. +// See LICENSE.txt for license information. + +package model + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +func dummyWebsocketHandler(t *testing.T) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + upgrader := &websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, + } + conn, err := upgrader.Upgrade(w, req, nil) + var buf []byte + for { + _, buf, err = conn.ReadMessage() + if err != nil { + break + } + t.Logf("%s\n", buf) + err = conn.WriteMessage(websocket.PingMessage, []byte("ping")) + if err != nil { + break + } + } + if _, ok := err.(*websocket.CloseError); !ok { + require.NoError(t, err) + } + } +} + +// TestWebSocketRace needs to be run with -race to verify that +// there is no race. +func TestWebSocketRace(t *testing.T) { + s := httptest.NewServer(dummyWebsocketHandler(t)) + defer s.Close() + + url := strings.Replace(s.URL, "http://", "ws://", 1) + cli, err := NewWebSocketClient4(url, "authToken") + require.Nil(t, err) + + cli.Listen() + + for i := 0; i < 10; i++ { + time.Sleep(500 * time.Millisecond) + cli.UserTyping("channel", "parentId") + } +}