Fix adding excluded users to the private room sharing tables when joining a room (#11143)

* We only need to fetch users in private rooms

* Filter out `user_id` at the top

* Discard excluded users in the top loop

We weren't doing this in the "First, if they're our user" branch so this
is a bugfix.

* The caller must check that `user_id` is included

This is in the docstring. There are two call sites:
- one in `_handle_room_publicity_change`, which explicitly checks before calling;
- and another in `_handle_room_membership_event`, which returns early if
  the user is excluded.

So this change is safe.

* Test joining a private room with an excluded user

* Tweak an existing test

* Changelog

* test docstring

* lint

3 files changed, 67 insertions, 29 deletions

diff --git a/changelog.d/11143.misc b/changelog.d/11143.misc
new file mode 100644
index 0000000000..496e44a9c0
--- /dev/null
+++ b/changelog.d/11143.misc
@@ -0,0 +1 @@
+Fix a long-standing bug where users excluded from the directory could still be added to the `users_who_share_private_rooms` table after a regular user joins a private room.
\ No newline at end of file
diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py
index 991fee7e58..a0eb45446f 100644
--- a/synapse/handlers/user_directory.py
+++ b/synapse/handlers/user_directory.py
@@ -373,31 +373,29 @@ class UserDirectoryHandler(StateDeltasHandler):
         is_public = await self.store.is_room_world_readable_or_publicly_joinable(
             room_id
         )
-        other_users_in_room = await self.store.get_users_in_room(room_id)
-
         if is_public:
             await self.store.add_users_in_public_rooms(room_id, (user_id,))
         else:
+            users_in_room = await self.store.get_users_in_room(room_id)
+            other_users_in_room = [
+                other
+                for other in users_in_room
+                if other != user_id
+                and (
+                    not self.is_mine_id(other)
+                    or await self.store.should_include_local_user_in_dir(other)
+                )
+            ]
             to_insert = set()
 
             # First, if they're our user then we need to update for every user
             if self.is_mine_id(user_id):
-                if await self.store.should_include_local_user_in_dir(user_id):
-                    for other_user_id in other_users_in_room:
-                        if user_id == other_user_id:
-                            continue
-
-                        to_insert.add((user_id, other_user_id))
+                for other_user_id in other_users_in_room:
+                    to_insert.add((user_id, other_user_id))
 
             # Next we need to update for every local user in the room
             for other_user_id in other_users_in_room:
-                if user_id == other_user_id:
-                    continue
-
-                include_other_user = self.is_mine_id(
-                    other_user_id
-                ) and await self.store.should_include_local_user_in_dir(other_user_id)
-                if include_other_user:
+                if self.is_mine_id(other_user_id):
                     to_insert.add((other_user_id, user_id))
 
             if to_insert:
diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index b9ad92b977..70c621b825 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -646,22 +646,20 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         u2_token = self.login(u2, "pass")
         u3 = self.register_user("user3", "pass")
 
-        # We do not add users to the directory until they join a room.
+        # u1 can't see u2 until they share a private room, or u1 is in a public room.
         s = self.get_success(self.handler.search_users(u1, "user2", 10))
         self.assertEqual(len(s["results"]), 0)
 
+        # Get u1 and u2 into a private room.
         room = self.helper.create_room_as(u1, is_public=False, tok=u1_token)
         self.helper.invite(room, src=u1, targ=u2, tok=u1_token)
         self.helper.join(room, user=u2, tok=u2_token)
 
         # Check we have populated the database correctly.
-        shares_private = self.get_success(
-            self.user_dir_helper.get_users_who_share_private_rooms()
-        )
-        public_users = self.get_success(
-            self.user_dir_helper.get_users_in_public_rooms()
+        users, public_users, shares_private = self.get_success(
+            self.user_dir_helper.get_tables()
         )
-
+        self.assertEqual(users, {u1, u2, u3})
         self.assertEqual(shares_private, {(u1, u2, room), (u2, u1, room)})
         self.assertEqual(public_users, set())
 
@@ -680,14 +678,11 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         # User 2 then leaves.
         self.helper.leave(room, user=u2, tok=u2_token)
 
-        # Check we have removed the values.
-        shares_private = self.get_success(
-            self.user_dir_helper.get_users_who_share_private_rooms()
-        )
-        public_users = self.get_success(
-            self.user_dir_helper.get_users_in_public_rooms()
+        # Check this is reflected in the DB.
+        users, public_users, shares_private = self.get_success(
+            self.user_dir_helper.get_tables()
         )
-
+        self.assertEqual(users, {u1, u2, u3})
         self.assertEqual(shares_private, set())
         self.assertEqual(public_users, set())
 
@@ -698,6 +693,50 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
         s = self.get_success(self.handler.search_users(u1, "user3", 10))
         self.assertEqual(len(s["results"]), 0)
 
+    def test_joining_private_room_with_excluded_user(self) -> None:
+        """
+        When a user excluded from the user directory, E say, joins a private
+        room, E will not appear in the `users_who_share_private_rooms` table.
+
+        When a normal user, U say, joins a private room containing E, then
+        U will appear in the `users_who_share_private_rooms` table, but E will
+        not.
+        """
+        # Setup a support and two normal users.
+        alice = self.register_user("alice", "pass")
+        alice_token = self.login(alice, "pass")
+        bob = self.register_user("bob", "pass")
+        bob_token = self.login(bob, "pass")
+        support = "@support1:test"
+        self.get_success(
+            self.store.register_user(
+                user_id=support, password_hash=None, user_type=UserTypes.SUPPORT
+            )
+        )
+
+        # Alice makes a room. Inject the support user into the room.
+        room = self.helper.create_room_as(alice, is_public=False, tok=alice_token)
+        self.get_success(inject_member_event(self.hs, room, support, "join"))
+        # Check the DB state. The support user should not be in the directory.
+        users, in_public, in_private = self.get_success(
+            self.user_dir_helper.get_tables()
+        )
+        self.assertEqual(users, {alice, bob})
+        self.assertEqual(in_public, set())
+        self.assertEqual(in_private, set())
+
+        # Then invite Bob, who accepts.
+        self.helper.invite(room, alice, bob, tok=alice_token)
+        self.helper.join(room, bob, tok=bob_token)
+
+        # Check the DB state. The support user should not be in the directory.
+        users, in_public, in_private = self.get_success(
+            self.user_dir_helper.get_tables()
+        )
+        self.assertEqual(users, {alice, bob})
+        self.assertEqual(in_public, set())
+        self.assertEqual(in_private, {(alice, bob, room), (bob, alice, room)})
+
     def test_spam_checker(self) -> None:
         """
         A user which fails the spam checks will not appear in search results.