Add get_media_upload_limits_for_user and on_media_upload_limit_exceeded callbacks to module API

changelog.d/18848.feature

@@ -0,0 +1 @@
+Add `get_media_upload_limits_for_user` and `on_media_upload_limit_exceeded` module API callbacks for media repository.

docs/modules/media_repository_callbacks.md

@@ -64,3 +64,55 @@ If multiple modules implement this callback, they will be considered in order. I
 returns `True`, Synapse falls through to the next one. The value of the first callback that
 returns `False` will be used. If this happens, Synapse will not call any of the subsequent
 implementations of this callback.
+
+### `get_media_upload_limits_for_user`
+
+_First introduced in Synapse v1.139.0_
+
+```python
+async def get_media_upload_limits_for_user(user_id: str, size: int) -> Optional[List[MediaUploadLimit]]
+```
+
+**<span style="color:red">
+Caution: This callback is currently experimental . The method signature or behaviour
+may change without notice.
+</span>**
+
+Called when processing a request to store content in the media repository.
+
+The arguments passed to this callback are:
+
+* `user_id`: The Matrix user ID of the user (e.g. `@alice:example.com`) making the request.
+
+If the callback returns a list then it will be used as the limits to be applied to the request.
+
+If an empty list is returned then no limits are applied.
+
+If multiple modules implement this callback, they will be considered in order. If a
+callback returns `None`, Synapse falls through to the next one. The value of the first
+callback that does not return `None` will be used. If this happens, Synapse will not call
+any of the subsequent implementations of this callback.
+
+If no module returns a non-`None` value then the default [media upload limits config](https://element-hq.github.io/synapse/latest/usage/configuration/config_documentation.html#media_upload_limits) will be used.
+
+### `on_media_upload_limit_exceeded`
+
+_First introduced in Synapse v1.139.0_
+
+```python
+async def on_media_upload_limit_exceeded(user_id: str, limit: MediaUploadLimit, sent_bytes: int, attempted_bytes: int) -> None
+```
+
+**<span style="color:red">
+Caution: This callback is currently experimental . The method signature or behaviour
+may change without notice.
+</span>**
+
+Called when a user attempts to upload media that would exceed a configured media upload limit.
+
+The arguments passed to this callback are:
+
+* `user_id`: The Matrix user ID of the user (e.g. `@alice:example.com`) making the request.
+* `limit`: The `MediaUploadLimit` that was reached.
+* `sent_bytes`: The number of bytes already sent during the period of the limit.
+* `attempted_bytes`: The number of bytes that the user attempted to send.

synapse/media/media_repository.py

@@ -179,11 +179,13 @@ def __init__(self, hs: "HomeServer"):
 
         # We get the media upload limits and sort them in descending order of
         # time period, so that we can apply some optimizations.
-        self.media_upload_limits = hs.config.media.media_upload_limits
-        self.media_upload_limits.sort(
+        self.default_media_upload_limits = hs.config.media.media_upload_limits
+        self.default_media_upload_limits.sort(
             key=lambda limit: limit.time_period_ms, reverse=True
         )
 
+        self.media_repository_callbacks = hs.get_module_api_callbacks().media_repository
+
     def _start_update_recently_accessed(self) -> Deferred:
         return run_as_background_process(
             "update_recently_accessed_media",
@@ -340,13 +342,24 @@ async def create_or_update_content(
 
         # Check that the user has not exceeded any of the media upload limits.
 
+        # Use limits from module API if provided
+        media_upload_limits = (
+            await self.media_repository_callbacks.get_media_upload_limits_for_user(
+                auth_user.to_string()
+            )
+        )
+
+        # Otherwise use the default limits from config
+        if media_upload_limits is None:
+            # Note: the media upload limits are sorted so larger time periods are
+            # first.
+            media_upload_limits = self.default_media_upload_limits
+
         # This is the total size of media uploaded by the user in the last
         # `time_period_ms` milliseconds, or None if we haven't checked yet.
         uploaded_media_size: Optional[int] = None
 
-        # Note: the media upload limits are sorted so larger time periods are
-        # first.
-        for limit in self.media_upload_limits:
+        for limit in media_upload_limits:
             # We only need to check the amount of media uploaded by the user in
             # this latest (smaller) time period if the amount of media uploaded
             # in a previous (larger) time period is above the limit.
@@ -363,6 +376,12 @@ async def create_or_update_content(
                 )
 
             if uploaded_media_size + content_length > limit.max_bytes:
+                await self.media_repository_callbacks.on_media_upload_limit_exceeded(
+                    user_id=auth_user.to_string(),
+                    limit=limit,
+                    sent_bytes=uploaded_media_size,
+                    attempted_bytes=content_length,
+                )
                 raise SynapseError(
                     400, "Media upload limit exceeded", Codes.RESOURCE_LIMIT_EXCEEDED
                 )

synapse/module_api/__init__.py

@@ -50,6 +50,7 @@
 from synapse.api.errors import SynapseError
 from synapse.api.presence import UserPresenceState
 from synapse.config import ConfigError
+from synapse.config.repository import MediaUploadLimit
 from synapse.events import EventBase
 from synapse.events.presence_router import (
     GET_INTERESTED_USERS_CALLBACK,
@@ -94,7 +95,9 @@
 )
 from synapse.module_api.callbacks.media_repository_callbacks import (
     GET_MEDIA_CONFIG_FOR_USER_CALLBACK,
+    GET_MEDIA_UPLOAD_LIMITS_FOR_USER_CALLBACK,
     IS_USER_ALLOWED_TO_UPLOAD_MEDIA_OF_SIZE_CALLBACK,
+    ON_MEDIA_UPLOAD_LIMIT_EXCEEDED_CALLBACK,
 )
 from synapse.module_api.callbacks.ratelimit_callbacks import (
     GET_RATELIMIT_OVERRIDE_FOR_USER_CALLBACK,
@@ -205,6 +208,7 @@
     "RoomAlias",
     "UserProfile",
     "RatelimitOverride",
+    "MediaUploadLimit",
 ]
 
 logger = logging.getLogger(__name__)
@@ -462,13 +466,21 @@ def register_media_repository_callbacks(
         is_user_allowed_to_upload_media_of_size: Optional[
             IS_USER_ALLOWED_TO_UPLOAD_MEDIA_OF_SIZE_CALLBACK
         ] = None,
+        get_media_upload_limits_for_user: Optional[
+            GET_MEDIA_UPLOAD_LIMITS_FOR_USER_CALLBACK
+        ] = None,
+        on_media_upload_limit_exceeded: Optional[
+            ON_MEDIA_UPLOAD_LIMIT_EXCEEDED_CALLBACK
+        ] = None,
     ) -> None:
         """Registers callbacks for media repository capabilities.
         Added in Synapse v1.132.0.
         """
         return self._callbacks.media_repository.register_callbacks(
             get_media_config_for_user=get_media_config_for_user,
             is_user_allowed_to_upload_media_of_size=is_user_allowed_to_upload_media_of_size,
+            get_media_upload_limits_for_user=get_media_upload_limits_for_user,
+            on_media_upload_limit_exceeded=on_media_upload_limit_exceeded,
         )
 
     def register_third_party_rules_callbacks(

synapse/module_api/callbacks/media_repository_callbacks.py

@@ -15,6 +15,7 @@
 import logging
 from typing import TYPE_CHECKING, Awaitable, Callable, List, Optional
 
+from synapse.config.repository import MediaUploadLimit
 from synapse.types import JsonDict
 from synapse.util.async_helpers import delay_cancellation
 from synapse.util.metrics import Measure
@@ -28,6 +29,14 @@
 
 IS_USER_ALLOWED_TO_UPLOAD_MEDIA_OF_SIZE_CALLBACK = Callable[[str, int], Awaitable[bool]]
 
+GET_MEDIA_UPLOAD_LIMITS_FOR_USER_CALLBACK = Callable[
+    [str], Awaitable[Optional[List[MediaUploadLimit]]]
+]
+
+ON_MEDIA_UPLOAD_LIMIT_EXCEEDED_CALLBACK = Callable[
+    [str, MediaUploadLimit, int, int], Awaitable[None]
+]
+
 
 class MediaRepositoryModuleApiCallbacks:
     def __init__(self, hs: "HomeServer") -> None:
@@ -39,13 +48,25 @@ def __init__(self, hs: "HomeServer") -> None:
         self._is_user_allowed_to_upload_media_of_size_callbacks: List[
             IS_USER_ALLOWED_TO_UPLOAD_MEDIA_OF_SIZE_CALLBACK
         ] = []
+        self._get_media_upload_limits_for_user_callbacks: List[
+            GET_MEDIA_UPLOAD_LIMITS_FOR_USER_CALLBACK
+        ] = []
+        self._on_media_upload_limit_exceeded_callbacks: List[
+            ON_MEDIA_UPLOAD_LIMIT_EXCEEDED_CALLBACK
+        ] = []
 
     def register_callbacks(
         self,
         get_media_config_for_user: Optional[GET_MEDIA_CONFIG_FOR_USER_CALLBACK] = None,
         is_user_allowed_to_upload_media_of_size: Optional[
             IS_USER_ALLOWED_TO_UPLOAD_MEDIA_OF_SIZE_CALLBACK
         ] = None,
+        get_media_upload_limits_for_user: Optional[
+            GET_MEDIA_UPLOAD_LIMITS_FOR_USER_CALLBACK
+        ] = None,
+        on_media_upload_limit_exceeded: Optional[
+            ON_MEDIA_UPLOAD_LIMIT_EXCEEDED_CALLBACK
+        ] = None,
     ) -> None:
         """Register callbacks from module for each hook."""
         if get_media_config_for_user is not None:
@@ -56,6 +77,16 @@ def register_callbacks(
                 is_user_allowed_to_upload_media_of_size
             )
 
+        if get_media_upload_limits_for_user is not None:
+            self._get_media_upload_limits_for_user_callbacks.append(
+                get_media_upload_limits_for_user
+            )
+
+        if on_media_upload_limit_exceeded is not None:
+            self._on_media_upload_limit_exceeded_callbacks.append(
+                on_media_upload_limit_exceeded
+            )
+
     async def get_media_config_for_user(self, user_id: str) -> Optional[JsonDict]:
         for callback in self._get_media_config_for_user_callbacks:
             with Measure(
@@ -83,3 +114,41 @@ async def is_user_allowed_to_upload_media_of_size(
                 return res
 
         return True
+
+    async def get_media_upload_limits_for_user(
+        self, user_id: str
+    ) -> Optional[List[MediaUploadLimit]]:
+        for callback in self._get_media_upload_limits_for_user_callbacks:
+            with Measure(
+                self.clock,
+                name=f"{callback.__module__}.{callback.__qualname__}",
+                server_name=self.server_name,
+            ):
+                res: Optional[List[MediaUploadLimit]] = await delay_cancellation(
+                    callback(user_id)
+                )
+            if res is not None:  # to allow [] to be returned meaning no limit
+                return res
+
+        return None
+
+    async def on_media_upload_limit_exceeded(
+        self,
+        user_id: str,
+        limit: MediaUploadLimit,
+        sent_bytes: int,
+        attempted_bytes: int,
+    ) -> None:
+        for callback in self._on_media_upload_limit_exceeded_callbacks:
+            with Measure(
+                self.clock,
+                name=f"{callback.__module__}.{callback.__qualname__}",
+                server_name=self.server_name,
+            ):
+                # Use a copy of the data in case the module modifies it
+                limit_copy = MediaUploadLimit(
+                    max_bytes=limit.max_bytes, time_period_ms=limit.time_period_ms
+                )
+                await delay_cancellation(
+                    callback(user_id, limit_copy, sent_bytes, attempted_bytes)
+                )