Re-introduce federation /download endpoint (#17350)

1 files changed, 77 insertions, 1 deletions

diff --git a/synapse/media/_base.py b/synapse/media/_base.py
index 3fbed6062f..7ad0b7c3cf 100644
--- a/synapse/media/_base.py
+++ b/synapse/media/_base.py
@@ -25,7 +25,16 @@ import os
 import urllib
 from abc import ABC, abstractmethod
 from types import TracebackType
-from typing import Awaitable, Dict, Generator, List, Optional, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Awaitable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Tuple,
+    Type,
+)
 
 import attr
 
@@ -37,8 +46,13 @@ from synapse.api.errors import Codes, cs_error
 from synapse.http.server import finish_request, respond_with_json
 from synapse.http.site import SynapseRequest
 from synapse.logging.context import make_deferred_yieldable
+from synapse.util import Clock
 from synapse.util.stringutils import is_ascii
 
+if TYPE_CHECKING:
+    from synapse.storage.databases.main.media_repository import LocalMedia
+
+
 logger = logging.getLogger(__name__)
 
 # list all text content types that will have the charset default to UTF-8 when
@@ -260,6 +274,68 @@ def _can_encode_filename_as_token(x: str) -> bool:
     return True
 
 
+async def respond_with_multipart_responder(
+    clock: Clock,
+    request: SynapseRequest,
+    responder: "Optional[Responder]",
+    media_info: "LocalMedia",
+) -> None:
+    """
+    Responds to requests originating from the federation media `/download` endpoint by
+    streaming a multipart/mixed response
+
+    Args:
+        clock:
+        request: the federation request to respond to
+        responder: the responder which will send the response
+        media_info: metadata about the media item
+    """
+    if not responder:
+        respond_404(request)
+        return
+
+    # If we have a responder we *must* use it as a context manager.
+    with responder:
+        if request._disconnected:
+            logger.warning(
+                "Not sending response to request %s, already disconnected.", request
+            )
+            return
+
+        from synapse.media.media_storage import MultipartFileConsumer
+
+        # note that currently the json_object is just {}, this will change when linked media
+        # is implemented
+        multipart_consumer = MultipartFileConsumer(
+            clock, request, media_info.media_type, {}, media_info.media_length
+        )
+
+        logger.debug("Responding to media request with responder %s", responder)
+        if media_info.media_length is not None:
+            content_length = multipart_consumer.content_length()
+            assert content_length is not None
+            request.setHeader(b"Content-Length", b"%d" % (content_length,))
+
+        request.setHeader(
+            b"Content-Type",
+            b"multipart/mixed; boundary=%s" % multipart_consumer.boundary,
+        )
+
+        try:
+            await responder.write_to_consumer(multipart_consumer)
+        except Exception as e:
+            # The majority of the time this will be due to the client having gone
+            # away. Unfortunately, Twisted simply throws a generic exception at us
+            # in that case.
+            logger.warning("Failed to write to consumer: %s %s", type(e), e)
+
+            # Unregister the producer, if it has one, so Twisted doesn't complain
+            if request.producer:
+                request.unregisterProducer()
+
+    finish_request(request)
+
+
 async def respond_with_responder(
     request: SynapseRequest,
     responder: "Optional[Responder]",