From 759162db2bab5fa42a608dfed922460f49a10943 Mon Sep 17 00:00:00 2001
From: Sumner Evans <sumner@beeper.com>
Date: Fri, 15 Apr 2022 09:05:11 -0600
Subject: [PATCH] media: move muxing of versioned endpoints to
 MediaRepositoryResource

This allows for better handling of all of the versioned media endpoints

Signed-off-by: Sumner Evans <sumner@beeper.com>
---
 synapse/api/urls.py                       |  4 +---
 synapse/app/generic_worker.py             |  8 ++-----
 synapse/app/homeserver.py                 | 12 ++--------
 synapse/rest/media/v1/media_repository.py | 29 +++++++++++++++++++++--
 4 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/synapse/api/urls.py b/synapse/api/urls.py
index bd49fa6a5f03..9881d736bd2a 100644
--- a/synapse/api/urls.py
+++ b/synapse/api/urls.py
@@ -29,9 +29,7 @@
 FEDERATION_UNSTABLE_PREFIX = FEDERATION_PREFIX + "/unstable"
 STATIC_PREFIX = "/_matrix/static"
 SERVER_KEY_V2_PREFIX = "/_matrix/key/v2"
-MEDIA_R0_PREFIX = "/_matrix/media/r0"
-MEDIA_V3_PREFIX = "/_matrix/media/v3"
-LEGACY_MEDIA_PREFIX = "/_matrix/media/v1"
+MEDIA_PREFIX = "/_matrix/media"
 
 
 class ConsentURIBuilder:
diff --git a/synapse/app/generic_worker.py b/synapse/app/generic_worker.py
index 0a6dd618f682..8b06e8f4bf6a 100644
--- a/synapse/app/generic_worker.py
+++ b/synapse/app/generic_worker.py
@@ -27,9 +27,7 @@
 from synapse.api.urls import (
     CLIENT_API_PREFIX,
     FEDERATION_PREFIX,
-    LEGACY_MEDIA_PREFIX,
-    MEDIA_R0_PREFIX,
-    MEDIA_V3_PREFIX,
+    MEDIA_PREFIX,
     SERVER_KEY_V2_PREFIX,
 )
 from synapse.app import _base
@@ -340,9 +338,7 @@ def _listen_http(self, listener_config: ListenerConfig) -> None:
 
                         resources.update(
                             {
-                                MEDIA_R0_PREFIX: media_repo,
-                                MEDIA_V3_PREFIX: media_repo,
-                                LEGACY_MEDIA_PREFIX: media_repo,
+                                MEDIA_PREFIX: media_repo,
                                 "/_synapse/admin": admin_resource,
                             }
                         )
diff --git a/synapse/app/homeserver.py b/synapse/app/homeserver.py
index 4c6c0658ab14..6003ddfc0626 100644
--- a/synapse/app/homeserver.py
+++ b/synapse/app/homeserver.py
@@ -30,9 +30,7 @@
 from synapse.api.urls import (
     CLIENT_API_PREFIX,
     FEDERATION_PREFIX,
-    LEGACY_MEDIA_PREFIX,
-    MEDIA_R0_PREFIX,
-    MEDIA_V3_PREFIX,
+    MEDIA_PREFIX,
     SERVER_KEY_V2_PREFIX,
     STATIC_PREFIX,
 )
@@ -245,13 +243,7 @@ def _configure_named_resource(
         if name in ["media", "federation", "client"]:
             if self.config.server.enable_media_repo:
                 media_repo = self.get_media_repository_resource()
-                resources.update(
-                    {
-                        MEDIA_R0_PREFIX: media_repo,
-                        MEDIA_V3_PREFIX: media_repo,
-                        LEGACY_MEDIA_PREFIX: media_repo,
-                    }
-                )
+                resources[MEDIA_PREFIX] = media_repo
             elif name == "media":
                 raise ConfigError(
                     "'media' resource conflicts with enable_media_repo=False"
diff --git a/synapse/rest/media/v1/media_repository.py b/synapse/rest/media/v1/media_repository.py
index 3e5d6c629418..b39616652b2e 100644
--- a/synapse/rest/media/v1/media_repository.py
+++ b/synapse/rest/media/v1/media_repository.py
@@ -16,6 +16,7 @@
 import logging
 import os
 import shutil
+from enum import Enum
 from io import BytesIO
 from typing import IO, TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 
@@ -949,7 +950,14 @@ async def _remove_local_media_from_disk(
         return removed_media, len(removed_media)
 
 
-class MediaRepositoryResource(Resource):
+class MediaVersion(Enum):
+    R0 = b"r0"
+    V3 = b"v3"
+    LEGACY = b"v1"
+    UNSTABLE = b"unstable"
+
+
+class VersionedMediaRepositoryResource(Resource):
     """File uploading and downloading.
 
     Uploads are POSTed to a resource which returns a token which is used to GET
@@ -994,7 +1002,7 @@ class MediaRepositoryResource(Resource):
     within a given rectangle.
     """
 
-    def __init__(self, hs: "HomeServer"):
+    def __init__(self, hs: "HomeServer", version: MediaVersion):
         # If we're not configured to use it, raise if we somehow got here.
         if not hs.config.media.can_load_media_repo:
             raise ConfigError("Synapse is not configured to use a media repo.")
@@ -1013,3 +1021,20 @@ def __init__(self, hs: "HomeServer"):
                 PreviewUrlResource(hs, media_repo, media_repo.media_storage),
             )
         self.putChild(b"config", MediaConfigResource(hs))
+
+
+class MediaRepositoryResource(Resource):
+    """
+    Base media repository resource. This handles all /_matrix/media requests
+    and muxes them between the versioned media repository endpoints.
+    """
+
+    def __init__(self, hs: "HomeServer"):
+        # If we're not configured to use it, raise if we somehow got here.
+        if not hs.config.media.can_load_media_repo:
+            raise ConfigError("Synapse is not configured to use a media repo.")
+
+        super().__init__()
+
+        for version in MediaVersion:
+            self.putChild(version.value, VersionedMediaRepositoryResource(hs, version))