snapshot of https://github.com/mattermost/mattermost.git (master, 2026-09-16; blobs > 131072B omitted)

server/public/model/websocket_client_test.go

@@ -0,0 +1,221 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package model
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/vmihailenco/msgpack/v5"
+)
+
+func dummyWebsocketHandler(t *testing.T, wg *sync.WaitGroup) http.HandlerFunc {
+	return func(w http.ResponseWriter, req *http.Request) {
+		defer wg.Done()
+		upgrader := &websocket.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+		}
+		conn, err := upgrader.Upgrade(w, req, nil)
+		require.NoError(t, err)
+		var buf []byte
+		for {
+			err = conn.SetReadDeadline(time.Now().Add(time.Second))
+			if err != nil {
+				break
+			}
+			_, buf, err = conn.ReadMessage()
+			if err != nil {
+				break
+			}
+			t.Logf("%s\n", buf)
+			err = conn.WriteMessage(websocket.PingMessage, []byte("ping"))
+			if err != nil {
+				break
+			}
+		}
+	}
+}
+
+// TestWebSocketRace needs to be run with -race to verify that
+// there is no race.
+func TestWebSocketRace(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	s := httptest.NewServer(dummyWebsocketHandler(t, &wg))
+	defer func() {
+		s.Close()
+		wg.Wait()
+	}()
+
+	url := strings.Replace(s.URL, "http://", "ws://", 1)
+	cli, err := NewWebSocketClient4(url, "authToken")
+	require.NoError(t, err)
+
+	cli.Listen()
+
+	for range 10 {
+		time.Sleep(500 * time.Millisecond)
+		cli.UserTyping("channel", "parentId")
+	}
+}
+
+func TestWebSocketClose(t *testing.T) {
+	// This fails in SuddenClose because we check for closing the writeChan
+	// only after waiting the closure of Event and Response channels.
+	// Therefore, it is still racy and can panic. There is no use chasing this
+	// again because it will be completely overhauled in v6.
+	t.Skip("Skipping the test. Will be changed in v6.")
+	var wg sync.WaitGroup
+	wg.Add(1)
+	s := httptest.NewServer(dummyWebsocketHandler(t, &wg))
+	defer func() {
+		s.Close()
+		wg.Wait()
+	}()
+
+	url := strings.Replace(s.URL, "http://", "ws://", 1)
+
+	// Check whether the Event and Response channels
+	// have been closed or not.
+	waitClose := func(doneChan chan struct{}) int {
+		numClosed := 0
+		timeout := time.After(300 * time.Millisecond)
+		for {
+			select {
+			case <-doneChan:
+				numClosed++
+				if numClosed == 2 {
+					return numClosed
+				}
+			case <-timeout:
+				require.Fail(t, "timed out waiting for channels to be closed")
+				return numClosed
+			}
+		}
+	}
+
+	checkWriteChan := func(writeChan chan writeMessage) {
+		defer func() {
+			if x := recover(); x == nil {
+				require.Fail(t, "should have panicked due to closing a closed channel")
+			}
+		}()
+		close(writeChan)
+	}
+
+	waitForResponses := func(doneChan chan struct{}, cli *WebSocketClient) {
+		go func() {
+			for range cli.EventChannel { //nolint:revive
+			}
+			doneChan <- struct{}{}
+		}()
+		go func() {
+			for range cli.ResponseChannel { //nolint:revive
+			}
+			doneChan <- struct{}{}
+		}()
+	}
+
+	t.Run("SuddenClose", func(t *testing.T) {
+		cli, err := NewWebSocketClient4(url, "authToken")
+		require.NoError(t, err)
+
+		cli.Listen()
+
+		doneChan := make(chan struct{}, 2)
+		waitForResponses(doneChan, cli)
+
+		cli.UserTyping("channelId", "parentId")
+		cli.Conn.Close()
+
+		numClosed := waitClose(doneChan)
+		assert.Equal(t, 2, numClosed, "unexpected number of channels closed")
+
+		// Check whether the write channel was closed or not.
+		checkWriteChan(cli.writeChan)
+
+		require.NotNil(t, cli.ListenError, "non-nil listen error")
+		assert.Equal(t, "model.websocket_client.connect_fail.app_error", cli.ListenError.Id, "unexpected error id")
+	})
+
+	t.Run("ExplicitClose", func(t *testing.T) {
+		cli, err := NewWebSocketClient4(url, "authToken")
+		require.NoError(t, err)
+
+		cli.Listen()
+
+		doneChan := make(chan struct{}, 2)
+		waitForResponses(doneChan, cli)
+
+		cli.UserTyping("channelId", "parentId")
+		cli.Close()
+
+		numClosed := waitClose(doneChan)
+		assert.Equal(t, 2, numClosed, "unexpected number of channels closed")
+
+		// Check whether the write channel was closed or not.
+		checkWriteChan(cli.writeChan)
+	})
+}
+
+func binaryWebsocketHandler(t *testing.T, clientData map[string]any, doneCh chan struct{}) http.HandlerFunc {
+	return func(w http.ResponseWriter, req *http.Request) {
+		defer close(doneCh)
+		upgrader := &websocket.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+		}
+		conn, err := upgrader.Upgrade(w, req, nil)
+		require.NoError(t, err)
+		defer conn.Close()
+
+		for {
+			msgType, buf, err := conn.ReadMessage()
+			require.NoError(t, err)
+			if msgType == websocket.BinaryMessage {
+				require.Equal(t, msgType, websocket.BinaryMessage)
+				wsReq := &WebSocketRequest{}
+				err = msgpack.Unmarshal(buf, wsReq)
+				require.NoError(t, err)
+				require.Equal(t, clientData, wsReq.Data)
+				break
+			}
+		}
+	}
+}
+
+func TestWebSocketSendBinaryMessage(t *testing.T) {
+	clientData := map[string]any{
+		"data": []byte("some data to send as binary"),
+	}
+
+	doneCh := make(chan struct{})
+	s := httptest.NewServer(binaryWebsocketHandler(t, clientData, doneCh))
+	defer s.Close()
+
+	url := strings.Replace(s.URL, "http://", "ws://", 1)
+	cli, err := NewWebSocketClient4(url, "authToken")
+	require.NoError(t, err)
+	cli.Listen()
+	defer cli.Close()
+
+	err = cli.SendBinaryMessage("binaryAction", map[string]any{
+		"unmarshable": func() {},
+	})
+	require.Error(t, err)
+
+	err = cli.SendBinaryMessage("binaryAction", clientData)
+	require.NoError(t, err)
+
+	// This is to make sure the message is handled prior to exiting.
+	<-doneCh
+}