Add a Synapse Module for configuring presence update routing (#9491)

At the moment, if you'd like to share presence between local or remote users, those users must be sharing a room together. This isn't always the most convenient or useful situation though.

This PR adds a module to Synapse that will allow deployments to set up extra logic on where presence updates should be routed. The module must implement two methods, `get_users_for_states` and `get_interested_users`. These methods are given presence updates or user IDs and must return information that Synapse will use to grant passing presence updates around.

A method is additionally added to `ModuleApi` which allows triggering a set of users to receive the current, online presence information for all users they are considered interested in. This is the equivalent of that user receiving presence information during an initial sync. 

The goal of this module is to be fairly generic and useful for a variety of applications, with hard requirements being:

* Sending state for a specific set or all known users to a defined set of local and remote users.
* The ability to trigger an initial sync for specific users, so they receive all current state.

3 files changed, 569 insertions, 13 deletions

diff --git a/tests/events/test_presence_router.py b/tests/events/test_presence_router.py
new file mode 100644
index 0000000000..c6e547f11c
--- /dev/null
+++ b/tests/events/test_presence_router.py
@@ -0,0 +1,386 @@
+# -*- coding: utf-8 -*-
+# Copyright 2021 The Matrix.org Foundation C.I.C.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# 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 Dict, Iterable, List, Optional, Set, Tuple, Union
+
+from mock import Mock
+
+import attr
+
+from synapse.api.constants import EduTypes
+from synapse.events.presence_router import PresenceRouter
+from synapse.federation.units import Transaction
+from synapse.handlers.presence import UserPresenceState
+from synapse.module_api import ModuleApi
+from synapse.rest import admin
+from synapse.rest.client.v1 import login, presence, room
+from synapse.types import JsonDict, StreamToken, create_requester
+
+from tests.handlers.test_sync import generate_sync_config
+from tests.unittest import FederatingHomeserverTestCase, TestCase, override_config
+
+
+@attr.s
+class PresenceRouterTestConfig:
+    users_who_should_receive_all_presence = attr.ib(type=List[str], default=[])
+
+
+class PresenceRouterTestModule:
+    def __init__(self, config: PresenceRouterTestConfig, module_api: ModuleApi):
+        self._config = config
+        self._module_api = module_api
+
+    async def get_users_for_states(
+        self, state_updates: Iterable[UserPresenceState]
+    ) -> Dict[str, Set[UserPresenceState]]:
+        users_to_state = {
+            user_id: set(state_updates)
+            for user_id in self._config.users_who_should_receive_all_presence
+        }
+        return users_to_state
+
+    async def get_interested_users(
+        self, user_id: str
+    ) -> Union[Set[str], PresenceRouter.ALL_USERS]:
+        if user_id in self._config.users_who_should_receive_all_presence:
+            return PresenceRouter.ALL_USERS
+
+        return set()
+
+    @staticmethod
+    def parse_config(config_dict: dict) -> PresenceRouterTestConfig:
+        """Parse a configuration dictionary from the homeserver config, do
+        some validation and return a typed PresenceRouterConfig.
+
+        Args:
+            config_dict: The configuration dictionary.
+
+        Returns:
+            A validated config object.
+        """
+        # Initialise a typed config object
+        config = PresenceRouterTestConfig()
+
+        config.users_who_should_receive_all_presence = config_dict.get(
+            "users_who_should_receive_all_presence"
+        )
+
+        return config
+
+
+class PresenceRouterTestCase(FederatingHomeserverTestCase):
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+        room.register_servlets,
+        presence.register_servlets,
+    ]
+
+    def make_homeserver(self, reactor, clock):
+        return self.setup_test_homeserver(
+            federation_transport_client=Mock(spec=["send_transaction"]),
+        )
+
+    def prepare(self, reactor, clock, homeserver):
+        self.sync_handler = self.hs.get_sync_handler()
+        self.module_api = homeserver.get_module_api()
+
+    @override_config(
+        {
+            "presence": {
+                "presence_router": {
+                    "module": __name__ + ".PresenceRouterTestModule",
+                    "config": {
+                        "users_who_should_receive_all_presence": [
+                            "@presence_gobbler:test",
+                        ]
+                    },
+                }
+            },
+            "send_federation": True,
+        }
+    )
+    def test_receiving_all_presence(self):
+        """Test that a user that does not share a room with another other can receive
+        presence for them, due to presence routing.
+        """
+        # Create a user who should receive all presence of others
+        self.presence_receiving_user_id = self.register_user(
+            "presence_gobbler", "monkey"
+        )
+        self.presence_receiving_user_tok = self.login("presence_gobbler", "monkey")
+
+        # And two users who should not have any special routing
+        self.other_user_one_id = self.register_user("other_user_one", "monkey")
+        self.other_user_one_tok = self.login("other_user_one", "monkey")
+        self.other_user_two_id = self.register_user("other_user_two", "monkey")
+        self.other_user_two_tok = self.login("other_user_two", "monkey")
+
+        # Put the other two users in a room with each other
+        room_id = self.helper.create_room_as(
+            self.other_user_one_id, tok=self.other_user_one_tok
+        )
+
+        self.helper.invite(
+            room_id,
+            self.other_user_one_id,
+            self.other_user_two_id,
+            tok=self.other_user_one_tok,
+        )
+        self.helper.join(room_id, self.other_user_two_id, tok=self.other_user_two_tok)
+        # User one sends some presence
+        send_presence_update(
+            self,
+            self.other_user_one_id,
+            self.other_user_one_tok,
+            "online",
+            "boop",
+        )
+
+        # Check that the presence receiving user gets user one's presence when syncing
+        presence_updates, sync_token = sync_presence(
+            self, self.presence_receiving_user_id
+        )
+        self.assertEqual(len(presence_updates), 1)
+
+        presence_update = presence_updates[0]  # type: UserPresenceState
+        self.assertEqual(presence_update.user_id, self.other_user_one_id)
+        self.assertEqual(presence_update.state, "online")
+        self.assertEqual(presence_update.status_msg, "boop")
+
+        # Have all three users send presence
+        send_presence_update(
+            self,
+            self.other_user_one_id,
+            self.other_user_one_tok,
+            "online",
+            "user_one",
+        )
+        send_presence_update(
+            self,
+            self.other_user_two_id,
+            self.other_user_two_tok,
+            "online",
+            "user_two",
+        )
+        send_presence_update(
+            self,
+            self.presence_receiving_user_id,
+            self.presence_receiving_user_tok,
+            "online",
+            "presence_gobbler",
+        )
+
+        # Check that the presence receiving user gets everyone's presence
+        presence_updates, _ = sync_presence(
+            self, self.presence_receiving_user_id, sync_token
+        )
+        self.assertEqual(len(presence_updates), 3)
+
+        # But that User One only get itself and User Two's presence
+        presence_updates, _ = sync_presence(self, self.other_user_one_id)
+        self.assertEqual(len(presence_updates), 2)
+
+        found = False
+        for update in presence_updates:
+            if update.user_id == self.other_user_two_id:
+                self.assertEqual(update.state, "online")
+                self.assertEqual(update.status_msg, "user_two")
+                found = True
+
+        self.assertTrue(found)
+
+    @override_config(
+        {
+            "presence": {
+                "presence_router": {
+                    "module": __name__ + ".PresenceRouterTestModule",
+                    "config": {
+                        "users_who_should_receive_all_presence": [
+                            "@presence_gobbler1:test",
+                            "@presence_gobbler2:test",
+                            "@far_away_person:island",
+                        ]
+                    },
+                }
+            },
+            "send_federation": True,
+        }
+    )
+    def test_send_local_online_presence_to_with_module(self):
+        """Tests that send_local_presence_to_users sends local online presence to a set
+        of specified local and remote users, with a custom PresenceRouter module enabled.
+        """
+        # Create a user who will send presence updates
+        self.other_user_id = self.register_user("other_user", "monkey")
+        self.other_user_tok = self.login("other_user", "monkey")
+
+        # And another two users that will also send out presence updates, as well as receive
+        # theirs and everyone else's
+        self.presence_receiving_user_one_id = self.register_user(
+            "presence_gobbler1", "monkey"
+        )
+        self.presence_receiving_user_one_tok = self.login("presence_gobbler1", "monkey")
+        self.presence_receiving_user_two_id = self.register_user(
+            "presence_gobbler2", "monkey"
+        )
+        self.presence_receiving_user_two_tok = self.login("presence_gobbler2", "monkey")
+
+        # Have all three users send some presence updates
+        send_presence_update(
+            self,
+            self.other_user_id,
+            self.other_user_tok,
+            "online",
+            "I'm online!",
+        )
+        send_presence_update(
+            self,
+            self.presence_receiving_user_one_id,
+            self.presence_receiving_user_one_tok,
+            "online",
+            "I'm also online!",
+        )
+        send_presence_update(
+            self,
+            self.presence_receiving_user_two_id,
+            self.presence_receiving_user_two_tok,
+            "unavailable",
+            "I'm in a meeting!",
+        )
+
+        # Mark each presence-receiving user for receiving all user presence
+        self.get_success(
+            self.module_api.send_local_online_presence_to(
+                [
+                    self.presence_receiving_user_one_id,
+                    self.presence_receiving_user_two_id,
+                ]
+            )
+        )
+
+        # Perform a sync for each user
+
+        # The other user should only receive their own presence
+        presence_updates, _ = sync_presence(self, self.other_user_id)
+        self.assertEqual(len(presence_updates), 1)
+
+        presence_update = presence_updates[0]  # type: UserPresenceState
+        self.assertEqual(presence_update.user_id, self.other_user_id)
+        self.assertEqual(presence_update.state, "online")
+        self.assertEqual(presence_update.status_msg, "I'm online!")
+
+        # Whereas both presence receiving users should receive everyone's presence updates
+        presence_updates, _ = sync_presence(self, self.presence_receiving_user_one_id)
+        self.assertEqual(len(presence_updates), 3)
+        presence_updates, _ = sync_presence(self, self.presence_receiving_user_two_id)
+        self.assertEqual(len(presence_updates), 3)
+
+        # Test that sending to a remote user works
+        remote_user_id = "@far_away_person:island"
+
+        # Note that due to the remote user being in our module's
+        # users_who_should_receive_all_presence config, they would have
+        # received user presence updates already.
+        #
+        # Thus we reset the mock, and try sending all online local user
+        # presence again
+        self.hs.get_federation_transport_client().send_transaction.reset_mock()
+
+        # Broadcast local user online presence
+        self.get_success(
+            self.module_api.send_local_online_presence_to([remote_user_id])
+        )
+
+        # Check that the expected presence updates were sent
+        expected_users = [
+            self.other_user_id,
+            self.presence_receiving_user_one_id,
+            self.presence_receiving_user_two_id,
+        ]
+
+        calls = (
+            self.hs.get_federation_transport_client().send_transaction.call_args_list
+        )
+        for call in calls:
+            federation_transaction = call.args[0]  # type: Transaction
+
+            # Get the sent EDUs in this transaction
+            edus = federation_transaction.get_dict()["edus"]
+
+            for edu in edus:
+                # Make sure we're only checking presence-type EDUs
+                if edu["edu_type"] != EduTypes.Presence:
+                    continue
+
+                # EDUs can contain multiple presence updates
+                for presence_update in edu["content"]["push"]:
+                    # Check for presence updates that contain the user IDs we're after
+                    expected_users.remove(presence_update["user_id"])
+
+                    # Ensure that no offline states are being sent out
+                    self.assertNotEqual(presence_update["presence"], "offline")
+
+        self.assertEqual(len(expected_users), 0)
+
+
+def send_presence_update(
+    testcase: TestCase,
+    user_id: str,
+    access_token: str,
+    presence_state: str,
+    status_message: Optional[str] = None,
+) -> JsonDict:
+    # Build the presence body
+    body = {"presence": presence_state}
+    if status_message:
+        body["status_msg"] = status_message
+
+    # Update the user's presence state
+    channel = testcase.make_request(
+        "PUT", "/presence/%s/status" % (user_id,), body, access_token=access_token
+    )
+    testcase.assertEqual(channel.code, 200)
+
+    return channel.json_body
+
+
+def sync_presence(
+    testcase: TestCase,
+    user_id: str,
+    since_token: Optional[StreamToken] = None,
+) -> Tuple[List[UserPresenceState], StreamToken]:
+    """Perform a sync request for the given user and return the user presence updates
+    they've received, as well as the next_batch token.
+
+    This method assumes testcase.sync_handler points to the homeserver's sync handler.
+
+    Args:
+        testcase: The testcase that is currently being run.
+        user_id: The ID of the user to generate a sync response for.
+        since_token: An optional token to indicate from at what point to sync from.
+
+    Returns:
+        A tuple containing a list of presence updates, and the sync response's
+        next_batch token.
+    """
+    requester = create_requester(user_id)
+    sync_config = generate_sync_config(requester.user.to_string())
+    sync_result = testcase.get_success(
+        testcase.sync_handler.wait_for_sync_for_user(
+            requester, sync_config, since_token
+        )
+    )
+
+    return sync_result.presence, sync_result.next_batch
diff --git a/tests/handlers/test_sync.py b/tests/handlers/test_sync.py
index e62586142e..8e950f25c5 100644
--- a/tests/handlers/test_sync.py
+++ b/tests/handlers/test_sync.py
@@ -37,7 +37,7 @@ class SyncTestCase(tests.unittest.HomeserverTestCase):
     def test_wait_for_sync_for_user_auth_blocking(self):
         user_id1 = "@user1:test"
         user_id2 = "@user2:test"
