Allow per connection WebSocket broadcasts (#19993)

2022-04-22 15:53:42 +02:00
--- a/app/web_conn.go
+++ b/app/web_conn.go
@@ -737,6 +737,11 @@ func (wc *WebConn) shouldSendEvent(msg *model.WebSocketEvent) bool {
 		}
 	}

+	// If the event is destined to a specific connection
+	if msg.GetBroadcast().ConnectionId != "" {
+		return wc.GetConnectionID() == msg.GetBroadcast().ConnectionId
+	}
+
 	// If the event is destined to a specific user
 	if msg.GetBroadcast().UserId != "" {
 		return wc.UserId == msg.GetBroadcast().UserId
--- a/app/web_conn_test.go
+++ b/app/web_conn_test.go
@@ -36,6 +36,8 @@ func TestWebConnShouldSendEvent(t *testing.T) {
 		T:      i18n.T,
 	}

+	user1ConnID := model.NewId()
+	basicUserWc.SetConnectionID(user1ConnID)
 	basicUserWc.SetSession(session)
 	basicUserWc.SetSessionToken(session.Token)
 	basicUserWc.SetSessionExpiresAt(session.ExpiresAt)
@@ -55,6 +57,8 @@ func TestWebConnShouldSendEvent(t *testing.T) {
 		T:      i18n.T,
 	}

+	user2ConnID := model.NewId()
+	basicUser2Wc.SetConnectionID(user2ConnID)
 	basicUser2Wc.SetSession(session2)
 	basicUser2Wc.SetSessionToken(session2.Token)
 	basicUser2Wc.SetSessionExpiresAt(session2.ExpiresAt)
@@ -68,10 +72,33 @@ func TestWebConnShouldSendEvent(t *testing.T) {
 		T:      i18n.T,
 	}

+	adminConnID := model.NewId()
+	adminUserWc.SetConnectionID(adminConnID)
 	adminUserWc.SetSession(session3)
 	adminUserWc.SetSessionToken(session3.Token)
 	adminUserWc.SetSessionExpiresAt(session3.ExpiresAt)

+	session4, err := th.App.CreateSession(&model.Session{UserId: th.BasicUser.Id, Roles: th.BasicUser.GetRawRoles(), TeamMembers: []*model.TeamMember{
+		{
+			UserId: th.BasicUser.Id,
+			TeamId: th.BasicTeam.Id,
+			Roles:  model.TeamUserRoleId,
+		},
+	}})
+	require.Nil(t, err)
+
+	basicUserWc2 := &WebConn{
+		App:    th.App,
+		UserId: th.BasicUser.Id,
+		T:      i18n.T,
+	}
+
+	user1Conn2ID := model.NewId()
+	basicUserWc2.SetConnectionID(user1Conn2ID)
+	basicUserWc2.SetSession(session4)
+	basicUserWc2.SetSessionToken(session4.Token)
+	basicUserWc2.SetSessionExpiresAt(session4.ExpiresAt)
+
 	// By default, only BasicUser and BasicUser2 get added to the BasicTeam.
 	th.LinkUserToTeam(th.SystemAdminUser, th.BasicTeam)

@@ -80,18 +107,21 @@ func TestWebConnShouldSendEvent(t *testing.T) {
 	th.AddUserToChannel(th.SystemAdminUser, channel2)

 	cases := []struct {
-		Description   string
-		Broadcast     *model.WebsocketBroadcast
-		User1Expected bool
-		User2Expected bool
-		AdminExpected bool
+		Description        string
+		Broadcast          *model.WebsocketBroadcast
+		User1Expected      bool
+		User2Expected      bool
+		AdminExpected      bool
+		User1Conn2Expected bool
 	}{
-		{"should send to all", &model.WebsocketBroadcast{}, true, true, true},
-		{"should only send to basic user", &model.WebsocketBroadcast{UserId: th.BasicUser.Id}, true, false, false},
-		{"should omit basic user 2", &model.WebsocketBroadcast{OmitUsers: map[string]bool{th.BasicUser2.Id: true}}, true, false, true},
-		{"should only send to admin", &model.WebsocketBroadcast{ContainsSensitiveData: true}, false, false, true},
-		{"should only send to non-admins", &model.WebsocketBroadcast{ContainsSanitizedData: true}, true, true, false},
-		{"should send to nobody", &model.WebsocketBroadcast{ContainsSensitiveData: true, ContainsSanitizedData: true}, false, false, false},
+		{"should send to all", &model.WebsocketBroadcast{}, true, true, true, true},
+		{"should only send to basic user", &model.WebsocketBroadcast{UserId: th.BasicUser.Id}, true, false, false, true},
+		{"should only send to basic user conn 1", &model.WebsocketBroadcast{ConnectionId: user1ConnID}, true, false, false, false},
+		{"should only send to basic user conn 2", &model.WebsocketBroadcast{ConnectionId: user1Conn2ID}, false, false, false, true},
+		{"should omit basic user 2", &model.WebsocketBroadcast{OmitUsers: map[string]bool{th.BasicUser2.Id: true}}, true, false, true, true},
+		{"should only send to admin", &model.WebsocketBroadcast{ContainsSensitiveData: true}, false, false, true, false},
+		{"should only send to non-admins", &model.WebsocketBroadcast{ContainsSanitizedData: true}, true, true, false, true},
+		{"should send to nobody", &model.WebsocketBroadcast{ContainsSensitiveData: true, ContainsSanitizedData: true}, false, false, false, false},
 		// needs more cases to get full coverage
 	}

@@ -114,6 +144,11 @@ func TestWebConnShouldSendEvent(t *testing.T) {
 			} else {
 				assert.False(t, adminUserWc.shouldSendEvent(event), "did not expect admin")
 			}
+			if c.User1Conn2Expected {
+				assert.True(t, basicUserWc2.shouldSendEvent(event), "expected user 1 conn 2")
+			} else {
+				assert.False(t, basicUserWc2.shouldSendEvent(event), "did not expect user 1 conn 2")
+			}
 		})
 	}

