Add a module callback to react to account data changes (#12327)

Co-authored-by: Sean Quah <8349537+squahtx@users.noreply.github.com>
Co-authored-by: Andrew Morgan <1342360+anoadragon453@users.noreply.github.com>

7 files changed, 250 insertions, 2 deletions

diff --git a/changelog.d/12327.feature b/changelog.d/12327.feature
new file mode 100644
index 0000000000..4fe294f1b1
--- /dev/null
+++ b/changelog.d/12327.feature
@@ -0,0 +1 @@
+Add a module callback to react to account data changes.
diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md
index 21f80efc99..6aa48e1919 100644
--- a/docs/SUMMARY.md
+++ b/docs/SUMMARY.md
@@ -45,6 +45,7 @@
         - [Account validity callbacks](modules/account_validity_callbacks.md)
         - [Password auth provider callbacks](modules/password_auth_provider_callbacks.md)
         - [Background update controller callbacks](modules/background_update_controller_callbacks.md)
+        - [Account data callbacks](modules/account_data_callbacks.md)
         - [Porting a legacy module to the new interface](modules/porting_legacy_module.md)
     - [Workers](workers.md)
       - [Using `synctl` with Workers](synctl_workers.md)
diff --git a/docs/modules/account_data_callbacks.md b/docs/modules/account_data_callbacks.md
new file mode 100644
index 0000000000..25de911627
--- /dev/null
+++ b/docs/modules/account_data_callbacks.md
@@ -0,0 +1,106 @@
+# Account data callbacks
+
+Account data callbacks allow module developers to react to changes of the account data
+of local users. Account data callbacks can be registered using the module API's
+`register_account_data_callbacks` method.
+
+## Callbacks
+
+The available account data callbacks are:
+
+### `on_account_data_updated`
+
+_First introduced in Synapse v1.57.0_
+
+```python
+async def on_account_data_updated(
+    user_id: str,
+    room_id: Optional[str],
+    account_data_type: str,
+    content: "synapse.module_api.JsonDict",
+) -> None:
+```
+
+Called after user's account data has been updated. The module is given the
+Matrix ID of the user whose account data is changing, the room ID the data is associated
+with, the type associated with the change, as well as the new content. If the account
+data is not associated with a specific room, then the room ID is `None`.
+
+This callback is triggered when new account data is added or when the data associated with
+a given type (and optionally room) changes. This includes deletion, since in Matrix,
+deleting account data consists of replacing the data associated with a given type
+(and optionally room) with an empty dictionary (`{}`).
+
+Note that this doesn't trigger when changing the tags associated with a room, as these are
+processed separately by Synapse.
+
+If multiple modules implement this callback, Synapse runs them all in order.
+
+## Example
+
+The example below is a module that implements the `on_account_data_updated` callback, and
+sends an event to an audit room when a user changes their account data.
+
+```python
+import json
+import attr
+from typing import Any, Dict, Optional
+
+from synapse.module_api import JsonDict, ModuleApi
+from synapse.module_api.errors import ConfigError
+
+
+@attr.s(auto_attribs=True)
+class CustomAccountDataConfig:
+    audit_room: str
+    sender: str
+
+
+class CustomAccountDataModule:
+    def __init__(self, config: CustomAccountDataConfig, api: ModuleApi):
+        self.api = api
+        self.config = config
+
+        self.api.register_account_data_callbacks(
+            on_account_data_updated=self.log_new_account_data,
+        )
+
+    @staticmethod
+    def parse_config(config: Dict[str, Any]) -> CustomAccountDataConfig:
+        def check_in_config(param: str):
+            if param not in config:
+                raise ConfigError(f"'{param}' is required")
+
+        check_in_config("audit_room")
+        check_in_config("sender")
+
+        return CustomAccountDataConfig(
+            audit_room=config["audit_room"],
+            sender=config["sender"],
+        )
+
+    async def log_new_account_data(
+        self,
+        user_id: str,
+        room_id: Optional[str],
+        account_data_type: str,
+        content: JsonDict,
+    ) -> None:
+        content_raw = json.dumps(content)
+        msg_content = f"{user_id} has changed their account data for type {account_data_type} to: {content_raw}"
+
+        if room_id is not None:
+            msg_content += f" (in room {room_id})"
+
+        await self.api.create_and_send_event_into_room(
+            {
+                "room_id": self.config.audit_room,
+                "sender": self.config.sender,
+                "type": "m.room.message",
+                "content": {
+                    "msgtype": "m.text",
+                    "body": msg_content
+                }
+            }
+        )
+```
diff --git a/docs/modules/writing_a_module.md b/docs/modules/writing_a_module.md
index e7c0ffad58..e6303b739e 100644
--- a/docs/modules/writing_a_module.md
+++ b/docs/modules/writing_a_module.md
@@ -33,7 +33,7 @@ A module can implement the following static method:
 
 ```python
 @staticmethod
-def parse_config(config: dict) -> dict
+def parse_config(config: dict) -> Any
 ```
 
 This method is given a dictionary resulting from parsing the YAML configuration for the
diff --git a/synapse/handlers/account_data.py b/synapse/handlers/account_data.py
index 177b4f8991..4af9fbc5d1 100644
--- a/synapse/handlers/account_data.py
+++ b/synapse/handlers/account_data.py
@@ -12,8 +12,9 @@
 # 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.
+import logging
 import random
-from typing import TYPE_CHECKING, Collection, List, Optional, Tuple
+from typing import TYPE_CHECKING, Awaitable, Callable, Collection, List, Optional, Tuple
 
 from synapse.replication.http.account_data import (
     ReplicationAddTagRestServlet,
@@ -27,6 +28,12 @@ from synapse.types import JsonDict, UserID
 if TYPE_CHECKING:
     from synapse.server import HomeServer
 
+logger = logging.getLogger(__name__)
+
+ON_ACCOUNT_DATA_UPDATED_CALLBACK = Callable[
+    [str, Optional[str], str, JsonDict], Awaitable
+]
+
 
 class AccountDataHandler:
     def __init__(self, hs: "HomeServer"):
@@ -40,6 +47,44 @@ class AccountDataHandler:
         self._remove_tag_client = ReplicationRemoveTagRestServlet.make_client(hs)
         self._account_data_writers = hs.config.worker.writers.account_data
 
+        self._on_account_data_updated_callbacks: List[
+            ON_ACCOUNT_DATA_UPDATED_CALLBACK
+        ] = []
+
+    def register_module_callbacks(
+        self, on_account_data_updated: Optional[ON_ACCOUNT_DATA_UPDATED_CALLBACK] = None
+    ) -> None:
+        """Register callbacks from modules."""
+        if on_account_data_updated is not None:
+            self._on_account_data_updated_callbacks.append(on_account_data_updated)
+
+    async def _notify_modules(
+        self,
+        user_id: str,
+        room_id: Optional[str],
+        account_data_type: str,
+        content: JsonDict,
+    ) -> None:
+        """Notifies modules about new account data changes.
+
+        A change can be either a new account data type being added, or the content
+        associated with a type being changed. Account data for a given type is removed by
+        changing the associated content to an empty dictionary.
+
+        Note that this is not called when the tags associated with a room change.
+
+        Args:
+            user_id: The user whose account data is changing.
+            room_id: The ID of the room the account data change concerns, if any.
+            account_data_type: The type of the account data.
+            content: The content that is now associated with this type.
+        """
+        for callback in self._on_account_data_updated_callbacks:
+            try:
+                await callback(user_id, room_id, account_data_type, content)
+            except Exception as e:
+                logger.exception("Failed to run module callback %s: %s", callback, e)
+
     async def add_account_data_to_room(
         self, user_id: str, room_id: str, account_data_type: str, content: JsonDict
     ) -> int:
@@ -63,6 +108,8 @@ class AccountDataHandler:
                 "account_data_key", max_stream_id, users=[user_id]
             )
 
+            await self._notify_modules(user_id, room_id, account_data_type, content)
+
             return max_stream_id
         else:
             response = await self._room_data_client(
@@ -96,6 +143,9 @@ class AccountDataHandler:
             self._notifier.on_new_event(
                 "account_data_key", max_stream_id, users=[user_id]
             )
+
+            await self._notify_modules(user_id, None, account_data_type, content)
+
             return max_stream_id
         else:
             response = await self._user_data_client(
diff --git a/synapse/module_api/__init__.py b/synapse/module_api/__init__.py
index f7f95bae29..9a61593ff5 100644
--- a/synapse/module_api/__init__.py
+++ b/synapse/module_api/__init__.py
@@ -65,6 +65,7 @@ from synapse.events.third_party_rules import (
     ON_THREEPID_BIND_CALLBACK,
     ON_USER_DEACTIVATION_STATUS_CHANGED_CALLBACK,
 )
+from synapse.handlers.account_data import ON_ACCOUNT_DATA_UPDATED_CALLBACK
 from synapse.handlers.account_validity import (
     IS_USER_EXPIRED_CALLBACK,
     ON_LEGACY_ADMIN_REQUEST,
@@ -216,6 +217,7 @@ class ModuleApi:
         self._third_party_event_rules = hs.get_third_party_event_rules()
         self._password_auth_provider = hs.get_password_auth_provider()
         self._presence_router = hs.get_presence_router()
+        self._account_data_handler = hs.get_account_data_handler()
 
     #################################################################################
     # The following methods should only be called during the module's initialisation.
@@ -376,6 +378,19 @@ class ModuleApi:
                 min_batch_size=min_batch_size,
             )
 
+    def register_account_data_callbacks(
+        self,
+        *,
+        on_account_data_updated: Optional[ON_ACCOUNT_DATA_UPDATED_CALLBACK] = None,
+    ) -> None:
+        """Registers account data callbacks.
+
+        Added in Synapse 1.57.0.
+        """
+        return self._account_data_handler.register_module_callbacks(
+            on_account_data_updated=on_account_data_updated,
+        )
+
     def register_web_resource(self, path: str, resource: Resource) -> None:
         """Registers a web resource to be served at the given path.
 
diff --git a/tests/rest/client/test_account_data.py b/tests/rest/client/test_account_data.py
new file mode 100644
index 0000000000..d5b0640e7a
--- /dev/null
+++ b/tests/rest/client/test_account_data.py
@@ -0,0 +1,75 @@
+# Copyright 2022 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 unittest.mock import Mock
+
+from synapse.rest import admin
+from synapse.rest.client import account_data, login, room
+
+from tests import unittest
+from tests.test_utils import make_awaitable
+
+
+class AccountDataTestCase(unittest.HomeserverTestCase):
+    servlets = [
+        admin.register_servlets,
+        login.register_servlets,
+        room.register_servlets,
+        account_data.register_servlets,
+    ]
+
+    def test_on_account_data_updated_callback(self) -> None:
+        """Tests that the on_account_data_updated module callback is called correctly when
+        a user's account data changes.
+        """
+        mocked_callback = Mock(return_value=make_awaitable(None))
+        self.hs.get_account_data_handler()._on_account_data_updated_callbacks.append(
+            mocked_callback
+        )
+
+        user_id = self.register_user("user", "password")
+        tok = self.login("user", "password")
+        account_data_type = "org.matrix.foo"
+        account_data_content = {"bar": "baz"}
+
+        # Change the user's global account data.
+        channel = self.make_request(
+            "PUT",
+            f"/user/{user_id}/account_data/{account_data_type}",
+            account_data_content,
+            access_token=tok,
+        )
+
+        # Test that the callback is called with the user ID, the new account data, and
+        # None as the room ID.
+        self.assertEqual(channel.code, 200, channel.result)
+        mocked_callback.assert_called_once_with(
+            user_id, None, account_data_type, account_data_content
+        )
+
+        # Change the user's room-specific account data.
+        room_id = self.helper.create_room_as(user_id, tok=tok)
+        channel = self.make_request(
+            "PUT",
+            f"/user/{user_id}/rooms/{room_id}/account_data/{account_data_type}",
+            account_data_content,
+            access_token=tok,
+        )
+
+        # Test that the callback is called with the user ID, the room ID and the new
+        # account data.
+        self.assertEqual(channel.code, 200, channel.result)
+        self.assertEqual(mocked_callback.call_count, 2)
+        mocked_callback.assert_called_with(
+            user_id, room_id, account_data_type, account_data_content
+        )