-        sync_config = self._generate_sync_config(user_id1)
+        sync_config = generate_sync_config(user_id1)
         requester = create_requester(user_id1)
 
         self.reactor.advance(100)  # So we get not 0 time
@@ -60,7 +60,7 @@ class SyncTestCase(tests.unittest.HomeserverTestCase):
 
         self.auth_blocking._hs_disabled = False
 
-        sync_config = self._generate_sync_config(user_id2)
+        sync_config = generate_sync_config(user_id2)
         requester = create_requester(user_id2)
 
         e = self.get_failure(
@@ -69,11 +69,12 @@ class SyncTestCase(tests.unittest.HomeserverTestCase):
         )
         self.assertEquals(e.value.errcode, Codes.RESOURCE_LIMIT_EXCEEDED)
 
-    def _generate_sync_config(self, user_id):
-        return SyncConfig(
-            user=UserID(user_id.split(":")[0][1:], user_id.split(":")[1]),
-            filter_collection=DEFAULT_FILTER_COLLECTION,
-            is_guest=False,
-            request_key="request_key",
-            device_id="device_id",
-        )
+
+def generate_sync_config(user_id: str) -> SyncConfig:
+    return SyncConfig(
+        user=UserID(user_id.split(":")[0][1:], user_id.split(":")[1]),
+        filter_collection=DEFAULT_FILTER_COLLECTION,
+        is_guest=False,
+        request_key="request_key",
+        device_id="device_id",
+    )
diff --git a/tests/module_api/test_api.py b/tests/module_api/test_api.py
index edacd1b566..1d1fceeecf 100644
--- a/tests/module_api/test_api.py
+++ b/tests/module_api/test_api.py
@@ -14,25 +14,37 @@
 # limitations under the License.
 from mock import Mock
 
