MM-54238 Add WebSocket broadcast hook and don't broadcast when other users were mentioned (#24641)

* MM-54238 Initial implementation

* MM-54238 Move websocket hook into app package

* MM-54238 Add tests for mentions in posted websocket messages

* Fix styling

* Fix other styling

* Idiomatic ID naming for new code

* Fix more styles

* Separate hooks to add mentions and followers

* Improved error handling for invalid types in hooks

* Rename HasChanges to ShouldProcess

* Pass broadcast hooks through hubStart

* Add test helper for asserting json unmarshaling

* Fix missing arguments in tests

* Ensure broadcast hooks are sent across the cluster and not to users

* Ensure tests actually cover following a post

* Fix code broken by merge

* Go vet again...

* Deep copy event before processing it with hooks

* Replace RemoveBroadcastHooks with WithoutBroadcastHooks

* Address feedback

* Add helper to fix type information for hook args

* Wrap WebSocketEvent and simplify BroadcastHook

* Address feedback

* Address feedback

server/channels/app/platform/helper_test.go

@@ -187,7 +187,7 @@ func setupTestHelper(dbStore store.Store, enterprise bool, includeCacheLayer boo
 		th.Service.SetLicense(nil)
 	}
 
-	err = th.Service.Start()
+	err = th.Service.Start(nil)
 	if err != nil {
 		panic(err)
 	}

server/channels/app/platform/service.go

@@ -344,8 +344,8 @@ func New(sc ServiceConfig, options ...Option) (*PlatformService, error) {
 	return ps, nil
 }
 
