From 7f4aeb7091105e813d6ffa27d7d9c6e481d93b3d Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Wed, 3 Aug 2022 18:12:42 +0200
Subject: [PATCH] Change query to include invited users

---
 syncapi/internal/keychange.go                 | 33 +++++++++----------
 syncapi/internal/keychange_test.go            |  1 +
 .../postgres/current_room_state_table.go      |  2 +-
 .../sqlite3/current_room_state_table.go       |  2 +-
 4 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/syncapi/internal/keychange.go b/syncapi/internal/keychange.go
index 7955b24f9..4bf54cae0 100644
--- a/syncapi/internal/keychange.go
+++ b/syncapi/internal/keychange.go
@@ -28,8 +28,6 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-const DeviceListLogName = "dl"
-
 // DeviceOTKCounts adds one-time key counts to the /sync response
 func DeviceOTKCounts(ctx context.Context, keyAPI keyapi.SyncKeyAPI, userID, deviceID string, res *types.Response) error {
 	var queryRes keyapi.QueryOneTimeKeysResponse
@@ -94,16 +92,13 @@ func DeviceListCatchup(
 	queryRes.UserIDs = append(queryRes.UserIDs, joinUserIDs...)
 	queryRes.UserIDs = append(queryRes.UserIDs, leaveUserIDs...)
 	queryRes.UserIDs = util.UniqueStrings(queryRes.UserIDs)
-	var sharedUsersMap map[string]int
-	sharedUsersMap, queryRes.UserIDs = filterSharedUsers(ctx, db, userID, queryRes.UserIDs)
+	sharedUsersMap := filterSharedUsers(ctx, db, userID, queryRes.UserIDs)
 	userSet := make(map[string]bool)
 	for _, userID := range res.DeviceLists.Changed {
-		if sharedUsersMap[userID] > 0 {
-			userSet[userID] = true
-		}
+		userSet[userID] = true
 	}
-	for _, userID := range queryRes.UserIDs {
-		if !userSet[userID] && sharedUsersMap[userID] > 0 {
+	for userID, count := range sharedUsersMap {
+		if !userSet[userID] && count > 0 {
 			res.DeviceLists.Changed = append(res.DeviceLists.Changed, userID)
 			hasNew = true
 			userSet[userID] = true
@@ -226,25 +221,27 @@ func TrackChangedUsers(
 // it down to include only users who the requesting user shares a room with.
 func filterSharedUsers(
 	ctx context.Context, db storage.SharedUsers, userID string, usersWithChangedKeys []string,
-) (map[string]int, []string) {
+) map[string]int {
 	sharedUsersMap := make(map[string]int, len(usersWithChangedKeys))
-	for _, userID := range usersWithChangedKeys {
-		sharedUsersMap[userID] = 0
+	for _, changedUserID := range usersWithChangedKeys {
+		sharedUsersMap[changedUserID] = 0
+		if changedUserID == userID {
+			// We forcibly put ourselves in this list because we should be notified about our own device updates
+			// and if we are in 0 rooms then we don't technically share any room with ourselves so we wouldn't
+			// be notified about key changes.
+			sharedUsersMap[userID] = 1
+		}
 	}
 	sharedUsers, err := db.SharedUsers(ctx, userID, usersWithChangedKeys)
 	if err != nil {
 		util.GetLogger(ctx).WithError(err).Errorf("db.SharedUsers failed: %s", err)
 		// default to all users so we do needless queries rather than miss some important device update
-		return nil, usersWithChangedKeys
+		return sharedUsersMap
 	}
 	for _, userID := range sharedUsers {
 		sharedUsersMap[userID]++
 	}
-	// We forcibly put ourselves in this list because we should be notified about our own device updates
-	// and if we are in 0 rooms then we don't technically share any room with ourselves so we wouldn't
-	// be notified about key changes.
-	sharedUsersMap[userID] = 1
-	return sharedUsersMap, sharedUsers
+	return sharedUsersMap
 }
 
 func joinedRooms(res *types.Response, userID string) []string {
diff --git a/syncapi/internal/keychange_test.go b/syncapi/internal/keychange_test.go
index 79ed440e7..6bfc91edd 100644
--- a/syncapi/internal/keychange_test.go
+++ b/syncapi/internal/keychange_test.go
@@ -129,6 +129,7 @@ type wantCatchup struct {
 }
 
 func assertCatchup(t *testing.T, hasNew bool, syncResponse *types.Response, want wantCatchup) {
+	t.Helper()
 	if hasNew != want.hasNew {
 		t.Errorf("got hasNew=%v want %v", hasNew, want.hasNew)
 	}
diff --git a/syncapi/storage/postgres/current_room_state_table.go b/syncapi/storage/postgres/current_room_state_table.go
index 678652031..58f404511 100644
--- a/syncapi/storage/postgres/current_room_state_table.go
+++ b/syncapi/storage/postgres/current_room_state_table.go
@@ -112,7 +112,7 @@ const selectEventsWithEventIDsSQL = "" +
 const selectSharedUsersSQL = "" +
 	"SELECT state_key FROM syncapi_current_room_state WHERE room_id = ANY(" +
 	"	SELECT room_id FROM syncapi_current_room_state WHERE state_key = $1 AND membership='join'" +
-	") AND state_key = ANY($2) AND membership='join';"
+	") AND state_key = ANY($2) AND membership IN ('join', 'invite');"
 
 type currentRoomStateStatements struct {
 	upsertRoomStateStmt                *sql.Stmt
diff --git a/syncapi/storage/sqlite3/current_room_state_table.go b/syncapi/storage/sqlite3/current_room_state_table.go
index 76dbd34f1..3a10b2325 100644
--- a/syncapi/storage/sqlite3/current_room_state_table.go
+++ b/syncapi/storage/sqlite3/current_room_state_table.go
@@ -96,7 +96,7 @@ const selectEventsWithEventIDsSQL = "" +
 const selectSharedUsersSQL = "" +
 	"SELECT state_key FROM syncapi_current_room_state WHERE room_id IN(" +
 	"	SELECT room_id FROM syncapi_current_room_state WHERE state_key = $1 AND membership='join'" +
-	") AND state_key IN ($2) AND membership='join';"
+	") AND state_key IN ($2) AND membership IN ('join', 'invite');"
 
 type currentRoomStateStatements struct {
 	db                                 *sql.DB