+from synapse.api.constants import EduTypes
 from synapse.events import EventBase
+from synapse.federation.units import Transaction
+from synapse.handlers.presence import UserPresenceState
 from synapse.rest import admin
-from synapse.rest.client.v1 import login, room
+from synapse.rest.client.v1 import login, presence, room
 from synapse.types import create_requester
 
-from tests.unittest import HomeserverTestCase
+from tests.events.test_presence_router import send_presence_update, sync_presence
+from tests.test_utils.event_injection import inject_member_event
+from tests.unittest import FederatingHomeserverTestCase, override_config
 
 
-class ModuleApiTestCase(HomeserverTestCase):
+class ModuleApiTestCase(FederatingHomeserverTestCase):
     servlets = [
         admin.register_servlets,
         login.register_servlets,
         room.register_servlets,
+        presence.register_servlets,
     ]
 
     def prepare(self, reactor, clock, homeserver):
         self.store = homeserver.get_datastore()
         self.module_api = homeserver.get_module_api()
         self.event_creation_handler = homeserver.get_event_creation_handler()
+        self.sync_handler = homeserver.get_sync_handler()
+
+    def make_homeserver(self, reactor, clock):
+        return self.setup_test_homeserver(
+            federation_transport_client=Mock(spec=["send_transaction"]),
+        )
 
     def test_can_register_user(self):
         """Tests that an external module can register a user"""