-func (ps *PlatformService) Start() error {
-	ps.hubStart()
+func (ps *PlatformService) Start(broadcastHooks map[string]BroadcastHook) error {
+	ps.hubStart(broadcastHooks)
 
 	ps.configListenerId = ps.AddConfigListener(func(_, _ *model.Config) {
 		ps.regenerateClientConfig()

server/channels/app/platform/web_broadcast_hook.go

@@ -0,0 +1,87 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package platform
+
+import (
+	"github.com/mattermost/mattermost/server/public/model"
+	"github.com/mattermost/mattermost/server/public/shared/mlog"
+)
+
+type BroadcastHook interface {
+	// Process takes a WebSocket event and modifies it in some way. It is passed a HookedWebSocketEvent which allows
+	// safe modification of the event.
+	Process(msg *HookedWebSocketEvent, webConn *WebConn, args map[string]any) error
+}
+
+func (h *Hub) runBroadcastHooks(msg *model.WebSocketEvent, webConn *WebConn, hookIDs []string, hookArgs []map[string]any) *model.WebSocketEvent {
+	if len(hookIDs) == 0 {
+		return msg
+	}
+
+	hookedEvent := MakeHookedWebSocketEvent(msg)
+
+	for i, hookID := range hookIDs {
+		hook := h.broadcastHooks[hookID]
+		args := hookArgs[i]
+		if hook == nil {
+			mlog.Warn("runBroadcastHooks: Unable to find broadcast hook", mlog.String("hook_id", hookID))
+			continue
+		}
+
+		hook.Process(hookedEvent, webConn, args)
+	}
+
+	return hookedEvent.Event()
+}
+
+// HookedWebSocketEvent is a wrapper for model.WebSocketEvent that is intended to provide a similar interface, except
+// it ensures the original WebSocket event is not modified.
+type HookedWebSocketEvent struct {
+	original *model.WebSocketEvent
+	copy     *model.WebSocketEvent
+}
+
+func MakeHookedWebSocketEvent(event *model.WebSocketEvent) *HookedWebSocketEvent {
+	return &HookedWebSocketEvent{
+		original: event,
+	}
+}
+
+func (he *HookedWebSocketEvent) Add(key string, value any) {
+	he.copyIfNecessary()
+
+	he.copy.Add(key, value)
+}
+
+func (he *HookedWebSocketEvent) EventType() string {
+	if he.copy == nil {
+		return he.original.EventType()
+	}
+
+	return he.copy.EventType()
+}
+
+// Get returns a value from the WebSocket event data. You should never mutate a value returned by this method.
+func (he *HookedWebSocketEvent) Get(key string) any {
+	if he.copy == nil {
+		return he.original.GetData()[key]
+	}
+
+	return he.copy.GetData()[key]
+}
+
+// copyIfNecessary should be called by any mutative method to ensure that the copy is instantiated.
+func (he *HookedWebSocketEvent) copyIfNecessary() {
+	if he.copy == nil {
+		he.copy = he.original.RemovePrecomputedJSON()
+	}
+}
+
+func (he *HookedWebSocketEvent) Event() *model.WebSocketEvent {
+	if he.copy == nil {
+		return he.original
+	}
+
+	return he.copy
+}

server/channels/app/platform/web_broadcast_hook_test.go

@@ -0,0 +1,206 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package platform
+
+import (
+	"testing"
+
+	"github.com/mattermost/mattermost/server/public/model"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const broadcastTest = "test_broadcast_hook"
+
+type testBroadcastHook struct{}
+
+func (h *testBroadcastHook) Process(msg *HookedWebSocketEvent, webConn *WebConn, args map[string]any) error {
+	if args["makes_changes"].(bool) {
+		changesMade, _ := msg.Get("changes_made").(int)
+		msg.Add("changes_made", changesMade+1)
+	}
+
+	return nil
+}
+
+func TestRunBroadcastHooks(t *testing.T) {
+	hub := &Hub{
+		broadcastHooks: map[string]BroadcastHook{
+			broadcastTest: &testBroadcastHook{},
+		},
+	}
+	webConn := &WebConn{}
+
+	t.Run("should not allocate a new object when no hooks are passed", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+
+		result := hub.runBroadcastHooks(event, webConn, nil, nil)
+
+		assert.Same(t, event, result)
+	})
+
+	t.Run("should not allocate a new object when a hook is not making changes", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+
+		hookIDs := []string{
+			broadcastTest,
+		}
+		hookArgs := []map[string]any{
+			{
+				"makes_changes": false,
+			},
+		}
+
+		result := hub.runBroadcastHooks(event, webConn, hookIDs, hookArgs)
+
+		assert.Same(t, event, result)
+	})
+
+	t.Run("should allocate a new object and remove when a hook makes changes", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+
+		hookIDs := []string{
+			broadcastTest,
+		}
+		hookArgs := []map[string]any{
+			{
+				"makes_changes": true,
+			},
+		}
+
+		result := hub.runBroadcastHooks(event, webConn, hookIDs, hookArgs)
+
+		assert.NotSame(t, event, result)
+		assert.NotSame(t, event.GetData(), result.GetData())
+		assert.Equal(t, map[string]any{}, event.GetData())
+		assert.Equal(t, result.GetData(), map[string]any{
+			"changes_made": 1,
+		})
+	})
+
+	t.Run("should not allocate a new object when multiple hooks are not making changes", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+
+		hookIDs := []string{
+			broadcastTest,
+			broadcastTest,
+			broadcastTest,
+		}
+		hookArgs := []map[string]any{
+			{
+				"makes_changes": false,
+			},
+			{
+				"makes_changes": false,
+			},
+			{
+				"makes_changes": false,
+			},
+		}
+
+		result := hub.runBroadcastHooks(event, webConn, hookIDs, hookArgs)
+
+		assert.Same(t, event, result)
+	})
+
+	t.Run("should be able to make changes from only one of make hooks", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+
+		var hookIDs []string
+		var hookArgs []map[string]any
+		for i := 0; i < 10; i++ {
+			hookIDs = append(hookIDs, broadcastTest)
+			hookArgs = append(hookArgs, map[string]any{
+				"makes_changes": i == 6,
+			})
+		}
+
+		result := hub.runBroadcastHooks(event, webConn, hookIDs, hookArgs)
+
+		assert.NotSame(t, event, result)
+		assert.NotSame(t, event.GetData(), result.GetData())
+		assert.Equal(t, event.GetData(), map[string]any{})
+		assert.Equal(t, result.GetData(), map[string]any{
+			"changes_made": 1,
+		})
+	})
+
+	t.Run("should be able to make changes from multiple hooks", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+
+		var hookIDs []string
+		var hookArgs []map[string]any
+		for i := 0; i < 10; i++ {
+			hookIDs = append(hookIDs, broadcastTest)
+			hookArgs = append(hookArgs, map[string]any{
+				"makes_changes": true,
+			})
+		}
+
+		result := hub.runBroadcastHooks(event, webConn, hookIDs, hookArgs)
+
+		assert.NotSame(t, event, result)
+		assert.NotSame(t, event.GetData(), result.GetData())
+		assert.Equal(t, event.GetData(), map[string]any{})
+		assert.Equal(t, result.GetData(), map[string]any{
+			"changes_made": 10,
+		})
+	})
+
+	t.Run("should not remove precomputed JSON when a hook doesn't make changes", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+		event = event.PrecomputeJSON()
+
+		// Ensure that the event has precomputed JSON because changes aren't included when ToJSON is called again
+		originalJSON, _ := event.ToJSON()
+		event.Add("data", 1234)
+		eventJSON, _ := event.ToJSON()
+		require.Equal(t, string(originalJSON), string(eventJSON))
+
+		hookIDs := []string{
+			broadcastTest,
+		}
+		hookArgs := []map[string]any{
+			{
+				"makes_changes": false,
+			},
+		}
+
+		result := hub.runBroadcastHooks(event, webConn, hookIDs, hookArgs)
+
+		eventJSON, _ = event.ToJSON()
+		assert.Equal(t, string(originalJSON), string(eventJSON))
+
+		resultJSON, _ := result.ToJSON()
+		assert.Equal(t, originalJSON, resultJSON)
+	})
+
+	t.Run("should remove precomputed JSON when a hook makes changes", func(t *testing.T) {
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+		event = event.PrecomputeJSON()
+
+		// Ensure that the event has precomputed JSON because changes aren't included when ToJSON is called again
+		originalJSON, _ := event.ToJSON()
+		event.Add("data", 1234)
+		eventJSON, _ := event.ToJSON()
+		require.Equal(t, originalJSON, eventJSON)
+
+		hookIDs := []string{
+			broadcastTest,
+		}
+		hookArgs := []map[string]any{
+			{
+				"makes_changes": true,
+			},
+		}
+
+		result := hub.runBroadcastHooks(event, webConn, hookIDs, hookArgs)
+
+		eventJSON, _ = event.ToJSON()
+		assert.Equal(t, string(originalJSON), string(eventJSON))
+
+		resultJSON, _ := result.ToJSON()
+		assert.NotEqual(t, originalJSON, resultJSON)
+	})
+}

