Split out BaseMediaResource into MediaRepository

This is so that a single MediaRepository can be shared across all
resources, rather than having a "copy" per resource.

In particular this allows us to guard against both the thumbnail and
download resource triggering a download of remote content at the same
time.

6 files changed, 180 insertions, 138 deletions

diff --git a/synapse/rest/media/v1/base_resource.py b/synapse/rest/media/v1/base_resource.py
index 2b1938dc8e..ade4e28034 100644
--- a/synapse/rest/media/v1/base_resource.py
+++ b/synapse/rest/media/v1/base_resource.py
@@ -23,7 +23,6 @@ from synapse.api.errors import (
 )
 
 from twisted.internet import defer, threads
-from twisted.web.resource import Resource
 from twisted.protocols.basic import FileSender
 
 from synapse.util.async import ObservableDeferred
@@ -60,11 +59,66 @@ def parse_media_id(request):
         )
 
 
-class BaseMediaResource(Resource):
-    isLeaf = True
+def respond_404(request):
+    respond_with_json(
+        request, 404,
+        cs_error(
+            "Not found %r" % (request.postpath,),
+            code=Codes.NOT_FOUND,
+        ),
+        send_cors=True
+    )
+
+
+@defer.inlineCallbacks
+def respond_with_file(request, media_type, file_path,
+                      file_size=None, upload_name=None):
+    logger.debug("Responding with %r", file_path)
+
+    if os.path.isfile(file_path):
+        request.setHeader(b"Content-Type", media_type.encode("UTF-8"))
+        if upload_name:
+            if is_ascii(upload_name):
+                request.setHeader(
+                    b"Content-Disposition",
+                    b"inline; filename=%s" % (
+                        urllib.quote(upload_name.encode("utf-8")),
+                    ),
+                )
+            else:
+                request.setHeader(
+                    b"Content-Disposition",
+                    b"inline; filename*=utf-8''%s" % (
+                        urllib.quote(upload_name.encode("utf-8")),
+                    ),
+                )
+
+        # cache for at least a day.
+        # XXX: we might want to turn this off for data we don't want to
+        # recommend caching as it's sensitive or private - or at least
+        # select private. don't bother setting Expires as all our
+        # clients are smart enough to be happy with Cache-Control
+        request.setHeader(
+            b"Cache-Control", b"public,max-age=86400,s-maxage=86400"
+        )
+        if file_size is None:
+            stat = os.stat(file_path)
+            file_size = stat.st_size
+
+        request.setHeader(
+            b"Content-Length", b"%d" % (file_size,)
+        )
+
+        with open(file_path, "rb") as f:
+            yield FileSender().beginFileTransfer(f, request)
 
+        finish_request(request)
+    else:
+        respond_404(request)
+
+
+class MediaRepository(object):
     def __init__(self, hs, filepaths):
-        Resource.__init__(self)
         self.auth = hs.get_auth()
         self.client = MatrixFederationHttpClient(hs)
         self.clock = hs.get_clock()
@@ -72,30 +126,48 @@ class BaseMediaResource(Resource):
         self.store = hs.get_datastore()
         self.max_upload_size = hs.config.max_upload_size
         self.max_image_pixels = hs.config.max_image_pixels
-        self.max_spider_size = hs.config.max_spider_size
         self.filepaths = filepaths
-        self.version_string = hs.version_string
         self.downloads = {}
         self.dynamic_thumbnails = hs.config.dynamic_thumbnails
         self.thumbnail_requirements = hs.config.thumbnail_requirements
 
-    def _respond_404(self, request):
-        respond_with_json(
-            request, 404,
-            cs_error(
-                "Not found %r" % (request.postpath,),
-                code=Codes.NOT_FOUND,
-            ),
-            send_cors=True
-        )
-
     @staticmethod
     def _makedirs(filepath):
         dirname = os.path.dirname(filepath)
         if not os.path.exists(dirname):
             os.makedirs(dirname)
 
