Start handlers for new media endpoints when media resource configured (#17483)

This is in response to issue #17473. 
Not all the necessary handlers to deal with media requests are started
now when configuring synapse to use a media worker as per the [example
config](https://element-hq.github.io/synapse/latest/workers.html#synapseappmedia_repository).
The new media endpoints introduced with authenticated media fall under
the `client` & `federation` handlers in synapse.
This PR starts up handlers for the new media endpoints if a worker has
been configured with only the `media` resource type.

### Pull Request Checklist

<!-- Please read
https://element-hq.github.io/synapse/latest/development/contributing_guide.html
before submitting your pull request -->

* [X] Pull request is based on the develop branch
* [x] Pull request includes a [changelog
file](https://element-hq.github.io/synapse/latest/development/contributing_guide.html#changelog).
The entry should:
- Be a short description of your change which makes sense to users.
"Fixed a bug that prevented receiving messages from other servers."
instead of "Moved X method from `EventStore` to `EventWorkerStore`.".
  - Use markdown where necessary, mostly for `code blocks`.
  - End with either a period (.) or an exclamation mark (!).
  - Start with a capital letter.
- Feel free to credit yourself, by adding a sentence "Contributed by
@github_username." or "Contributed by [Your Name]." to the end of the
entry.
* [X] [Code
style](https://element-hq.github.io/synapse/latest/code_style.html) is
correct
(run the
[linters](https://element-hq.github.io/synapse/latest/development/contributing_guide.html#run-the-linters))

---------

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

5 files changed, 131 insertions, 70 deletions

diff --git a/synapse/app/generic_worker.py b/synapse/app/generic_worker.py
index 248622fa92..53f1859256 100644
--- a/synapse/app/generic_worker.py
+++ b/synapse/app/generic_worker.py
@@ -206,6 +206,21 @@ class GenericWorkerServer(HomeServer):
                                 "/_synapse/admin": admin_resource,
                             }
                         )
+
+                        if "federation" not in res.names:
+                            # Only load the federation media resource separately if federation
+                            # resource is not specified since federation resource includes media
+                            # resource.
+                            resources[FEDERATION_PREFIX] = TransportLayerServer(
+                                self, servlet_groups=["media"]
+                            )
+                        if "client" not in res.names:
+                            # Only load the client media resource separately if client
+                            # resource is not specified since client resource includes media
+                            # resource.
+                            resources[CLIENT_API_PREFIX] = ClientRestResource(
+                                self, servlet_groups=["media"]
+                            )
                     else:
                         logger.warning(
                             "A 'media' listener is configured but the media"
diff --git a/synapse/app/homeserver.py b/synapse/app/homeserver.py
index e114ab7ec4..2a824e8457 100644
--- a/synapse/app/homeserver.py
+++ b/synapse/app/homeserver.py
@@ -101,6 +101,12 @@ class SynapseHomeServer(HomeServer):
                     # Skip loading openid resource if federation is defined
                     # since federation resource will include openid
                     continue
+                if name == "media" and (
+                    "federation" in res.names or "client" in res.names
+                ):
+                    # Skip loading media resource if federation or client are defined
+                    # since federation & client resources will include media
+                    continue
                 if name == "health":
                     # Skip loading, health resource is always included
                     continue
@@ -231,6 +237,14 @@ class SynapseHomeServer(HomeServer):
                     "'media' resource conflicts with enable_media_repo=False"
                 )
 
+        if name == "media":
+            resources[FEDERATION_PREFIX] = TransportLayerServer(
+                self, servlet_groups=["media"]
+            )
+            resources[CLIENT_API_PREFIX] = ClientRestResource(
+                self, servlet_groups=["media"]
+            )
+
         if name in ["keys", "federation"]:
             resources[SERVER_KEY_PREFIX] = KeyResource(self)
 
diff --git a/synapse/federation/transport/server/__init__.py b/synapse/federation/transport/server/__init__.py
index 72599bb204..43102567db 100644
--- a/synapse/federation/transport/server/__init__.py
+++ b/synapse/federation/transport/server/__init__.py
@@ -271,6 +271,10 @@ SERVLET_GROUPS: Dict[str, Iterable[Type[BaseFederationServlet]]] = {
     "federation": FEDERATION_SERVLET_CLASSES,
     "room_list": (PublicRoomList,),
     "openid": (OpenIdUserInfo,),
+    "media": (
+        FederationMediaDownloadServlet,
+        FederationMediaThumbnailServlet,
+    ),
 }
 
 
diff --git a/synapse/federation/transport/server/federation.py b/synapse/federation/transport/server/federation.py
index b075a86f68..20f87c885e 100644
--- a/synapse/federation/transport/server/federation.py
+++ b/synapse/federation/transport/server/federation.py
@@ -912,6 +912,4 @@ FEDERATION_SERVLET_CLASSES: Tuple[Type[BaseFederationServlet], ...] = (
     FederationV1SendKnockServlet,
     FederationMakeKnockServlet,
     FederationAccountStatusServlet,
-    FederationMediaDownloadServlet,
-    FederationMediaThumbnailServlet,
 )
diff --git a/synapse/rest/__init__.py b/synapse/rest/__init__.py
index 1aa9ea3877..c5cdc36955 100644
--- a/synapse/rest/__init__.py
+++ b/synapse/rest/__init__.py
@@ -18,7 +18,8 @@
 # [This file includes modifications made by New Vector Limited]
 #
 #
-from typing import TYPE_CHECKING, Callable
+import logging
+from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional, Tuple
 
 from synapse.http.server import HttpServer, JsonResource
 from synapse.rest import admin
@@ -67,11 +68,64 @@ from synapse.rest.client import (
     voip,
 )
 
+logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from synapse.server import HomeServer
 
 RegisterServletsFunc = Callable[["HomeServer", HttpServer], None]
 
+CLIENT_SERVLET_FUNCTIONS: Tuple[RegisterServletsFunc, ...] = (
+    versions.register_servlets,
+    initial_sync.register_servlets,
+    room.register_deprecated_servlets,
+    events.register_servlets,
+    room.register_servlets,
+    login.register_servlets,
+    profile.register_servlets,
+    presence.register_servlets,
+    directory.register_servlets,
+    voip.register_servlets,
+    pusher.register_servlets,
+    push_rule.register_servlets,
+    logout.register_servlets,
+    sync.register_servlets,
+    filter.register_servlets,
+    account.register_servlets,
+    register.register_servlets,
+    auth.register_servlets,
+    receipts.register_servlets,
+    read_marker.register_servlets,
+    room_keys.register_servlets,
+    keys.register_servlets,
+    tokenrefresh.register_servlets,
+    tags.register_servlets,
+    account_data.register_servlets,
+    reporting.register_servlets,
+    openid.register_servlets,
+    notifications.register_servlets,
+    devices.register_servlets,
+    thirdparty.register_servlets,
+    sendtodevice.register_servlets,
+    user_directory.register_servlets,
+    room_upgrade_rest_servlet.register_servlets,
+    capabilities.register_servlets,
+    account_validity.register_servlets,
+    relations.register_servlets,
+    password_policy.register_servlets,
+    knock.register_servlets,
+    appservice_ping.register_servlets,
+    admin.register_servlets_for_client_rest_resource,
+    mutual_rooms.register_servlets,
+    login_token_request.register_servlets,
+    rendezvous.register_servlets,
+    auth_issuer.register_servlets,
+)
+
+SERVLET_GROUPS: Dict[str, Iterable[RegisterServletsFunc]] = {
+    "client": CLIENT_SERVLET_FUNCTIONS,
+}
+
 
 class ClientRestResource(JsonResource):
     """Matrix Client API REST resource.
@@ -83,80 +137,56 @@ class ClientRestResource(JsonResource):
        * etc
     """
 
-    def __init__(self, hs: "HomeServer"):
+    def __init__(self, hs: "HomeServer", servlet_groups: Optional[List[str]] = None):
         JsonResource.__init__(self, hs, canonical_json=False)
-        self.register_servlets(self, hs)
+        if hs.config.media.can_load_media_repo:
+            # This import is here to prevent a circular import failure
+            from synapse.rest.client import media
+
+            SERVLET_GROUPS["media"] = (media.register_servlets,)
+        self.register_servlets(self, hs, servlet_groups)
 
     @staticmethod
-    def register_servlets(client_resource: HttpServer, hs: "HomeServer") -> None:
+    def register_servlets(
+        client_resource: HttpServer,
+        hs: "HomeServer",
+        servlet_groups: Optional[Iterable[str]] = None,
+    ) -> None:
         # Some servlets are only registered on the main process (and not worker
         # processes).
         is_main_process = hs.config.worker.worker_app is None
 
-        versions.register_servlets(hs, client_resource)
-
-        # Deprecated in r0
-        initial_sync.register_servlets(hs, client_resource)
-        room.register_deprecated_servlets(hs, client_resource)
-
-        # Partially deprecated in r0
-        events.register_servlets(hs, client_resource)
-
-        room.register_servlets(hs, client_resource)
-        login.register_servlets(hs, client_resource)
-        profile.register_servlets(hs, client_resource)
-        presence.register_servlets(hs, client_resource)
-        directory.register_servlets(hs, client_resource)
-        voip.register_servlets(hs, client_resource)
-        if is_main_process:
-            pusher.register_servlets(hs, client_resource)
-        push_rule.register_servlets(hs, client_resource)
-        if is_main_process:
-            logout.register_servlets(hs, client_resource)
-        sync.register_servlets(hs, client_resource)
-        filter.register_servlets(hs, client_resource)
-        account.register_servlets(hs, client_resource)
-        register.register_servlets(hs, client_resource)
-        if is_main_process:
-            auth.register_servlets(hs, client_resource)
-        receipts.register_servlets(hs, client_resource)
-        read_marker.register_servlets(hs, client_resource)
-        room_keys.register_servlets(hs, client_resource)
-        keys.register_servlets(hs, client_resource)
-        if is_main_process:
-            tokenrefresh.register_servlets(hs, client_resource)
-        tags.register_servlets(hs, client_resource)
-        account_data.register_servlets(hs, client_resource)
-        if is_main_process:
-            reporting.register_servlets(hs, client_resource)
-            openid.register_servlets(hs, client_resource)
-        notifications.register_servlets(hs, client_resource)
-        devices.register_servlets(hs, client_resource)
-        if is_main_process:
-            thirdparty.register_servlets(hs, client_resource)
-        sendtodevice.register_servlets(hs, client_resource)
-        user_directory.register_servlets(hs, client_resource)
-        if is_main_process:
-            room_upgrade_rest_servlet.register_servlets(hs, client_resource)
-        capabilities.register_servlets(hs, client_resource)
-        if is_main_process:
-            account_validity.register_servlets(hs, client_resource)
-        relations.register_servlets(hs, client_resource)
-        password_policy.register_servlets(hs, client_resource)
-        knock.register_servlets(hs, client_resource)
-        appservice_ping.register_servlets(hs, client_resource)
-        if hs.config.media.can_load_media_repo:
-            from synapse.rest.client import media
+        if not servlet_groups:
+            servlet_groups = SERVLET_GROUPS.keys()
 
-            media.register_servlets(hs, client_resource)
+        for servlet_group in servlet_groups:
+            # Fail on unknown servlet groups.
+            if servlet_group not in SERVLET_GROUPS:
+                if servlet_group == "media":
+                    logger.warn(
+                        "media.can_load_media_repo needs to be configured for the media servlet to be available"
+                    )
+                raise RuntimeError(
+                    f"Attempting to register unknown client servlet: '{servlet_group}'"
+                )
 
-        # moving to /_synapse/admin
-        if is_main_process:
-            admin.register_servlets_for_client_rest_resource(hs, client_resource)
+            for servletfunc in SERVLET_GROUPS[servlet_group]:
+                if not is_main_process and servletfunc in [
+                    pusher.register_servlets,
+                    logout.register_servlets,
+                    auth.register_servlets,
+                    tokenrefresh.register_servlets,
+                    reporting.register_servlets,
+                    openid.register_servlets,
+                    thirdparty.register_servlets,
+                    room_upgrade_rest_servlet.register_servlets,
+                    account_validity.register_servlets,
+                    admin.register_servlets_for_client_rest_resource,
+                    mutual_rooms.register_servlets,
+                    login_token_request.register_servlets,
+                    rendezvous.register_servlets,
+                    auth_issuer.register_servlets,
+                ]:
+                    continue
 
-        # unstable
-        if is_main_process:
-            mutual_rooms.register_servlets(hs, client_resource)
-            login_token_request.register_servlets(hs, client_resource)
-            rendezvous.register_servlets(hs, client_resource)
-            auth_issuer.register_servlets(hs, client_resource)
+                servletfunc(hs, client_resource)