@@ -205,3 +217,160 @@ class ModuleApiTestCase(HomeserverTestCase):
             )
         )
         self.assertFalse(is_in_public_rooms)
+
+    # The ability to send federation is required by send_local_online_presence_to.
+    @override_config({"send_federation": True})
+    def test_send_local_online_presence_to(self):
+        """Tests that send_local_presence_to_users sends local online presence to local users."""
+        # Create a user who will send presence updates
+        self.presence_receiver_id = self.register_user("presence_receiver", "monkey")
+        self.presence_receiver_tok = self.login("presence_receiver", "monkey")
+
+        # And another user that will send presence updates out
+        self.presence_sender_id = self.register_user("presence_sender", "monkey")
+        self.presence_sender_tok = self.login("presence_sender", "monkey")
+
+        # Put them in a room together so they will receive each other's presence updates
+        room_id = self.helper.create_room_as(
+            self.presence_receiver_id,
+            tok=self.presence_receiver_tok,
+        )
+        self.helper.join(room_id, self.presence_sender_id, tok=self.presence_sender_tok)
+
+        # Presence sender comes online
+        send_presence_update(
+            self,
+            self.presence_sender_id,
+            self.presence_sender_tok,
+            "online",
+            "I'm online!",
+        )
+
+        # Presence receiver should have received it
+        presence_updates, sync_token = sync_presence(self, self.presence_receiver_id)
+        self.assertEqual(len(presence_updates), 1)
+
+        presence_update = presence_updates[0]  # type: UserPresenceState
+        self.assertEqual(presence_update.user_id, self.presence_sender_id)
+        self.assertEqual(presence_update.state, "online")
+
+        # Syncing again should result in no presence updates
+        presence_updates, sync_token = sync_presence(
+            self, self.presence_receiver_id, sync_token
+        )
+        self.assertEqual(len(presence_updates), 0)
+
+        # Trigger sending local online presence
+        self.get_success(
+            self.module_api.send_local_online_presence_to(
+                [
+                    self.presence_receiver_id,
+                ]
+            )
+        )
+
+        # Presence receiver should have received online presence again
+        presence_updates, sync_token = sync_presence(
+            self, self.presence_receiver_id, sync_token
+        )
+        self.assertEqual(len(presence_updates), 1)
+
+        presence_update = presence_updates[0]  # type: UserPresenceState
+        self.assertEqual(presence_update.user_id, self.presence_sender_id)
+        self.assertEqual(presence_update.state, "online")
+
+        # Presence sender goes offline
+        send_presence_update(
+            self,
+            self.presence_sender_id,
+            self.presence_sender_tok,
+            "offline",
+            "I slink back into the darkness.",
+        )
+
+        # Trigger sending local online presence
+        self.get_success(
+            self.module_api.send_local_online_presence_to(
+                [
+                    self.presence_receiver_id,
+                ]
+            )
+        )
+
+        # Presence receiver should *not* have received offline state
+        presence_updates, sync_token = sync_presence(
+            self, self.presence_receiver_id, sync_token
+        )
+        self.assertEqual(len(presence_updates), 0)
+
+    @override_config({"send_federation": True})
+    def test_send_local_online_presence_to_federation(self):
+        """Tests that send_local_presence_to_users sends local online presence to remote users."""
+        # Create a user who will send presence updates
+        self.presence_sender_id = self.register_user("presence_sender", "monkey")
+        self.presence_sender_tok = self.login("presence_sender", "monkey")
+
+        # And a room they're a part of
+        room_id = self.helper.create_room_as(
+            self.presence_sender_id,
+            tok=self.presence_sender_tok,
+        )
+
+        # Mark them as online
+        send_presence_update(
+            self,
+            self.presence_sender_id,
+            self.presence_sender_tok,
+            "online",
+            "I'm online!",
+        )
+
+        # Make up a remote user to send presence to
+        remote_user_id = "@far_away_person:island"
+
+        # Create a join membership event for the remote user into the room.
+        # This allows presence information to flow from one user to the other.
+        self.get_success(
+            inject_member_event(
+                self.hs,
+                room_id,
+                sender=remote_user_id,
+                target=remote_user_id,
+                membership="join",
+            )
+        )
+
+        # The remote user would have received the existing room members' presence
+        # when they joined the room.
+        #
+        # Thus we reset the mock, and try sending online local user
+        # presence again
+        self.hs.get_federation_transport_client().send_transaction.reset_mock()
+
+        # Broadcast local user online presence
+        self.get_success(
+            self.module_api.send_local_online_presence_to([remote_user_id])
+        )
+
+        # Check that a presence update was sent as part of a federation transaction
+        found_update = False
+        calls = (
+            self.hs.get_federation_transport_client().send_transaction.call_args_list
+        )
+        for call in calls:
+            federation_transaction = call.args[0]  # type: Transaction
+
+            # Get the sent EDUs in this transaction
+            edus = federation_transaction.get_dict()["edus"]
+
+            for edu in edus:
+                # Make sure we're only checking presence-type EDUs
+                if edu["edu_type"] != EduTypes.Presence:
+                    continue
+
+                # EDUs can contain multiple presence updates
+                for presence_update in edu["content"]["push"]:
+                    if presence_update["user_id"] == self.presence_sender_id:
+                        found_update = True
+
+        self.assertTrue(found_update)