-    def _get_remote_media(self, server_name, media_id):
+    @defer.inlineCallbacks
+    def create_content(self, media_type, upload_name, content, content_length,
+                       auth_user):
+        media_id = random_string(24)
+
+        fname = self.filepaths.local_media_filepath(media_id)
+        self._makedirs(fname)
+
+        # This shouldn't block for very long because the content will have
+        # already been uploaded at this point.
+        with open(fname, "wb") as f:
+            f.write(content)
+
+        yield self.store.store_local_media(
+            media_id=media_id,
+            media_type=media_type,
+            time_now_ms=self.clock.time_msec(),
+            upload_name=upload_name,
+            media_length=content_length,
+            user_id=auth_user,
+        )
+        media_info = {
+            "media_type": media_type,
+            "media_length": content_length,
+        }
+
+        yield self._generate_local_thumbnails(media_id, media_info)
+
+        defer.returnValue("mxc://%s/%s" % (self.server_name, media_id))
+
+    def get_remote_media(self, server_name, media_id):
         key = (server_name, media_id)
         download = self.downloads.get(key)
         if download is None:
@@ -197,52 +269,6 @@ class BaseMediaResource(Resource):
 
         defer.returnValue(media_info)
 
-    @defer.inlineCallbacks
-    def _respond_with_file(self, request, media_type, file_path,
-                           file_size=None, upload_name=None):
-        logger.debug("Responding with %r", file_path)
-
-        if os.path.isfile(file_path):
-            request.setHeader(b"Content-Type", media_type.encode("UTF-8"))
-            if upload_name:
-                if is_ascii(upload_name):
-                    request.setHeader(
-                        b"Content-Disposition",
-                        b"inline; filename=%s" % (
-                            urllib.quote(upload_name.encode("utf-8")),
-                        ),
-                    )
-                else:
-                    request.setHeader(
-                        b"Content-Disposition",
-                        b"inline; filename*=utf-8''%s" % (
-                            urllib.quote(upload_name.encode("utf-8")),
-                        ),
-                    )
-
-            # cache for at least a day.
-            # XXX: we might want to turn this off for data we don't want to
-            # recommend caching as it's sensitive or private - or at least
-            # select private. don't bother setting Expires as all our
-            # clients are smart enough to be happy with Cache-Control
-            request.setHeader(
-                b"Cache-Control", b"public,max-age=86400,s-maxage=86400"
-            )
-            if file_size is None:
-                stat = os.stat(file_path)
-                file_size = stat.st_size
-
-            request.setHeader(
-                b"Content-Length", b"%d" % (file_size,)
-            )
-
-            with open(file_path, "rb") as f:
-                yield FileSender().beginFileTransfer(f, request)
-
-            finish_request(request)
-        else:
-            self._respond_404(request)
-
     def _get_thumbnail_requirements(self, media_type):
         return self.thumbnail_requirements.get(media_type, ())
 
@@ -269,8 +295,8 @@ class BaseMediaResource(Resource):
         return t_len
 
     @defer.inlineCallbacks
