diff --git a/app/web_conn.go b/app/web_conn.go index 29f9fe558f..549d785b55 100644 --- 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 diff --git a/app/web_conn_test.go b/app/web_conn_test.go index ec313d4d87..be9dbdc56f 100644 --- 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") + } }) } diff --git a/app/web_hub.go b/app/web_hub.go index bbaf4fc5e6..0537e7af49 100644 --- 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 { diff --git a/app/web_hub_test.go b/app/web_hub_test.go index 7031434c85..10dd13c695 100644 --- 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) diff --git a/model/websocket_message.go b/model/websocket_message.go index 38e42bb2e9..343d73050c 100644 --- 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