--- a/app/web_hub.go
+++ b/app/web_hub.go
@@ -539,13 +539,20 @@ func (h *Hub) Start() {
 						}
 					}
 				}
-				if msg.GetBroadcast().UserId != "" {
+
+				if connID := msg.GetBroadcast().ConnectionId; connID != "" {
+					if webConn := connIndex.byConnectionId[connID]; webConn != nil {
+						broadcast(webConn)
+						continue
+					}
+				} else if msg.GetBroadcast().UserId != "" {
 					candidates := connIndex.ForUser(msg.GetBroadcast().UserId)
 					for _, webConn := range candidates {
 						broadcast(webConn)
 					}
 					continue
 				}
+
 				candidates := connIndex.All()
 				for webConn := range candidates {
 					broadcast(webConn)
@@ -596,7 +603,8 @@ type hubConnectionIndex struct {
 	byUserId map[string][]*WebConn
 	// byConnection serves the dual purpose of storing the index of the webconn
 	// in the value of byUserId map, and also to get all connections.
-	byConnection map[*WebConn]int
+	byConnection   map[*WebConn]int
+	byConnectionId map[string]*WebConn
 	// staleThreshold is the limit beyond which inactive connections
 	// will be deleted.
 	staleThreshold time.Duration
@@ -606,6 +614,7 @@ func newHubConnectionIndex(interval time.Duration) *hubConnectionIndex {
 	return &hubConnectionIndex{
 		byUserId:       make(map[string][]*WebConn),
 		byConnection:   make(map[*WebConn]int),
+		byConnectionId: make(map[string]*WebConn),
 		staleThreshold: interval,
 	}
 }
@@ -613,6 +622,7 @@ func newHubConnectionIndex(interval time.Duration) *hubConnectionIndex {
 func (i *hubConnectionIndex) Add(wc *WebConn) {
 	i.byUserId[wc.UserId] = append(i.byUserId[wc.UserId], wc)
 	i.byConnection[wc] = len(i.byUserId[wc.UserId]) - 1
+	i.byConnectionId[wc.GetConnectionID()] = wc
 }

 func (i *hubConnectionIndex) Remove(wc *WebConn) {
@@ -633,6 +643,7 @@ func (i *hubConnectionIndex) Remove(wc *WebConn) {
 	i.byConnection[last] = userConnIndex

 	delete(i.byConnection, wc)
+	delete(i.byConnectionId, wc.GetConnectionID())
 }

 func (i *hubConnectionIndex) Has(wc *WebConn) bool {
--- a/app/web_hub_test.go
+++ b/app/web_hub_test.go
@@ -228,20 +228,24 @@ func TestHubConnIndex(t *testing.T) {
 		App:    th.App,
 		UserId: model.NewId(),
 	}
+	wc1.SetConnectionID(model.NewId())

 	// User2
 	wc2 := &WebConn{
 		App:    th.App,
 		UserId: model.NewId(),
 	}
+	wc2.SetConnectionID(model.NewId())
 	wc3 := &WebConn{
 		App:    th.App,
 		UserId: wc2.UserId,
 	}
+	wc3.SetConnectionID(model.NewId())
 	wc4 := &WebConn{
 		App:    th.App,
 		UserId: wc2.UserId,
 	}
+	wc4.SetConnectionID(model.NewId())

 	connIndex.Add(wc1)
 	connIndex.Add(wc2)
@@ -283,7 +287,7 @@ func TestHubConnIndex(t *testing.T) {
 	t.Run("RemoveEndUser2", func(t *testing.T) {
 		connIndex.Remove(wc4) // Remove from end from user2

-		assert.ElementsMatch(t, connIndex.ForUser(wc2.UserId), []*WebConn{wc4})
+		assert.ElementsMatch(t, connIndex.ForUser(wc2.UserId), []*WebConn{wc2})
 		assert.ElementsMatch(t, connIndex.ForUser(wc1.UserId), []*WebConn{})
 		assert.True(t, connIndex.Has(wc2))
 		assert.False(t, connIndex.Has(wc3))
@@ -292,6 +296,62 @@ func TestHubConnIndex(t *testing.T) {
 	})
 }

+func TestHubConnIndexByConnectionId(t *testing.T) {
+	th := Setup(t)
+	defer th.TearDown()
+
+	connIndex := newHubConnectionIndex(1 * time.Second)
+
+	// User1
+	wc1ID := model.NewId()
+	wc1 := &WebConn{
+		App:    th.App,
+		UserId: model.NewId(),
+	}
+	wc1.SetConnectionID(wc1ID)
+
+	// User2
+	wc2ID := model.NewId()
+	wc2 := &WebConn{
+		App:    th.App,
+		UserId: model.NewId(),
+	}
+	wc2.SetConnectionID(wc2ID)
+
+	wc3ID := model.NewId()
+	wc3 := &WebConn{
+		App:    th.App,
+		UserId: wc2.UserId,
+	}
+	wc3.SetConnectionID(wc3ID)
+
+	t.Run("no connections", func(t *testing.T) {
+		assert.False(t, connIndex.Has(wc1))
+		assert.False(t, connIndex.Has(wc2))
+		assert.False(t, connIndex.Has(wc3))
+		assert.Empty(t, connIndex.byConnectionId)
+	})
+
+	t.Run("adding", func(t *testing.T) {
+		connIndex.Add(wc1)
+		connIndex.Add(wc3)
+
+		assert.Len(t, connIndex.byConnectionId, 2)
+		assert.Equal(t, wc1, connIndex.byConnectionId[wc1ID])
+		assert.Equal(t, wc3, connIndex.byConnectionId[wc3ID])
+		assert.Equal(t, (*WebConn)(nil), connIndex.byConnectionId[wc2ID])
+	})
+
+	t.Run("removing", func(t *testing.T) {
+		connIndex.Remove(wc3)
+
+		assert.Len(t, connIndex.byConnectionId, 1)
+		assert.Equal(t, wc1, connIndex.byConnectionId[wc1ID])
+		assert.Equal(t, (*WebConn)(nil), connIndex.byConnectionId[wc3ID])
+		assert.Equal(t, (*WebConn)(nil), connIndex.byConnectionId[wc2ID])
+	})
+}
+
 func TestHubConnIndexInactive(t *testing.T) {
 	connIndex := newHubConnectionIndex(2 * time.Second)

--- a/model/websocket_message.go
+++ b/model/websocket_message.go
@@ -84,10 +84,11 @@ type WebSocketMessage interface {
 }

 type WebsocketBroadcast struct {
-	OmitUsers             map[string]bool `json:"omit_users"` // broadcast is omitted for users listed here
-	UserId                string          `json:"user_id"`    // broadcast only occurs for this user
-	ChannelId             string          `json:"channel_id"` // broadcast only occurs for users in this channel
-	TeamId                string          `json:"team_id"`    // broadcast only occurs for users in this team
+	OmitUsers             map[string]bool `json:"omit_users"`    // broadcast is omitted for users listed here
+	UserId                string          `json:"user_id"`       // broadcast only occurs for this user
+	ChannelId             string          `json:"channel_id"`    // broadcast only occurs for users in this channel
+	TeamId                string          `json:"team_id"`       // broadcast only occurs for users in this team
+	ConnectionId          string          `json:"connection_id"` // broadcast only occurs for this connection
 	ContainsSanitizedData bool            `json:"-"`
 	ContainsSensitiveData bool            `json:"-"`
 	// ReliableClusterSend indicates whether or not the message should