-    def _generate_local_exact_thumbnail(self, media_id, t_width, t_height,
-                                        t_method, t_type):
+    def generate_local_exact_thumbnail(self, media_id, t_width, t_height,
+                                       t_method, t_type):
         input_path = self.filepaths.local_media_filepath(media_id)
 
         t_path = self.filepaths.local_media_thumbnail(
@@ -292,8 +318,8 @@ class BaseMediaResource(Resource):
             defer.returnValue(t_path)
 
     @defer.inlineCallbacks
-    def _generate_remote_exact_thumbnail(self, server_name, file_id, media_id,
-                                         t_width, t_height, t_method, t_type):
+    def generate_remote_exact_thumbnail(self, server_name, file_id, media_id,
+                                        t_width, t_height, t_method, t_type):
         input_path = self.filepaths.remote_media_filepath(server_name, file_id)
 
         t_path = self.filepaths.remote_media_thumbnail(
diff --git a/synapse/rest/media/v1/download_resource.py b/synapse/rest/media/v1/download_resource.py
index 1aad6b3551..97f4e9b54b 100644
--- a/synapse/rest/media/v1/download_resource.py
+++ b/synapse/rest/media/v1/download_resource.py
@@ -13,7 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base_resource import BaseMediaResource, parse_media_id
+from .base_resource import parse_media_id, respond_with_file, respond_404
+from twisted.web.resource import Resource
 from synapse.http.server import request_handler
 
 from twisted.web.server import NOT_DONE_YET
@@ -24,7 +25,18 @@ import logging
 logger = logging.getLogger(__name__)
 
 
-class DownloadResource(BaseMediaResource):
+class DownloadResource(Resource):
+    isLeaf = True
+
+    def __init__(self, hs, media_repo):
+        Resource.__init__(self)
+
+        self.filepaths = media_repo.filepaths
+        self.media_repo = media_repo
+        self.server_name = hs.hostname
+        self.store = hs.get_datastore()
+        self.version_string = hs.version_string
+
     def render_GET(self, request):
         self._async_render_GET(request)
         return NOT_DONE_YET
@@ -44,7 +56,7 @@ class DownloadResource(BaseMediaResource):
     def _respond_local_file(self, request, media_id, name):
         media_info = yield self.store.get_local_media(media_id)
         if not media_info:
-            self._respond_404(request)
+            respond_404(request)
             return
 
         media_type = media_info["media_type"]
@@ -52,14 +64,14 @@ class DownloadResource(BaseMediaResource):
         upload_name = name if name else media_info["upload_name"]
         file_path = self.filepaths.local_media_filepath(media_id)
 
-        yield self._respond_with_file(
+        yield respond_with_file(
             request, media_type, file_path, media_length,
             upload_name=upload_name,
         )
 
     @defer.inlineCallbacks
     def _respond_remote_file(self, request, server_name, media_id, name):
-        media_info = yield self._get_remote_media(server_name, media_id)
+        media_info = yield self.media_repo.get_remote_media(server_name, media_id)
 
         media_type = media_info["media_type"]
         media_length = media_info["media_length"]
@@ -70,7 +82,7 @@ class DownloadResource(BaseMediaResource):
             server_name, filesystem_id
         )
 
-        yield self._respond_with_file(
+        yield respond_with_file(
             request, media_type, file_path, media_length,
             upload_name=upload_name,
         )
diff --git a/synapse/rest/media/v1/media_repository.py b/synapse/rest/media/v1/media_repository.py
index 77fb0313c5..e8fe3302b2 100644
--- a/synapse/rest/media/v1/media_repository.py
+++ b/synapse/rest/media/v1/media_repository.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .base_resource import MediaRepository
 from .upload_resource import UploadResource
 from .download_resource import DownloadResource
 from .thumbnail_resource import ThumbnailResource
@@ -75,9 +76,12 @@ class MediaRepositoryResource(Resource):
     def __init__(self, hs):
         Resource.__init__(self)
         filepaths = MediaFilePaths(hs.config.media_store_path)
-        self.putChild("upload", UploadResource(hs, filepaths))
-        self.putChild("download", DownloadResource(hs, filepaths))
-        self.putChild("thumbnail", ThumbnailResource(hs, filepaths))
+
+        media_repo = MediaRepository(hs, filepaths)
+
+        self.putChild("upload", UploadResource(hs, media_repo))
+        self.putChild("download", DownloadResource(hs, media_repo))
+        self.putChild("thumbnail", ThumbnailResource(hs, media_repo))
         self.putChild("identicon", IdenticonResource())
         if hs.config.url_preview_enabled:
-            self.putChild("preview_url", PreviewUrlResource(hs, filepaths))
+            self.putChild("preview_url", PreviewUrlResource(hs, media_repo))
diff --git a/synapse/rest/media/v1/preview_url_resource.py b/synapse/rest/media/v1/preview_url_resource.py
index 9bb7c72cfc..fecdf8ed86 100644
--- a/synapse/rest/media/v1/preview_url_resource.py
+++ b/synapse/rest/media/v1/preview_url_resource.py
@@ -13,10 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base_resource import BaseMediaResource
-
 from twisted.web.server import NOT_DONE_YET
 from twisted.internet import defer
+from twisted.web.resource import Resource
 
 from synapse.api.errors import (
     SynapseError, Codes,
@@ -41,11 +40,11 @@ import logging
 logger = logging.getLogger(__name__)
 
 
-class PreviewUrlResource(BaseMediaResource):
+class PreviewUrlResource(Resource):
     isLeaf = True
 
-    def __init__(self, hs, filepaths):
-        BaseMediaResource.__init__(self, hs, filepaths)
+    def __init__(self, hs, media_repo):
+        Resource.__init__(self)
         self.client = SpiderHttpClient(hs)
         if hasattr(hs.config, "url_preview_url_blacklist"):
             self.url_preview_url_blacklist = hs.config.url_preview_url_blacklist
@@ -61,6 +60,13 @@ class PreviewUrlResource(BaseMediaResource):
 
         self.downloads = {}
 
+        self.auth = hs.get_auth()
+        self.clock = hs.get_clock()
+        self.version_string = hs.version_string
+        self.filepaths = media_repo.filepaths
+        self.max_spider_size = hs.config.max_spider_size
+        self.server_name = hs.hostname
+
     def render_GET(self, request):
         self._async_render_GET(request)
         return NOT_DONE_YET
@@ -156,7 +162,7 @@ class PreviewUrlResource(BaseMediaResource):
         logger.debug("got media_info of '%s'" % media_info)
 
         if self._is_media(media_info['media_type']):
-            dims = yield self._generate_local_thumbnails(
+            dims = yield self.media_repo._generate_local_thumbnails(
                 media_info['filesystem_id'], media_info
             )
 
@@ -291,7 +297,7 @@ class PreviewUrlResource(BaseMediaResource):
 
             if self._is_media(image_info['media_type']):
                 # TODO: make sure we don't choke on white-on-transparent images
-                dims = yield self._generate_local_thumbnails(
+                dims = yield self.media_repo._generate_local_thumbnails(
                     image_info['filesystem_id'], image_info
                 )
                 if dims:
diff --git a/synapse/rest/media/v1/thumbnail_resource.py b/synapse/rest/media/v1/thumbnail_resource.py
index 40ef22459c..43c568b769 100644
--- a/synapse/rest/media/v1/thumbnail_resource.py
+++ b/synapse/rest/media/v1/thumbnail_resource.py
@@ -14,7 +14,8 @@
 # limitations under the License.
 
 
-from .base_resource import BaseMediaResource, parse_media_id
+from .base_resource import parse_media_id, respond_404, respond_with_file
+from twisted.web.resource import Resource
 from synapse.http.servlet import parse_string, parse_integer
 from synapse.http.server import request_handler
 
@@ -26,9 +27,19 @@ import logging
 logger = logging.getLogger(__name__)
 
 
-class ThumbnailResource(BaseMediaResource):
+class ThumbnailResource(Resource):
     isLeaf = True
 
+    def __init__(self, hs, media_repo):
+        Resource.__init__(self)
+
+        self.store = hs.get_datastore()
+        self.filepaths = media_repo.filepaths
+        self.media_repo = media_repo
+        self.dynamic_thumbnails = hs.config.dynamic_thumbnails
+        self.server_name = hs.hostname
+        self.version_string = hs.version_string
+
     def render_GET(self, request):
         self._async_render_GET(request)
         return NOT_DONE_YET
@@ -69,12 +80,12 @@ class ThumbnailResource(BaseMediaResource):
         media_info = yield self.store.get_local_media(media_id)
 
         if not media_info:
-            self._respond_404(request)
+            respond_404(request)
             return
 
         # if media_info["media_type"] == "image/svg+xml":
         #     file_path = self.filepaths.local_media_filepath(media_id)
-        #     yield self._respond_with_file(request, media_info["media_type"], file_path)
+        #     yield respond_with_file(request, media_info["media_type"], file_path)
         #     return
 
         thumbnail_infos = yield self.store.get_local_media_thumbnails(media_id)
@@ -91,7 +102,7 @@ class ThumbnailResource(BaseMediaResource):
             file_path = self.filepaths.local_media_thumbnail(
                 media_id, t_width, t_height, t_type, t_method,
             )
-            yield self._respond_with_file(request, t_type, file_path)
+            yield respond_with_file(request, t_type, file_path)
 
         else:
             yield self._respond_default_thumbnail(
@@ -105,12 +116,12 @@ class ThumbnailResource(BaseMediaResource):
         media_info = yield self.store.get_local_media(media_id)
 
         if not media_info:
-            self._respond_404(request)
+            respond_404(request)
             return
 
         # if media_info["media_type"] == "image/svg+xml":
         #     file_path = self.filepaths.local_media_filepath(media_id)
-        #     yield self._respond_with_file(request, media_info["media_type"], file_path)
+        #     yield respond_with_file(request, media_info["media_type"], file_path)
         #     return
 
         thumbnail_infos = yield self.store.get_local_media_thumbnails(media_id)
@@ -124,18 +135,18 @@ class ThumbnailResource(BaseMediaResource):
                 file_path = self.filepaths.local_media_thumbnail(
                     media_id, desired_width, desired_height, desired_type, desired_method,
                 )
-                yield self._respond_with_file(request, desired_type, file_path)
+                yield respond_with_file(request, desired_type, file_path)
                 return
 
         logger.debug("We don't have a local thumbnail of that size. Generating")
 
         # Okay, so we generate one.
-        file_path = yield self._generate_local_exact_thumbnail(
+        file_path = yield self.media_repo.generate_local_exact_thumbnail(
             media_id, desired_width, desired_height, desired_method, desired_type
         )
 
         if file_path:
-            yield self._respond_with_file(request, desired_type, file_path)
+            yield respond_with_file(request, desired_type, file_path)
         else:
             yield self._respond_default_thumbnail(
                 request, media_info, desired_width, desired_height,
@@ -146,11 +157,11 @@ class ThumbnailResource(BaseMediaResource):
     def _select_or_generate_remote_thumbnail(self, request, server_name, media_id,
                                              desired_width, desired_height,
                                              desired_method, desired_type):
-        media_info = yield self._get_remote_media(server_name, media_id)
+        media_info = yield self.media_repo.get_remote_media(server_name, media_id)
 
         # if media_info["media_type"] == "image/svg+xml":
         #     file_path = self.filepaths.remote_media_filepath(server_name, media_id)
-        #     yield self._respond_with_file(request, media_info["media_type"], file_path)
+        #     yield respond_with_file(request, media_info["media_type"], file_path)
         #     return
 
         thumbnail_infos = yield self.store.get_remote_media_thumbnails(
@@ -170,19 +181,19 @@ class ThumbnailResource(BaseMediaResource):
                     server_name, file_id, desired_width, desired_height,
                     desired_type, desired_method,
                 )
-                yield self._respond_with_file(request, desired_type, file_path)
+                yield respond_with_file(request, desired_type, file_path)
                 return
 
         logger.debug("We don't have a local thumbnail of that size. Generating")
 
         # Okay, so we generate one.
-        file_path = yield self._generate_remote_exact_thumbnail(
+        file_path = yield self.media_repo.generate_remote_exact_thumbnail(
             server_name, file_id, media_id, desired_width,
             desired_height, desired_method, desired_type
         )
 
         if file_path:
-            yield self._respond_with_file(request, desired_type, file_path)
+            yield respond_with_file(request, desired_type, file_path)
         else:
             yield self._respond_default_thumbnail(
                 request, media_info, desired_width, desired_height,
@@ -194,11 +205,11 @@ class ThumbnailResource(BaseMediaResource):
                                   height, method, m_type):
         # TODO: Don't download the whole remote file
         # We should proxy the thumbnail from the remote server instead.
-        media_info = yield self._get_remote_media(server_name, media_id)
+        media_info = yield self.media_repo.get_remote_media(server_name, media_id)
 
         # if media_info["media_type"] == "image/svg+xml":
         #     file_path = self.filepaths.remote_media_filepath(server_name, media_id)
-        #     yield self._respond_with_file(request, media_info["media_type"], file_path)
+        #     yield respond_with_file(request, media_info["media_type"], file_path)
         #     return
 
         thumbnail_infos = yield self.store.get_remote_media_thumbnails(
@@ -219,7 +230,7 @@ class ThumbnailResource(BaseMediaResource):
             file_path = self.filepaths.remote_media_thumbnail(
                 server_name, file_id, t_width, t_height, t_type, t_method,
             )
-            yield self._respond_with_file(request, t_type, file_path, t_length)
+            yield respond_with_file(request, t_type, file_path, t_length)
         else:
             yield self._respond_default_thumbnail(
                 request, media_info, width, height, method, m_type,
@@ -245,7 +256,7 @@ class ThumbnailResource(BaseMediaResource):
                 "_default", "_default",
             )
         if not thumbnail_infos:
-            self._respond_404(request)
+            respond_404(request)
             return
 
         thumbnail_info = self._select_thumbnail(
@@ -261,7 +272,7 @@ class ThumbnailResource(BaseMediaResource):
         file_path = self.filepaths.default_thumbnail(
             top_level_type, sub_type, t_width, t_height, t_type, t_method,
         )
-        yield self.respond_with_file(request, t_type, file_path, t_length)
+        yield respond_with_file(request, t_type, file_path, t_length)
 
     def _select_thumbnail(self, desired_width, desired_height, desired_method,
                           desired_type, thumbnail_infos):
diff --git a/synapse/rest/media/v1/upload_resource.py b/synapse/rest/media/v1/upload_resource.py
index 9c7ad4ae85..299e1f6e56 100644
--- a/synapse/rest/media/v1/upload_resource.py
+++ b/synapse/rest/media/v1/upload_resource.py
@@ -15,20 +15,33 @@
 
 from synapse.http.server import respond_with_json, request_handler
 
-from synapse.util.stringutils import random_string
 from synapse.api.errors import SynapseError
 
 from twisted.web.server import NOT_DONE_YET
 from twisted.internet import defer
 
-from .base_resource import BaseMediaResource
+from twisted.web.resource import Resource
 
 import logging
 
 logger = logging.getLogger(__name__)
 
 
-class UploadResource(BaseMediaResource):
+class UploadResource(Resource):
+    isLeaf = True
+
+    def __init__(self, hs, media_repo):
+        Resource.__init__(self)
+
+        self.media_repo = media_repo
+        self.filepaths = media_repo.filepaths
+        self.store = hs.get_datastore()
+        self.clock = hs.get_clock()
+        self.server_name = hs.hostname
+        self.auth = hs.get_auth()
+        self.max_upload_size = hs.config.max_upload_size
+        self.version_string = hs.version_string
+
     def render_POST(self, request):
         self._async_render_POST(request)
         return NOT_DONE_YET
@@ -37,36 +50,6 @@ class UploadResource(BaseMediaResource):
         respond_with_json(request, 200, {}, send_cors=True)
         return NOT_DONE_YET
 
-    @defer.inlineCallbacks
-    def create_content(self, media_type, upload_name, content, content_length,
-                       auth_user):
-        media_id = random_string(24)
-
-        fname = self.filepaths.local_media_filepath(media_id)
-        self._makedirs(fname)
-
-        # This shouldn't block for very long because the content will have
-        # already been uploaded at this point.
-        with open(fname, "wb") as f:
-            f.write(content)
-
-        yield self.store.store_local_media(
-            media_id=media_id,
-            media_type=media_type,
-            time_now_ms=self.clock.time_msec(),
-            upload_name=upload_name,
-            media_length=content_length,
-            user_id=auth_user,
-        )
-        media_info = {
-            "media_type": media_type,
-            "media_length": content_length,
-        }
-
-        yield self._generate_local_thumbnails(media_id, media_info)
-
-        defer.returnValue("mxc://%s/%s" % (self.server_name, media_id))
-
     @request_handler
     @defer.inlineCallbacks
     def _async_render_POST(self, request):
@@ -108,7 +91,7 @@ class UploadResource(BaseMediaResource):
         #     disposition = headers.getRawHeaders("Content-Disposition")[0]
         # TODO(markjh): parse content-dispostion
 
-        content_uri = yield self.create_content(
+        content_uri = yield self.media_repo.create_content(
             media_type, upload_name, request.content.read(),
             content_length, requester.user
         )