From 3d89f8fcd938db374694637e4a0fc5b407c0cb5a Mon Sep 17 00:00:00 2001
From: Sumner Evans <sumner@beeper.com>
Date: Fri, 9 Jun 2023 17:20:25 -0600
Subject: [PATCH] media/create: enforce limit on number of pending uploads

Signed-off-by: Sumner Evans <sumner@beeper.com>
---
 synapse/media/media_repository.py             | 20 ++++++++++++
 synapse/rest/media/create_resource.py         |  1 +
 .../databases/main/media_repository.py        | 31 +++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/synapse/media/media_repository.py b/synapse/media/media_repository.py
index 17397599affe..6e87ca9d4c6b 100644
--- a/synapse/media/media_repository.py
+++ b/synapse/media/media_repository.py
@@ -210,6 +210,26 @@ async def create_media_id(self, auth_user: UserID) -> Tuple[str, int]:
         )
         return f"mxc://{self.server_name}/{media_id}", unused_expires_at
 
+    @trace
+    async def reached_pending_media_limit(
+        self, auth_user: UserID, limit: int
+    ) -> Tuple[bool, int]:
+        """Check if the user is over the limit for pending media uploads.
+
+        Args:
+            auth_user: The user_id of the uploader
+            limit: The maximum number of pending media uploads a user is allowed to have
+
+        Returns:
+            A tuple with a boolean and an integer indicating whether the user has too
+            many pending media uploads and the timestamp at which the first pending
+            media will expire, respectively.
+        """
+        pending, first_expiration_ts = await self.store.count_pending_media(
+            user_id=auth_user
+        )
+        return pending >= limit, first_expiration_ts
+
     @trace
     async def verify_can_upload(self, media_id: str, auth_user: UserID) -> None:
         """Verify that the media ID can be uploaded to by the given user. This
diff --git a/synapse/rest/media/create_resource.py b/synapse/rest/media/create_resource.py
index 5b23a36ee079..87b90a92eb4b 100644
--- a/synapse/rest/media/create_resource.py
+++ b/synapse/rest/media/create_resource.py
@@ -38,6 +38,7 @@ def __init__(self, hs: "HomeServer", media_repo: "MediaRepository"):
         self.media_repo = media_repo
         self.clock = hs.get_clock()
         self.auth = hs.get_auth()
+        self.max_pending_media_uploads = hs.config.media.max_pending_media_uploads
 
         # A rate limiter for creating new media IDs.
         self._create_media_rate_limiter = Ratelimiter(
diff --git a/synapse/storage/databases/main/media_repository.py b/synapse/storage/databases/main/media_repository.py
index 8f0fd7059256..71d87149131f 100644
--- a/synapse/storage/databases/main/media_repository.py
+++ b/synapse/storage/databases/main/media_repository.py
@@ -27,6 +27,7 @@
 )
 
 from synapse.api.constants import Direction
+from synapse.api.errors import StoreError
 from synapse.logging.opentracing import trace
 from synapse.media._base import ThumbnailInfo
 from synapse.storage._base import SQLBaseStore
@@ -409,6 +410,36 @@ async def mark_local_media_as_safe(self, media_id: str, safe: bool = True) -> No
             desc="mark_local_media_as_safe",
         )
 
+    async def count_pending_media(self, user_id: UserID) -> Tuple[int, int]:
+        """Count the number of pending media for a user.
+
+        Returns:
+            A tuple of two integers: the total pending media requests and the earliest
+            expiration timestamp.
+        """
+
+        def get_pending_media_txn(txn: LoggingTransaction) -> Tuple[int, int]:
+            sql = """
+            SELECT COUNT(*), MIN(created_at)
+            FROM local_media_repository
+            WHERE user_id = ?
+                AND created_at > ?
+                AND media_length IS NULL
+            """
+            txn.execute(
+                sql,
+                (
+                    user_id.to_string(),
+                    self._clock.time_msec() - self.unused_expiration_time,
+                ),
+            )
+            row = txn.fetchone()
+            if not row:
+                raise StoreError(404, "Failed to count pending media for user")
+            return row[0], row[1] or 0
+
+        return await self.db_pool.runInteraction("get_url_cache", get_pending_media_txn)
+
     async def get_url_cache(self, url: str, ts: int) -> Optional[Dict[str, Any]]:
         """Get the media_id and ts for a cached URL as of the given timestamp
         Returns: