media/{download,thumbnail}: add support for stall parameter

The stall will default to 20s if MSC2246 is not enabled.

For remote media, the max_stall_ms parameter through to remote server.

This way, if servers over federation attempt to request the media, the
request will only 404 if the media is large (or the client doesn't
upload as soon as it can).

Signed-off-by: Sumner Evans <sumner@beeper.com>

Author: sumnerevans

synapse/rest/media/v1/_base.py

@@ -49,6 +49,8 @@
     "text/xml",
 ]
 
+DEFAULT_MSC2246_DELAY = 20_000
+
 
 def parse_media_id(request: Request) -> Tuple[str, str, Optional[str]]:
     """Parses the server name, media ID and optional file name from the request URI

synapse/rest/media/v1/download_resource.py

@@ -16,10 +16,10 @@
 from typing import TYPE_CHECKING
 
 from synapse.http.server import DirectServeJsonResource, set_cors_headers
-from synapse.http.servlet import parse_boolean
+from synapse.http.servlet import parse_boolean, parse_integer
 from synapse.http.site import SynapseRequest
 
-from ._base import parse_media_id, respond_404
+from ._base import DEFAULT_MSC2246_DELAY, parse_media_id, respond_404
 
 if TYPE_CHECKING:
     from synapse.rest.media.v1.media_repository import MediaRepository
@@ -35,6 +35,7 @@ def __init__(self, hs: "HomeServer", media_repo: "MediaRepository"):
         super().__init__()
         self.media_repo = media_repo
         self.server_name = hs.hostname
+        self.enable_msc2246 = hs.config.experimental.msc2246_enabled
 
     async def _async_render_GET(self, request: SynapseRequest) -> None:
         set_cors_headers(request)
@@ -50,13 +51,14 @@ async def _async_render_GET(self, request: SynapseRequest) -> None:
         )
         # Limited non-standard form of CSP for IE11
         request.setHeader(b"X-Content-Security-Policy", b"sandbox;")
-        request.setHeader(
-            b"Referrer-Policy",
-            b"no-referrer",
-        )
+        request.setHeader(b"Referrer-Policy", b"no-referrer")
         server_name, media_id, name = parse_media_id(request)
+        max_stall_ms = parse_integer(
+            request, "fi.mau.msc2246.max_stall_ms", default=DEFAULT_MSC2246_DELAY
+        )
+
         if server_name == self.server_name:
-            await self.media_repo.get_local_media(request, media_id, name)
+            await self.media_repo.get_local_media(request, media_id, name, max_stall_ms)
         else:
             allow_remote = parse_boolean(request, "allow_remote", default=True)
             if not allow_remote:
@@ -68,4 +70,6 @@ async def _async_render_GET(self, request: SynapseRequest) -> None:
                 respond_404(request)
                 return
 
-            await self.media_repo.get_remote_media(request, server_name, media_id, name)
+            await self.media_repo.get_remote_media(
+                request, server_name, media_id, name, max_stall_ms
+            )

synapse/rest/media/v1/media_repository.py

@@ -18,7 +18,7 @@
 import shutil
 from enum import Enum
 from io import BytesIO
-from typing import IO, TYPE_CHECKING, Dict, List, Optional, Set, Tuple
+from typing import IO, TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
 
 import twisted.internet.error
 import twisted.web.http
@@ -32,13 +32,14 @@
     NotFoundError,
     RequestSendFailed,
     SynapseError,
+    cs_error,
 )
 from synapse.config._base import ConfigError
 from synapse.config.repository import ThumbnailRequirement
+from synapse.http.server import respond_with_json
 from synapse.http.site import SynapseRequest
 from synapse.logging.context import defer_to_thread
 from synapse.metrics.background_process_metrics import run_as_background_process
-from synapse.rest.media.v1.create_resource import CreateResource
 from synapse.types import UserID
 from synapse.util.async_helpers import Linearizer
 from synapse.util.retryutils import NotRetryingDestination
@@ -53,6 +54,7 @@
     respond_with_responder,
 )
 from .config_resource import MediaConfigResource
+from .create_resource import CreateResource
 from .download_resource import DownloadResource
 from .filepath import MediaFilePaths
 from .media_storage import MediaStorage
@@ -288,8 +290,64 @@ async def create_content(
 
         return "mxc://%s/%s" % (self.server_name, media_id)
 
+    def respond_not_yet_uploaded(self, request: SynapseRequest) -> None:
+        not_uploaded_error = cs_error(
+            "Media has not been uploaded yet",
+            code="FI.MAU.MSC2246_NOT_YET_UPLOADED",
+            retry_after_ms=5_000,
+        )
+        respond_with_json(request, 404, not_uploaded_error, send_cors=True)
+
+    async def get_local_media_info(
+        self, request: SynapseRequest, media_id: str, max_stall_ms: int
+    ) -> Optional[Dict[str, Any]]:
+        """Gets the info dictionary for given local media ID. If the media has
+        not been uploaded yet, this function will wait up to ``max_stall_ms``
+        milliseconds for the media to be uploaded.
+
+        Args:
+            request: The incoming request.
+            media_id: The media ID of the content. (This is the same as
+                the file_id for local content.)
+            max_stall_ms: the maximum number of milliseconds to wait for the
+                media to be uploaded.
+
+        Returns:
+            Either the info dictionary for the given local media ID or
+            ``None``. If ``None``, then no further processing is necessary as
+            this function will send the necessary JSON response.
+        """
+        wait_until = self.clock.time_msec() + max_stall_ms
+        while True:
+            # Get the info for the media
+            media_info = await self.store.get_local_media(media_id)
+            if not media_info:
+                respond_404(request)
+                return None
+
+            if media_info["quarantined_by"]:
+                logger.info("Media is quarantined")
+                respond_404(request)
+                return None
+
+            # The file has been uploaded, so stop looping
+            if media_info.get("media_length") is not None:
+                return media_info
+
+            if self.clock.time_msec() >= wait_until:
+                break
+
+            await self.clock.sleep(0.5)
+
+        self.respond_not_yet_uploaded(request)
+        return None
+
     async def get_local_media(
-        self, request: SynapseRequest, media_id: str, name: Optional[str]
+        self,
+        request: SynapseRequest,
+        media_id: str,
+        name: Optional[str],
+        max_stall_ms: int,
     ) -> None:
         """Responds to requests for local media, if exists, or returns 404.
 
@@ -299,13 +357,14 @@ async def get_local_media(
                 the file_id for local content.)
             name: Optional name that, if specified, will be used as
                 the filename in the Content-Disposition header of the response.
+            max_stall_ms: the maximum number of milliseconds to wait for the
+                media to be uploaded.
 
         Returns:
             Resolves once a response has successfully been written to request
         """
-        media_info = await self.store.get_local_media(media_id)
-        if not media_info or media_info["quarantined_by"]:
-            respond_404(request)
+        media_info = await self.get_local_media_info(request, media_id, max_stall_ms)
+        if not media_info:
             return
 
         self.mark_recently_accessed(None, media_id)
@@ -330,6 +389,7 @@ async def get_remote_media(
         server_name: str,
         media_id: str,
         name: Optional[str],
+        max_stall_ms: int,
     ) -> None:
         """Respond to requests for remote media.
 
@@ -339,6 +399,8 @@ async def get_remote_media(
             media_id: The media ID of the content (as defined by the remote server).
             name: Optional name that, if specified, will be used as
                 the filename in the Content-Disposition header of the response.
+            max_stall_ms: the maximum number of milliseconds to wait for the
+                media to be uploaded.
 
         Returns:
             Resolves once a response has successfully been written to request
@@ -353,33 +415,37 @@ async def get_remote_media(
 
         # We linearize here to ensure that we don't try and download remote
         # media multiple times concurrently
-        key = (server_name, media_id)
-        async with self.remote_media_linearizer.queue(key):
+        async with self.remote_media_linearizer.queue((server_name, media_id)):
             responder, media_info = await self._get_remote_media_impl(
-                server_name, media_id
+                server_name, media_id, max_stall_ms
             )
 
-        # We deliberately stream the file outside the lock
-        if responder:
+        if responder and media_info:
             media_type = media_info["media_type"]
             media_length = media_info["media_length"]
             upload_name = name if name else media_info["upload_name"]
             await respond_with_responder(
                 request, responder, media_type, media_length, upload_name
             )
         else:
-            respond_404(request)
+            self.respond_not_yet_uploaded(request)
+            return
 
-    async def get_remote_media_info(self, server_name: str, media_id: str) -> dict:
+    async def get_remote_media_info(
+        self, server_name: str, media_id: str, max_stall_ms: int
+    ) -> Optional[dict]:
         """Gets the media info associated with the remote file, downloading
         if necessary.
 
         Args:
             server_name: Remote server_name where the media originated.
             media_id: The media ID of the content (as defined by the remote server).
+            max_stall_ms: the maximum number of milliseconds to wait for the
+                media to be uploaded.
 
         Returns:
-            The media info of the file
+            The media info of the file or ``None`` if the media wasn't uploaded
+            in time.
         """
         if (
             self.federation_domain_whitelist is not None
@@ -389,10 +455,9 @@ async def get_remote_media_info(self, server_name: str, media_id: str) -> dict:
 
         # We linearize here to ensure that we don't try and download remote
         # media multiple times concurrently
-        key = (server_name, media_id)
-        async with self.remote_media_linearizer.queue(key):
+        async with self.remote_media_linearizer.queue((server_name, media_id)):
             responder, media_info = await self._get_remote_media_impl(
-                server_name, media_id
+                server_name, media_id, max_stall_ms
             )
 
         # Ensure we actually use the responder so that it releases resources
@@ -403,7 +468,7 @@ async def get_remote_media_info(self, server_name: str, media_id: str) -> dict:
         return media_info
 
     async def _get_remote_media_impl(
-        self, server_name: str, media_id: str
+        self, server_name: str, media_id: str, max_stall_ms: int
     ) -> Tuple[Optional[Responder], dict]:
         """Looks for media in local cache, if not there then attempt to
         download from remote server.
@@ -412,6 +477,8 @@ async def _get_remote_media_impl(
             server_name (str): Remote server_name where the media originated.
             media_id (str): The media ID of the content (as defined by the
                 remote server).
+            max_stall_ms: the maximum number of milliseconds to wait for the
+                media to be uploaded.
 
         Returns:
             A tuple of responder and the media info of the file.
@@ -442,8 +509,7 @@ async def _get_remote_media_impl(
 
         try:
             media_info = await self._download_remote_file(
-                server_name,
-                media_id,
+                server_name, media_id, max_stall_ms
             )
         except SynapseError:
             raise
@@ -476,6 +542,7 @@ async def _download_remote_file(
         self,
         server_name: str,
         media_id: str,
+        max_stall_ms: int,
     ) -> dict:
         """Attempt to download the remote file from the given server name,
         using the given file_id as the local id.
@@ -485,7 +552,8 @@ async def _download_remote_file(
             media_id: The media ID of the content (as defined by the
                 remote server). This is different than the file_id, which is
                 locally generated.
-            file_id: Local file ID
+            max_stall_ms: the maximum number of milliseconds to wait for the
+                media to be uploaded.
 
         Returns:
             The media info of the file.
@@ -509,7 +577,8 @@ async def _download_remote_file(
                         # tell the remote server to 404 if it doesn't
                         # recognise the server_name, to make sure we don't
                         # end up with a routing loop.
-                        "allow_remote": "false"
+                        "allow_remote": "false",
+                        "fi.mau.msc2246.max_stall_ms": str(max_stall_ms),
                     },
                 )
             except RequestSendFailed as e: