Implement deep copy method for model.WebSocketEvent (#18977)

app/post.go

@@ -735,7 +735,9 @@ func (a *App) publishWebsocketEventForPermalinkPost(post *model.Post, message *m
 			}
 			return false, err
 		}
-		messageCopy := message.Copy()
+		// Using DeepCopy here to avoid a race condition
+		// between publishing the event and setting the "post" data value below.
+		messageCopy := message.DeepCopy()
 		broadcastCopy := messageCopy.GetBroadcast()
 		broadcastCopy.UserId = cm.UserId
 		messageCopy.SetBroadcast(broadcastCopy)

app/post_test.go

@@ -826,6 +826,56 @@ func TestCreatePost(t *testing.T) {
 
 		require.EqualValues(t, int64(1), val)
 	})
+
+	t.Run("MM-40016 should not panic with `concurrent map read and map write`", func(t *testing.T) {
+		th := Setup(t).InitBasic()
+		defer th.TearDown()
+
+		channelForPreview := th.CreateChannel(th.BasicTeam)
+
+		for i := 0; i < 20; i++ {
+			user := th.CreateUser()
+			th.LinkUserToTeam(user, th.BasicTeam)
+			th.AddUserToChannel(user, channelForPreview)
+		}
+
+		referencedPost := &model.Post{
+			ChannelId: th.BasicChannel.Id,
+			Message:   "hello world",
+			UserId:    th.BasicUser.Id,
+		}
+		referencedPost, err := th.App.CreatePost(th.Context, referencedPost, th.BasicChannel, false, false)
+		require.Nil(t, err)
+
+		th.App.UpdateConfig(func(cfg *model.Config) {
+			*cfg.ServiceSettings.SiteURL = "http://example.com"
+			*cfg.ServiceSettings.EnablePermalinkPreviews = true
+		})
+
+		permalink := fmt.Sprintf("%s/%s/pl/%s", *th.App.Config().ServiceSettings.SiteURL, th.BasicTeam.Name, referencedPost.Id)
+
+		previewPost := &model.Post{
+			ChannelId: channelForPreview.Id,
+			Message:   permalink,
+			UserId:    th.BasicUser.Id,
+		}
+
+		previewPost, err = th.App.CreatePost(th.Context, previewPost, channelForPreview, false, false)
+		require.Nil(t, err)
+
+		n := 1000
+		var wg sync.WaitGroup
+		wg.Add(n)
+		for i := 0; i < n; i++ {
+			go func() {
+				defer wg.Done()
+				post := previewPost.Clone()
+				th.App.UpdatePost(th.Context, post, false)
+			}()
+		}
+
+		wg.Wait()
+	})
 }
 
 func TestPatchPost(t *testing.T) {
@@ -997,7 +1047,6 @@ func TestCreatePostAsUser(t *testing.T) {
 	t.Run("logs warning for user not in channel", func(t *testing.T) {
 		th := Setup(t).InitBasic()
 		defer th.TearDown()
-
 		user := th.CreateUser()
 		th.LinkUserToTeam(user, th.BasicTeam)
 

model/websocket_message.go

@@ -90,12 +90,58 @@ type WebsocketBroadcast struct {
 	ContainsSensitiveData bool            `json:"-"`
 }
 
+func (wb *WebsocketBroadcast) copy() *WebsocketBroadcast {
+	if wb == nil {
+		return nil
+	}
+
+	var c WebsocketBroadcast
+	if wb.OmitUsers != nil {
+		c.OmitUsers = make(map[string]bool, len(wb.OmitUsers))
+		for k, v := range wb.OmitUsers {
+			c.OmitUsers[k] = v
+		}
+	}
+	c.UserId = wb.UserId
+	c.ChannelId = wb.ChannelId
+	c.TeamId = wb.TeamId
+	c.ContainsSanitizedData = wb.ContainsSanitizedData
+	c.ContainsSensitiveData = wb.ContainsSensitiveData
+
+	return &c
+}
+
 type precomputedWebSocketEventJSON struct {
 	Event     json.RawMessage
 	Data      json.RawMessage
 	Broadcast json.RawMessage
 }
 
+func (p *precomputedWebSocketEventJSON) copy() *precomputedWebSocketEventJSON {
+	if p == nil {
+		return nil
+	}
+
+	var c precomputedWebSocketEventJSON
+
+	if p.Event != nil {
+		c.Event = make([]byte, len(p.Event))
+		copy(c.Event, p.Event)
+	}
+
+	if p.Data != nil {
+		c.Data = make([]byte, len(p.Data))
+		copy(c.Data, p.Data)
+	}
+
+	if p.Broadcast != nil {
+		c.Broadcast = make([]byte, len(p.Broadcast))
+		copy(c.Broadcast, p.Broadcast)
+	}
+
+	return &c
+}
+
 // webSocketEventJSON mirrors WebSocketEvent to make some of its unexported fields serializable
 type webSocketEventJSON struct {
 	Event     string                 `json:"event"`
@@ -154,6 +200,25 @@ func (ev *WebSocketEvent) Copy() *WebSocketEvent {
 	return copy
 }
 
+func (ev *WebSocketEvent) DeepCopy() *WebSocketEvent {
+	var dataCopy map[string]interface{}
+	if ev.data != nil {
+		dataCopy = make(map[string]interface{}, len(ev.data))
+		for k, v := range ev.data {
+			dataCopy[k] = v
+		}
+	}
+
+	copy := &WebSocketEvent{
+		event:           ev.event,
+		data:            dataCopy,
+		broadcast:       ev.broadcast.copy(),
+		sequence:        ev.sequence,
+		precomputedJSON: ev.precomputedJSON.copy(),
+	}
+	return copy
+}
+
 func (ev *WebSocketEvent) GetData() map[string]interface{} {
 	return ev.data
 }

model/websocket_message_test.go

@@ -157,3 +157,80 @@ func BenchmarkWebSocketEvent_ToJSON(b *testing.B) {
 		}
 	})
 }
+
+func TestWebsocketBroadcastCopy(t *testing.T) {
+	w := &WebsocketBroadcast{}
+	require.Equal(t, w, w.copy())
+
+	w = nil
+	require.Equal(t, w, w.copy())
+
+	w = &WebsocketBroadcast{
+		OmitUsers: map[string]bool{
+			"aaa": true,
+			"bbb": true,
+			"ccc": false,
+		},
+		UserId:                "aaa",
+		ChannelId:             "bbb",
+		TeamId:                "ccc",
+		ContainsSanitizedData: true,
+		ContainsSensitiveData: true,
+	}
+	require.Equal(t, w, w.copy())
+}
+
+func TestPrecomputedWebSocketEventJSONCopy(t *testing.T) {
+	p := &precomputedWebSocketEventJSON{}
+	require.Equal(t, p, p.copy())
+
+	p = nil
+	require.Equal(t, p, p.copy())
+
+	p = &precomputedWebSocketEventJSON{
+		Event:     []byte{},
+		Data:      []byte{},
+		Broadcast: []byte{},
+	}
+	require.Equal(t, p, p.copy())
+
+	p = &precomputedWebSocketEventJSON{
+		Event:     []byte{'a', 'b', 'c'},
+		Data:      []byte{'d', 'e', 'f'},
+		Broadcast: []byte{'g', 'h', 'i'},
+	}
+	require.Equal(t, p, p.copy())
+}
+
+func TestWebSocketEventDeepCopy(t *testing.T) {
+	omitUsers := map[string]bool{
+		"user1": true,
+		"user2": false,
+	}
+
+	broadcast := &WebsocketBroadcast{
+		OmitUsers:             omitUsers,
+		UserId:                "aaa",
+		ChannelId:             "bbb",
+		TeamId:                "ccc",
+		ContainsSanitizedData: true,
+		ContainsSensitiveData: true,
+	}
+
+	ev := NewWebSocketEvent("test", "team", "channel", "user", omitUsers)
+
+	ev.Add("post", &Post{})
+	ev.SetBroadcast(broadcast)
+	ev = ev.PrecomputeJSON()
+
+	evCopy := ev.DeepCopy()
+	require.Equal(t, ev, evCopy)
+	require.NotSame(t, ev.data, evCopy.data)
+	require.NotSame(t, ev.broadcast, evCopy.broadcast)
+	require.NotSame(t, ev.precomputedJSON, evCopy.precomputedJSON)
+
+	ev.Add("post", &Post{
+		Id: "test",
+	})
+	require.NotEqual(t, ev.data, evCopy.data)
+}