server/channels/app/platform/web_hub.go

@@ -70,6 +70,7 @@ type Hub struct {
 	explicitStop    bool
 	checkRegistered chan *webConnSessionMessage
 	checkConn       chan *webConnCheckMessage
+	broadcastHooks  map[string]BroadcastHook
 }
 
 // newWebHub creates a new Hub.
@@ -90,7 +91,7 @@ func newWebHub(ps *PlatformService) *Hub {
 }
 
 // hubStart starts all the hubs.
-func (ps *PlatformService) hubStart() {
+func (ps *PlatformService) hubStart(broadcastHooks map[string]BroadcastHook) {
 	// Total number of hubs is twice the number of CPUs.
 	numberOfHubs := runtime.NumCPU() * 2
 	ps.logger.Info("Starting websocket hubs", mlog.Int("number_of_hubs", numberOfHubs))
@@ -100,6 +101,7 @@ func (ps *PlatformService) hubStart() {
 	for i := 0; i < numberOfHubs; i++ {
 		hubs[i] = newWebHub(ps)
 		hubs[i].connectionIndex = i
+		hubs[i].broadcastHooks = broadcastHooks
 		hubs[i].Start()
 	}
 	// Assigning to the hubs slice without any mutex is fine because it is only assigned once
@@ -492,14 +494,19 @@ func (h *Hub) Start() {
 				if metrics := h.platform.metricsIFace; metrics != nil {
 					metrics.DecrementWebSocketBroadcastBufferSize(strconv.Itoa(h.connectionIndex), 1)
 				}
+
+				// Remove the broadcast hook information before precomputing the JSON so that those aren't included in it
+				msg, broadcastHooks, broadcastHookArgs := msg.WithoutBroadcastHooks()
+
 				msg = msg.PrecomputeJSON()
+
 				broadcast := func(webConn *WebConn) {
 					if !connIndex.Has(webConn) {
 						return
 					}
 					if webConn.ShouldSendEvent(msg) {
 						select {
-						case webConn.send <- msg:
+						case webConn.send <- h.runBroadcastHooks(msg, webConn, broadcastHooks, broadcastHookArgs):
 						default:
 							// Don't log the warning if it's an inactive connection.
 							if webConn.active.Load() {

server/channels/app/platform/web_hub_test.go

@@ -4,6 +4,7 @@
 package platform
 
 import (
+	"bytes"
 	"encoding/json"
 	"net"
 	"net/http"
@@ -69,7 +70,7 @@ func TestHubStopWithMultipleConnections(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	th.Service.Start()
+	th.Service.Start(nil)
 	wc1 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
 	wc2 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
 	wc3 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
@@ -93,7 +94,7 @@ func TestHubStopRaceCondition(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	th.Service.Start()
+	th.Service.Start(nil)
 	wc1 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
 	defer wc1.Close()
 
@@ -476,7 +477,7 @@ func TestHubIsRegistered(t *testing.T) {
 	s := httptest.NewServer(dummyWebsocketHandler(t))
 	defer s.Close()
 
-	th.Service.Start()
+	th.Service.Start(nil)
 	wc1 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
 	wc2 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
 	wc3 := registerDummyWebConn(t, th, s.Listener.Addr(), session)
@@ -583,7 +584,7 @@ func BenchmarkGetHubForUserId(b *testing.B) {
 	th := Setup(b).InitBasic()
 	defer th.TearDown()
 
-	th.Service.Start()
+	th.Service.Start(nil)
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
@@ -618,3 +619,54 @@ func TestClusterBroadcast(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, clusterEvent.Broadcast, broadcast)
 }
+
+func TestClusterBroadcastHooks(t *testing.T) {
+	t.Run("should send broadcast hook information across cluster", func(t *testing.T) {
+		testCluster := &testlib.FakeClusterInterface{}
+
+		th := SetupWithCluster(t, testCluster)
+		defer th.TearDown()
+
+		hookID := broadcastTest
+		hookArgs := map[string]any{
+			"makes_changes": true,
+		}
+
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+		event.GetBroadcast().AddHook(hookID, hookArgs)
+
+		th.Service.Publish(event)
+
+		received, err := model.WebSocketEventFromJSON(bytes.NewReader(testCluster.GetMessages()[0].Data))
+
+		require.NoError(t, err)
+		assert.Equal(t, []string{hookID}, received.GetBroadcast().BroadcastHooks)
+		assert.Equal(t, []map[string]any{hookArgs}, received.GetBroadcast().BroadcastHookArgs)
+	})
+
+	t.Run("should not preserve type information for args", func(t *testing.T) {
+		// This behaviour isn't ideal, but this test confirms that it hasn't changed
+		testCluster := &testlib.FakeClusterInterface{}
+
+		th := SetupWithCluster(t, testCluster)
+		defer th.TearDown()
+
+		hookID := "test_broadcast_hook_with_args"
+		hookArgs := map[string]any{
+			"user":  &model.User{Id: "user1"},
+			"array": []string{"a", "b", "c"},
+		}
+
+		event := model.NewWebSocketEvent(model.WebsocketEventPosted, "", "", "", nil, "")
+		event.GetBroadcast().AddHook(hookID, hookArgs)
+
+		th.Service.Publish(event)
+
+		received, err := model.WebSocketEventFromJSON(bytes.NewReader(testCluster.GetMessages()[0].Data))
+
+		require.NoError(t, err)
+		assert.Equal(t, []string{hookID}, received.GetBroadcast().BroadcastHooks)
+		assert.IsType(t, map[string]any{}, received.GetBroadcast().BroadcastHookArgs[0]["user"])
+		assert.IsType(t, []any{}, received.GetBroadcast().BroadcastHookArgs[0]["array"])
+	})
+}