From b9638c96ef9e0d2fb345f253521b6c52cb3ae612 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Thu, 30 Sep 2021 14:14:38 +0100
Subject: [PATCH 01/12] Introduce `should_include_local_users_in_dir`

We exclude three kinds of local users from the user_directory tables. At
present we don't consistently exclude all three in the same places. This
commit introduces a new function to gather those exclusion conditions
together. Because we have to handle local and remote users in different
ways, I've made that function only consider the case of remote users.
It's the callers responsibility to make the local versus remote
distinction clear and correct.

A test fixup is required. The test now hits a path which makes db
queries against the users table. The expected rows were missing, because
we were using a dummy user that hadn't actually been registered.

Broken-off from #10914.

----

By my reading this makes these changes:

* When an app service user registers or changes their profile, they will
  _not_ be added to the user directory. (Previously only support and
  deactivated users were excluded). This is consistent with the logic that
  rebuilds the user directory. See also [the discussion
  here](https://github.com/matrix-org/synapse/pull/10914#discussion_r716859548).
* When rebuilding the directory, exclude support and disabled users from
  room sharing tables. Previously only appservice users were excluded.
* Exclude all three categories of local users when rebuilding the
  directory. Previously `_populate_user_directory_process_users` didn't do
  any exclusion.
---
 synapse/handlers/user_directory.py            | 27 +++++---------
 .../storage/databases/main/user_directory.py  | 37 ++++++++++++++++---
 tests/handlers/test_user_directory.py         | 26 ++++++++-----
 3 files changed, 56 insertions(+), 34 deletions(-)

diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py
index f4430ce3c9aa..18d8c8744e75 100644
--- a/synapse/handlers/user_directory.py
+++ b/synapse/handlers/user_directory.py
@@ -132,12 +132,7 @@ async def handle_local_profile_change(
         # FIXME(#3714): We should probably do this in the same worker as all
         # the other changes.
 
-        # Support users are for diagnostics and should not appear in the user directory.
-        is_support = await self.store.is_support_user(user_id)
-        # When change profile information of deactivated user it should not appear in the user directory.
-        is_deactivated = await self.store.get_user_deactivated_status(user_id)
-
-        if not (is_support or is_deactivated):
+        if await self.store.should_include_local_user_in_dir(user_id):
             await self.store.update_profile_in_user_dir(
                 user_id, profile.display_name, profile.avatar_url
             )
@@ -229,8 +224,10 @@ async def _handle_deltas(self, deltas: List[Dict[str, Any]]) -> None:
                     else:
                         logger.debug("Server is still in room: %r", room_id)
 
-                is_support = await self.store.is_support_user(state_key)
-                if not is_support:
+                include_in_dir = not self.is_mine_id(
+                    state_key
+                ) or await self.store.should_include_local_user_in_dir(state_key)
+                if include_in_dir:
                     if change is MatchChange.no_change:
                         # Handle any profile changes
                         await self._handle_profile_change(
@@ -356,13 +353,7 @@ async def _handle_new_user(
 
             # First, if they're our user then we need to update for every user
             if self.is_mine_id(user_id):
-
-                is_appservice = self.store.get_if_app_services_interested_in_user(
-                    user_id
-                )
-
-                # We don't care about appservice users.
-                if not is_appservice:
+                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
@@ -374,10 +365,10 @@ async def _handle_new_user(
                 if user_id == other_user_id:
                     continue
 
-                is_appservice = self.store.get_if_app_services_interested_in_user(
+                include_other_user = self.is_mine_id(
                     other_user_id
-                )
-                if self.is_mine_id(other_user_id) and not is_appservice:
+                ) and await self.store.should_include_local_user_in_dir(other_user_id)
+                if include_other_user:
                     to_insert.add((other_user_id, user_id))
 
             if to_insert:
diff --git a/synapse/storage/databases/main/user_directory.py b/synapse/storage/databases/main/user_directory.py
index c26e3e066f9d..f22423afca45 100644
--- a/synapse/storage/databases/main/user_directory.py
+++ b/synapse/storage/databases/main/user_directory.py
@@ -238,6 +238,10 @@ def _get_next_batch(
 
                 # Update each user in the user directory.
                 for user_id, profile in users_with_profile.items():
+                    if self.hs.is_mine_id(
+                        user_id
+                    ) and not await self.should_include_local_user_in_dir(user_id):
+                        continue
                     await self.update_profile_in_user_dir(
                         user_id, profile.display_name, profile.avatar_url
                     )
@@ -246,7 +250,9 @@ def _get_next_batch(
 
                 if is_public:
                     for user_id in users_with_profile:
-                        if self.get_if_app_services_interested_in_user(user_id):
+                        if self.hs.is_mine_id(
+                            user_id
+                        ) and not await self.should_include_local_user_in_dir(user_id):
                             continue
 
                         to_insert.add(user_id)
@@ -259,7 +265,7 @@ def _get_next_batch(
                         if not self.hs.is_mine_id(user_id):
                             continue
 
-                        if self.get_if_app_services_interested_in_user(user_id):
+                        if not await self.should_include_local_user_in_dir(user_id):
                             continue
 
                         for other_user_id in users_with_profile:
@@ -349,10 +355,11 @@ def _get_next_batch(txn: LoggingTransaction) -> Optional[List[str]]:
         )
 
         for user_id in users_to_work_on:
-            profile = await self.get_profileinfo(get_localpart_from_id(user_id))
-            await self.update_profile_in_user_dir(
-                user_id, profile.display_name, profile.avatar_url
-            )
+            if await self.should_include_local_user_in_dir(user_id):
+                profile = await self.get_profileinfo(get_localpart_from_id(user_id))
+                await self.update_profile_in_user_dir(
+                    user_id, profile.display_name, profile.avatar_url
+                )
 
             # We've finished processing a user. Delete it from the table.
             await self.db_pool.simple_delete_one(
@@ -369,6 +376,24 @@ def _get_next_batch(txn: LoggingTransaction) -> Optional[List[str]]:
 
         return len(users_to_work_on)
 
+    async def should_include_local_user_in_dir(self, user: str) -> bool:
+        """Certain classes of local user are omitted from the user directory.
+        Is this user one of them?
+        """
+        # App service users aren't usually contactable, so exclude them.
+        if self.get_if_app_services_interested_in_user(user):
+            # TODO we might want to make this configurable for each app service
+            return False
+
+        # Support users are for diagnostics and should not appear in the user directory.
+        if await self.is_support_user(user):
+            return False
+
+        # Deactivated users aren't contactable, so should not appear in the user directory.
+        if await self.get_user_deactivated_status(user):
+            return False
+        return True
+
     async def is_room_world_readable_or_publicly_joinable(self, room_id: str) -> bool:
         """Check if the room is either world_readable or publically joinable"""
 
diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index 2988befb21b4..866b418367a4 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -483,8 +483,6 @@ def _add_user_to_room(
 
 
 class TestUserDirSearchDisabled(unittest.HomeserverTestCase):
-    user_id = "@test:test"
-
     servlets = [
         user_directory.register_servlets,
         room.register_servlets,
@@ -504,16 +502,21 @@ def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
     def test_disabling_room_list(self) -> None:
         self.config.userdirectory.user_directory_search_enabled = True
 
-        # First we create a room with another user so that user dir is non-empty
-        # for our user
-        self.helper.create_room_as(self.user_id)
+        # Create two users and put them in the same room.
+        u1 = self.register_user("user1", "pass")
+        u1_token = self.login(u1, "pass")
         u2 = self.register_user("user2", "pass")
-        room = self.helper.create_room_as(self.user_id)
-        self.helper.join(room, user=u2)
+        u2_token = self.login(u2, "pass")
+
+        room = self.helper.create_room_as(u1, tok=u1_token)
+        self.helper.join(room, user=u2, tok=u2_token)
 
-        # Assert user directory is not empty
+        # Each should see the other when searching the user directory.
         channel = self.make_request(
-            "POST", b"user_directory/search", b'{"search_term":"user2"}'
+            "POST",
+            b"user_directory/search",
+            b'{"search_term":"user2"}',
+            access_token=u1_token,
         )
         self.assertEquals(200, channel.code, channel.result)
         self.assertTrue(len(channel.json_body["results"]) > 0)
@@ -521,7 +524,10 @@ def test_disabling_room_list(self) -> None:
         # Disable user directory and check search returns nothing
         self.config.userdirectory.user_directory_search_enabled = False
         channel = self.make_request(
-            "POST", b"user_directory/search", b'{"search_term":"user2"}'
+            "POST",
+            b"user_directory/search",
+            b'{"search_term":"user2"}',
+            access_token=u1_token,
         )
         self.assertEquals(200, channel.code, channel.result)
         self.assertTrue(len(channel.json_body["results"]) == 0)

From 3f969d1a6d453c51c367a3a0487dc92d2db541a2 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Thu, 30 Sep 2021 19:14:42 +0100
Subject: [PATCH 02/12] Introduce register_appservice_user to
 HomeserverTestCase

already used in one test suite, and I want to make use of it shortly for
the user_directory tests.
---
 tests/rest/client/test_login.py | 17 +++++------------
 tests/unittest.py               | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/tests/rest/client/test_login.py b/tests/rest/client/test_login.py
index 371615a015f3..3f248baca515 100644
--- a/tests/rest/client/test_login.py
+++ b/tests/rest/client/test_login.py
@@ -1064,13 +1064,6 @@ class AppserviceLoginRestServletTestCase(unittest.HomeserverTestCase):
         register.register_servlets,
     ]
 
-    def register_as_user(self, username):
-        self.make_request(
-            b"POST",
-            "/_matrix/client/r0/register?access_token=%s" % (self.service.token,),
-            {"username": username},
-        )
-
     def make_homeserver(self, reactor, clock):
         self.hs = self.setup_test_homeserver()
 
@@ -1107,7 +1100,7 @@ def make_homeserver(self, reactor, clock):
 
     def test_login_appservice_user(self):
         """Test that an appservice user can use /login"""
-        self.register_as_user(AS_USER)
+        self.register_appservice_user(AS_USER, self.service.token)
 
         params = {
             "type": login.LoginRestServlet.APPSERVICE_TYPE,
@@ -1121,7 +1114,7 @@ def test_login_appservice_user(self):
 
     def test_login_appservice_user_bot(self):
         """Test that the appservice bot can use /login"""
-        self.register_as_user(AS_USER)
+        self.register_appservice_user(AS_USER, self.service.token)
 
         params = {
             "type": login.LoginRestServlet.APPSERVICE_TYPE,
@@ -1135,7 +1128,7 @@ def test_login_appservice_user_bot(self):
 
     def test_login_appservice_wrong_user(self):
         """Test that non-as users cannot login with the as token"""
-        self.register_as_user(AS_USER)
+        self.register_appservice_user(AS_USER, self.service.token)
 
         params = {
             "type": login.LoginRestServlet.APPSERVICE_TYPE,
@@ -1149,7 +1142,7 @@ def test_login_appservice_wrong_user(self):
 
     def test_login_appservice_wrong_as(self):
         """Test that as users cannot login with wrong as token"""
-        self.register_as_user(AS_USER)
+        self.register_appservice_user(AS_USER, self.service.token)
 
         params = {
             "type": login.LoginRestServlet.APPSERVICE_TYPE,
@@ -1165,7 +1158,7 @@ def test_login_appservice_no_token(self):
         """Test that users must provide a token when using the appservice
         login method
         """
-        self.register_as_user(AS_USER)
+        self.register_appservice_user(AS_USER, self.service.token)
 
         params = {
             "type": login.LoginRestServlet.APPSERVICE_TYPE,
diff --git a/tests/unittest.py b/tests/unittest.py
index 5f93ebf1479a..293b4d856518 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -22,6 +22,7 @@
 import time
 from typing import Callable, Dict, Iterable, Optional, Tuple, Type, TypeVar, Union
 from unittest.mock import Mock, patch
+from urllib.parse import quote
 
 from canonicaljson import json
 
@@ -596,6 +597,25 @@ def register_user(
         user_id = channel.json_body["user_id"]
         return user_id
 
+    def register_appservice_user(
+        self,
+        username: str,
+        appservice_token: str,
+    ) -> str:
+        """Register an appservice user. Requires the client-facing registration API be registered.
+
+        Returns the MXID of the new user."""
+        urlencoded_token = quote(appservice_token)
+        channel = self.make_request(
+            "POST",
+            f"/_matrix/client/r0/register?access_token={urlencoded_token}",
+            {
+                "username": username,
+                "type": "m.login.application_service",
+            },
+        )
+        return channel.json_body["user_id"]
+
     def login(
         self,
         username,

From 76751aa15691035e594ae2c8b37cb8dc0d829fb0 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Thu, 30 Sep 2021 19:15:46 +0100
Subject: [PATCH 03/12] Check appservice profile changes are ignored

---
 tests/handlers/test_user_directory.py | 41 +++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index 866b418367a4..588b6e983d0b 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -20,7 +20,8 @@
 import synapse.rest.admin
 from synapse.api.constants import UserTypes
 from synapse.api.room_versions import RoomVersion, RoomVersions
-from synapse.rest.client import login, room, user_directory
+from synapse.appservice import ApplicationService
+from synapse.rest.client import login, register, room, user_directory
 from synapse.server import HomeServer
 from synapse.storage.roommember import ProfileInfo
 from synapse.types import create_requester
@@ -47,13 +48,29 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase):
     servlets = [
         login.register_servlets,
         synapse.rest.admin.register_servlets,
+        register.register_servlets,
         room.register_servlets,
     ]
 
     def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
         config = self.default_config()
         config["update_user_directory"] = True
-        return self.setup_test_homeserver(config=config)
+
+        self.appservice = ApplicationService(
+            token="i_am_an_app_service",
+            hostname="test",
+            id="1234",
+            namespaces={"users": [{"regex": r"@as_user.*", "exclusive": True}]},
+            sender="@as:test",
+        )
+
+        mock_load_appservices = Mock(return_value=[self.appservice])
+        with patch(
+            "synapse.storage.databases.main.appservice.load_appservices",
+            mock_load_appservices,
+        ):
+            hs = self.setup_test_homeserver(config=config)
+        return hs
 
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.store = hs.get_datastore()
@@ -125,6 +142,26 @@ def test_handle_local_profile_change_with_deactivated_user(self) -> None:
         profile = self.get_success(self.store.get_user_in_directory(r_user_id))
         self.assertTrue(profile is None)
 
+    def test_handle_local_profile_change_with_appservice_user(self) -> None:
+        # create user
+        as_user_id = self.register_appservice_user(
+            "as_user_alice", self.appservice.token
+        )
+
+        # profile is not in directory
+        profile = self.get_success(self.store.get_user_in_directory(as_user_id))
+        self.assertTrue(profile is None)
+
+        # update profile
+        profile_info = ProfileInfo(avatar_url="avatar_url", display_name="4L1c3")
+        self.get_success(
+            self.handler.handle_local_profile_change(as_user_id, profile_info)
+        )
+
+        # profile is still not in directory
+        profile = self.get_success(self.store.get_user_in_directory(as_user_id))
+        self.assertTrue(profile is None)
+
     def test_handle_user_deactivated_support_user(self) -> None:
         s_user_id = "@support:test"
         self.get_success(

From 010314d1e3cd41a9f5136e6fa4b3d7aa89b544f1 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Thu, 30 Sep 2021 21:00:02 +0100
Subject: [PATCH 04/12] test population excludes users it should

---
 tests/storage/test_user_directory.py | 95 +++++++++++++++++++++++++++-
 1 file changed, 93 insertions(+), 2 deletions(-)

diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index 74c8a8599e7d..26ecfc7fe95f 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -12,11 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import Dict, List, Set, Tuple
+from unittest.mock import Mock, patch
 
 from twisted.test.proto_helpers import MemoryReactor
 
+from synapse.api.constants import UserTypes
+from synapse.appservice import ApplicationService
 from synapse.rest import admin
-from synapse.rest.client import login, room
+from synapse.rest.client import login, register, room
 from synapse.server import HomeServer
 from synapse.storage import DataStore
 from synapse.util import Clock
@@ -64,6 +67,14 @@ async def get_users_who_share_private_rooms(self) -> List[Dict[str, str]]:
             ["user_id", "other_user_id", "room_id"],
         )
 
+    async def get_users_in_user_directory(self) -> Set[str]:
+        result = await self.store.db_pool.simple_select_list(
+            "user_directory",
+            None,
+            ["user_id"],
+        )
+        return {row["user_id"] for row in result}
+
 
 class UserDirectoryInitialPopulationTestcase(HomeserverTestCase):
     """Ensure that rebuilding the directory writes the correct data to the DB.
@@ -74,10 +85,28 @@ class UserDirectoryInitialPopulationTestcase(HomeserverTestCase):
 
     servlets = [
         login.register_servlets,
-        admin.register_servlets_for_client_rest_resource,
+        admin.register_servlets,
         room.register_servlets,
+        register.register_servlets,
     ]
 
+    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
+        self.appservice = ApplicationService(
+            token="i_am_an_app_service",
+            hostname="test",
+            id="1234",
+            namespaces={"users": [{"regex": r"@as_user.*", "exclusive": True}]},
+            sender="@as:test",
+        )
+
+        mock_load_appservices = Mock(return_value=[self.appservice])
+        with patch(
+            "synapse.storage.databases.main.appservice.load_appservices",
+            mock_load_appservices,
+        ):
+            hs = super().make_homeserver(reactor, clock)
+        return hs
+
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.store = hs.get_datastore()
         self.user_dir_helper = GetUserDirectoryTables(self.store)
@@ -204,6 +233,68 @@ def test_initial(self) -> None:
             {(u1, u3, private_room), (u3, u1, private_room)},
         )
 
+    def test_population_excludes_support_user(self) -> None:
+        support = "@support1:test"
+        self.get_success(
+            self.store.register_user(
+                user_id=support, password_hash=None, user_type=UserTypes.SUPPORT
+            )
+        )
+        # Check the support user is not in the directory.
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, set())
+
+        # TODO add support user to a public and private room. Check that
+        # users_in_public_rooms and users_who_share_private_rooms is empty.
+
+        # Rebuild the directory. It should still exclude the support user.
+        self._purge_and_rebuild_user_dir()
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, set())
+
+    def test_population_excludes_deactivated_user(self) -> None:
+        user = self.register_user("naughty", "pass")
+        admin = self.register_user("admin", "pass", admin=True)
+        admin_token = self.login(admin, "pass")
+
+        # Directory contains both users to start with.
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, {admin, user})
+
+        # Deactivate the user.
+        channel = self.make_request(
+            "PUT",
+            f"/_synapse/admin/v2/users/{user}",
+            access_token=admin_token,
+            content={"deactivated": True},
+        )
+        self.assertEqual(channel.code, 200)
+        self.assertEqual(channel.json_body["deactivated"], True)
+
+        # They should no longer be in the directory.
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, {admin})
+
+        # Rebuild the user dir. The deactivated user should still be missing.
+        self._purge_and_rebuild_user_dir()
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, {admin})
+
+    def test_population_excludes_appservice_user(self) -> None:
+        # Register an AS user.
+        self.register_appservice_user("as_user_potato", self.appservice.token)
+        # TODO put this user in a public and private room with someone else
+
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, set())
+        # TODO assert the room sharing tables are as expected
+
+        self._purge_and_rebuild_user_dir()
+
+        # TODO assert the room sharing tables are as expected
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, set())
+
 
 class UserDirectoryStoreTestCase(HomeserverTestCase):
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:

From a8a12a50977adbf8bb6116ad80687e7d7257d400 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Thu, 30 Sep 2021 21:20:06 +0100
Subject: [PATCH 05/12] Extra assertion in test_initial

---
 tests/storage/test_user_directory.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index 26ecfc7fe95f..618de46ca25e 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -233,6 +233,10 @@ def test_initial(self) -> None:
             {(u1, u3, private_room), (u3, u1, private_room)},
         )
 
+        # All three should have entries in the directory
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, {u1, u2, u3})
+
     def test_population_excludes_support_user(self) -> None:
         support = "@support1:test"
         self.get_success(

From 8b614d7ba9a5696ed03ca7bc537b76bc69a274f1 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Thu, 30 Sep 2021 21:27:12 +0100
Subject: [PATCH 06/12] Changelog

---
 changelog.d/10960.bugfix | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/10960.bugfix

diff --git a/changelog.d/10960.bugfix b/changelog.d/10960.bugfix
new file mode 100644
index 000000000000..b4f1c228ea0e
--- /dev/null
+++ b/changelog.d/10960.bugfix
@@ -0,0 +1 @@
+Fix a long-standing bug where rebuilding the user directory wouldn't exclude support and disabled users.
\ No newline at end of file

From 8c41d2e9e040925be67c6d39b88043df8d6ff781 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Fri, 1 Oct 2021 12:53:05 +0100
Subject: [PATCH 07/12] Improve docstring of `register_appservice_user`

---
 tests/unittest.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/tests/unittest.py b/tests/unittest.py
index 293b4d856518..25943522320b 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -602,10 +602,18 @@ def register_appservice_user(
         username: str,
         appservice_token: str,
     ) -> str:
-        """Register an appservice user. Requires the client-facing registration API be registered.
+        """Register an appservice user as an application service.
+        Requires the client-facing registration API be registered.
 
-        Returns the MXID of the new user."""
-        urlencoded_token = quote(appservice_token)
+        Args:
+            username: the user to be registered by an application service.
+                Should be a full username, i.e. ""@localpart:hostname" as opposed to just "localpart"
+            appservice_token: the acccess token for that application service.
+
+        Raises: if the request to '/register' does not return 200 OK.
+
+        Returns: the MXID of the new user.
+        """
         channel = self.make_request(
             "POST",
             f"/_matrix/client/r0/register?access_token={urlencoded_token}",

From 667ec6e30396b89836c8f2d67dc3af8e80037d4b Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Fri, 1 Oct 2021 12:54:25 +0100
Subject: [PATCH 08/12] Use access_token parameter

---
 tests/unittest.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/unittest.py b/tests/unittest.py
index 25943522320b..7746c98095cb 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -22,7 +22,6 @@
 import time
 from typing import Callable, Dict, Iterable, Optional, Tuple, Type, TypeVar, Union
 from unittest.mock import Mock, patch
-from urllib.parse import quote
 
 from canonicaljson import json
 
@@ -616,11 +615,12 @@ def register_appservice_user(
         """
         channel = self.make_request(
             "POST",
-            f"/_matrix/client/r0/register?access_token={urlencoded_token}",
+            "/_matrix/client/r0/register",
             {
                 "username": username,
                 "type": "m.login.application_service",
             },
+            access_token=appservice_token,
         )
         return channel.json_body["user_id"]
 

From c8bb5615bd7958e34aa249b276c4e021795c62aa Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Fri, 1 Oct 2021 12:54:34 +0100
Subject: [PATCH 09/12] Assert 200 response!

---
 tests/unittest.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unittest.py b/tests/unittest.py
index 7746c98095cb..1598ef7bf6af 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -622,6 +622,7 @@ def register_appservice_user(
             },
             access_token=appservice_token,
         )
+        self.assertEqual(channel.code, 200, channel.json_body)
         return channel.json_body["user_id"]
 
     def login(

From aeeb6973c395029c5aa0e3ae9dcb89bb6d74a5ad Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Fri, 1 Oct 2021 16:14:37 +0100
Subject: [PATCH 10/12] Add tests which interrogate room sharing tables

---
 tests/handlers/test_user_directory.py | 133 ++++++++++++++++++++++++++
 tests/storage/test_user_directory.py  |  99 ++++++++++++++-----
 2 files changed, 206 insertions(+), 26 deletions(-)

diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py
index 588b6e983d0b..b3c3af113b28 100644
--- a/tests/handlers/test_user_directory.py
+++ b/tests/handlers/test_user_directory.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Tuple
 from unittest.mock import Mock, patch
 from urllib.parse import quote
 
@@ -29,6 +30,7 @@
 
 from tests import unittest
 from tests.storage.test_user_directory import GetUserDirectoryTables
+from tests.test_utils.event_injection import inject_member_event
 from tests.unittest import override_config
 
 
@@ -79,6 +81,137 @@ def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.event_creation_handler = self.hs.get_event_creation_handler()
         self.user_dir_helper = GetUserDirectoryTables(self.store)
 
+    def test_normal_user_pair(self) -> None:
+        """Sanity check that the room-sharing tables are updated correctly."""
+        alice = self.register_user("alice", "pass")
+        alice_token = self.login(alice, "pass")
+        bob = self.register_user("bob", "pass")
+        bob_token = self.login(bob, "pass")
+
+        public = self.helper.create_room_as(
+            alice,
+            is_public=True,
+            extra_content={"visibility": "public"},
+            tok=alice_token,
+        )
+        private = self.helper.create_room_as(alice, is_public=False, tok=alice_token)
+        self.helper.invite(private, alice, bob, tok=alice_token)
+        self.helper.join(public, bob, tok=bob_token)
+        self.helper.join(private, bob, tok=bob_token)
+
+        # Alice also makes a second public room but no-one else joins
+        public2 = self.helper.create_room_as(
+            alice,
+            is_public=True,
+            extra_content={"visibility": "public"},
+            tok=alice_token,
+        )
+
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        in_public = self.get_success(self.user_dir_helper.get_users_in_public_rooms())
+        in_private = self.get_success(
+            self.user_dir_helper.get_users_who_share_private_rooms()
+        )
+
+        self.assertEqual(users, {alice, bob})
+        self.assertEqual(
+            set(in_public), {(alice, public), (bob, public), (alice, public2)}
+        )
+        self.assertEqual(
+            self.user_dir_helper._compress_shared(in_private),
+            {(alice, bob, private), (bob, alice, private)},
+        )
+
+    # The next three tests (test_population_excludes_*) all setup
+    #   - A normal user included in the user dir
+    #   - A public and private room created by that user
+    #   - A user excluded from the room dir, belonging to both rooms
+
+    # They match similar logic in storage/test_user_directory. But that tests
+    # rebuilding the directory; this tests updating it incrementally.
+
+    def test_excludes_support_user(self) -> None:
+        alice = self.register_user("alice", "pass")
+        alice_token = self.login(alice, "pass")
+        support = "@support1:test"
+        self.get_success(
+            self.store.register_user(
+                user_id=support, password_hash=None, user_type=UserTypes.SUPPORT
+            )
+        )
+
+        public, private = self._create_rooms_and_inject_memberships(
+            alice, alice_token, support
+        )
+        self._check_only_one_user_in_directory(alice, public)
+
+    def test_excludes_deactivated_user(self) -> None:
+        admin = self.register_user("admin", "pass", admin=True)
+        admin_token = self.login(admin, "pass")
+        user = self.register_user("naughty", "pass")
+
+        # Deactivate the user.
+        channel = self.make_request(
+            "PUT",
+            f"/_synapse/admin/v2/users/{user}",
+            access_token=admin_token,
+            content={"deactivated": True},
+        )
+        self.assertEqual(channel.code, 200)
+        self.assertEqual(channel.json_body["deactivated"], True)
+
+        # Join the deactivated user to rooms owned by the admin.
+        # Is this something that could actually happen outside of a test?
+        public, private = self._create_rooms_and_inject_memberships(
+            admin, admin_token, user
+        )
+        self._check_only_one_user_in_directory(admin, public)
+
+    def test_excludes_appservices_user(self) -> None:
+        # Register an AS user.
+        user = self.register_user("user", "pass")
+        token = self.login(user, "pass")
+        as_user = self.register_appservice_user("as_user_potato", self.appservice.token)
+
+        # Join the AS user to rooms owned by the normal user.
+        public, private = self._create_rooms_and_inject_memberships(
+            user, token, as_user
+        )
+        self._check_only_one_user_in_directory(user, public)
+
+    def _create_rooms_and_inject_memberships(
+        self, creator: str, token: str, joiner: str
+    ) -> Tuple[str, str]:
+        """Create a public and private room as a normal user.
+        Then get the `joiner` into those rooms.
+        """
+        # TODO: Duplicates the same-named method in UserDirectoryInitialPopulationTest.
+        public_room = self.helper.create_room_as(
+            creator,
+            is_public=True,
+            # See https://github.com/matrix-org/synapse/issues/10951
+            extra_content={"visibility": "public"},
+            tok=token,
+        )
+        private_room = self.helper.create_room_as(creator, is_public=False, tok=token)
+
+        # HACK: get the user into these rooms
+        self.get_success(inject_member_event(self.hs, public_room, joiner, "join"))
+        self.get_success(inject_member_event(self.hs, private_room, joiner, "join"))
+
+        return public_room, private_room
+
+    def _check_only_one_user_in_directory(self, user: str, public: str) -> None:
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        in_public = self.get_success(self.user_dir_helper.get_users_in_public_rooms())
+        in_private = self.get_success(
+            self.user_dir_helper.get_users_who_share_private_rooms()
+        )
+
+        self.assertEqual(users, {user})
+        self.assertEqual(set(in_public), {(user, public)})
+        self.assertEqual(in_private, [])
+
     def test_handle_local_profile_change_with_support_user(self) -> None:
         support_user_id = "@support:test"
         self.get_success(
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index 618de46ca25e..9f4410f68f0b 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -24,6 +24,7 @@
 from synapse.storage import DataStore
 from synapse.util import Clock
 
+from tests.test_utils.event_injection import inject_member_event
 from tests.unittest import HomeserverTestCase, override_config
 
 ALICE = "@alice:a"
@@ -237,34 +238,77 @@ def test_initial(self) -> None:
         users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
         self.assertEqual(users, {u1, u2, u3})
 
+    # The next three tests (test_population_excludes_*) all setup
+    #   - A normal user included in the user dir
+    #   - A public and private room created by that user
+    #   - A user excluded from the room dir, belonging to both rooms
+
+    # They match similar logic in handlers/test_user_directory.py But that tests
+    # updating the directory; this tests rebuilding it from scratch.
+
+    def _create_rooms_and_inject_memberships(
+        self, creator: str, token: str, joiner: str
+    ) -> Tuple[str, str]:
+        """Create a public and private room as a normal user.
+        Then get the `joiner` into those rooms.
+        """
+        public_room = self.helper.create_room_as(
+            creator,
+            is_public=True,
+            # See https://github.com/matrix-org/synapse/issues/10951
+            extra_content={"visibility": "public"},
+            tok=token,
+        )
+        private_room = self.helper.create_room_as(creator, is_public=False, tok=token)
+
+        # HACK: get the user into these rooms
+        self.get_success(inject_member_event(self.hs, public_room, joiner, "join"))
+        self.get_success(inject_member_event(self.hs, private_room, joiner, "join"))
+
+        return public_room, private_room
+
+    def _check_room_sharing_tables(
+        self, normal_user: str, public_room: str, private_room: str
+    ) -> None:
+        # After rebuilding the directory, we should only see the normal user.
+        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
+        self.assertEqual(users, {normal_user})
+        in_public_rooms = self.get_success(
+            self.user_dir_helper.get_users_in_public_rooms()
+        )
+        self.assertEqual(set(in_public_rooms), {(normal_user, public_room)})
+        in_private_rooms = self.get_success(
+            self.user_dir_helper.get_users_who_share_private_rooms()
+        )
+        self.assertEqual(in_private_rooms, [])
+
     def test_population_excludes_support_user(self) -> None:
+        # Create a normal and support user.
+        user = self.register_user("user", "pass")
+        token = self.login(user, "pass")
         support = "@support1:test"
         self.get_success(
             self.store.register_user(
                 user_id=support, password_hash=None, user_type=UserTypes.SUPPORT
             )
         )
-        # Check the support user is not in the directory.
-        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
-        self.assertEqual(users, set())
 
-        # TODO add support user to a public and private room. Check that
-        # users_in_public_rooms and users_who_share_private_rooms is empty.
+        # Join the support user to rooms owned by the normal user.
+        public, private = self._create_rooms_and_inject_memberships(
+            user, token, support
+        )
 
-        # Rebuild the directory. It should still exclude the support user.
+        # Rebuild the directory.
         self._purge_and_rebuild_user_dir()
-        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
-        self.assertEqual(users, set())
+
+        # Check the support user is not in the directory.
+        self._check_room_sharing_tables(user, public, private)
 
     def test_population_excludes_deactivated_user(self) -> None:
         user = self.register_user("naughty", "pass")
         admin = self.register_user("admin", "pass", admin=True)
         admin_token = self.login(admin, "pass")
 
-        # Directory contains both users to start with.
-        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
-        self.assertEqual(users, {admin, user})
-
         # Deactivate the user.
         channel = self.make_request(
             "PUT",
@@ -275,29 +319,32 @@ def test_population_excludes_deactivated_user(self) -> None:
         self.assertEqual(channel.code, 200)
         self.assertEqual(channel.json_body["deactivated"], True)
 
-        # They should no longer be in the directory.
-        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
-        self.assertEqual(users, {admin})
+        # Join the deactivated user to rooms owned by the admin.
+        # Is this something that could actually happen outside of a test?
+        public, private = self._create_rooms_and_inject_memberships(
+            admin, admin_token, user
+        )
 
-        # Rebuild the user dir. The deactivated user should still be missing.
+        # Rebuild the user dir. The deactivated user should be missing.
         self._purge_and_rebuild_user_dir()
-        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
-        self.assertEqual(users, {admin})
+        self._check_room_sharing_tables(admin, public, private)
 
     def test_population_excludes_appservice_user(self) -> None:
         # Register an AS user.
-        self.register_appservice_user("as_user_potato", self.appservice.token)
-        # TODO put this user in a public and private room with someone else
+        user = self.register_user("user", "pass")
+        token = self.login(user, "pass")
+        as_user = self.register_appservice_user("as_user_potato", self.appservice.token)
 
-        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
-        self.assertEqual(users, set())
-        # TODO assert the room sharing tables are as expected
+        # Join the AS user to rooms owned by the normal user.
+        public, private = self._create_rooms_and_inject_memberships(
+            user, token, as_user
+        )
 
+        # Rebuild the directory.
         self._purge_and_rebuild_user_dir()
 
-        # TODO assert the room sharing tables are as expected
-        users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
-        self.assertEqual(users, set())
+        # Check the AS user is not in the directory.
+        self._check_room_sharing_tables(user, public, private)
 
 
 class UserDirectoryStoreTestCase(HomeserverTestCase):

From c090ab6e0b76160f85b2d28e1bb04f0131686d16 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Fri, 1 Oct 2021 16:44:42 +0100
Subject: [PATCH 11/12] Correctly exclude from  users_who_share_private_room

---
 .../storage/databases/main/user_directory.py  | 25 ++++++++-----------
 1 file changed, 10 insertions(+), 15 deletions(-)

diff --git a/synapse/storage/databases/main/user_directory.py b/synapse/storage/databases/main/user_directory.py
index f22423afca45..5f538947ecf8 100644
--- a/synapse/storage/databases/main/user_directory.py
+++ b/synapse/storage/databases/main/user_directory.py
@@ -40,12 +40,10 @@
 
 logger = logging.getLogger(__name__)
 
-
 TEMP_TABLE = "_temp_populate_user_directory"
 
 
 class UserDirectoryBackgroundUpdateStore(StateDeltasStore):
-
     # How many records do we calculate before sending it to
     # add_users_who_share_private_rooms?
     SHARE_PRIVATE_WORKING_SET = 500
@@ -235,13 +233,16 @@ def _get_next_batch(
                 )
 
                 users_with_profile = await self.get_users_in_room_with_profiles(room_id)
+                # Throw away users excluded from the directory.
+                users_with_profile = {
+                    user_id: profile
+                    for user_id, profile in users_with_profile.items()
+                    if not self.hs.is_mine_id(user_id)
+                    or await self.should_include_local_user_in_dir(user_id)
+                }
 
                 # Update each user in the user directory.
                 for user_id, profile in users_with_profile.items():
-                    if self.hs.is_mine_id(
-                        user_id
-                    ) and not await self.should_include_local_user_in_dir(user_id):
-                        continue
                     await self.update_profile_in_user_dir(
                         user_id, profile.display_name, profile.avatar_url
                     )
@@ -250,11 +251,6 @@ def _get_next_batch(
 
                 if is_public:
                     for user_id in users_with_profile:
-                        if self.hs.is_mine_id(
-                            user_id
-                        ) and not await self.should_include_local_user_in_dir(user_id):
-                            continue
-
                         to_insert.add(user_id)
 
                     if to_insert:
@@ -262,12 +258,12 @@ def _get_next_batch(
                         to_insert.clear()
                 else:
                     for user_id in users_with_profile:
+                        # We want the set of pairs (L, M) where L and M are
+                        # in `users_with_profile` and L is local.
+                        # Do so by looking for the local user L first.
                         if not self.hs.is_mine_id(user_id):
                             continue
 
-                        if not await self.should_include_local_user_in_dir(user_id):
-                            continue
-
                         for other_user_id in users_with_profile:
                             if user_id == other_user_id:
                                 continue
@@ -562,7 +558,6 @@ async def update_user_directory_stream_pos(self, stream_id: Optional[int]) -> No
 
 
 class UserDirectoryStore(UserDirectoryBackgroundUpdateStore):
-
     # How many records do we calculate before sending it to
     # add_users_who_share_private_rooms?
     SHARE_PRIVATE_WORKING_SET = 500

From 3dc823f16c3f9443ad2be9fea8bbfa63e7a27d5e Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Mon, 4 Oct 2021 12:20:01 +0100
Subject: [PATCH 12/12] `setup` is not a verb

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
---
 tests/storage/test_user_directory.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index 9f4410f68f0b..6884ca9b7a3f 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -238,7 +238,7 @@ def test_initial(self) -> None:
         users = self.get_success(self.user_dir_helper.get_users_in_user_directory())
         self.assertEqual(users, {u1, u2, u3})
 
-    # The next three tests (test_population_excludes_*) all setup
+    # The next three tests (test_population_excludes_*) all set up
     #   - A normal user included in the user dir
     #   - A public and private room created by that user
     #   - A user excluded from the room dir, belonging to both rooms