Merge branch 'hs/portal-limit' into element-master

mautrix_telegram/commands/portal/bridge.py

@@ -92,6 +92,11 @@ async def bridge(evt: CommandEvent) -> EventID:
                                "`$cmdprefix+sp delete-and-continue`. To unbridge the portal "
                                "without kicking Matrix users, use `$cmdprefix+sp unbridge-and-"
                                "continue`. To cancel, use `$cmdprefix+sp cancel`")
+
+    if await po.Portal.reached_portal_limit():
+        return await evt.reply("This bridge has reached the maximum number of rooms that "
+                               "can be bridged.")
+
     evt.sender.command_status = {
         "next": confirm_bridge,
         "action": "Room bridging",

mautrix_telegram/config.py

@@ -138,6 +138,7 @@ def do_update(self, helper: ConfigUpdateHelper) -> None:
         copy("bridge.backfill.missed_limit")
         copy("bridge.backfill.disable_notifications")
         copy("bridge.backfill.normal_groups")
+        copy("bridge.max_portal_rooms")
 
         copy("bridge.initial_power_level_overrides.group")
         copy("bridge.initial_power_level_overrides.user")

mautrix_telegram/db/portal.py

@@ -15,7 +15,7 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from typing import Optional, Iterable
 
-from sqlalchemy import Column, BigInteger, String, Boolean, Text, func, sql
+from sqlalchemy import Column, BigInteger, String, Boolean, Text, Table, MetaData, func, sql, select
 
 from mautrix.types import RoomID, ContentURI
 from mautrix.util.db import Base
@@ -53,6 +53,11 @@ def get_by_tgid(cls, tgid: TelegramID, tg_receiver: TelegramID) -> Optional['Por
     def find_private_chats(cls, tg_receiver: TelegramID) -> Iterable['Portal']:
         yield from cls._select_all(cls.c.tg_receiver == tg_receiver, cls.c.peer_type == "user")
 
+    @classmethod
+    def count(cls) -> int:
+        count = cls.db.execute(select([func.count('*')]).select_from(Table("portal", MetaData()))).scalar()
+        return count
+
     @classmethod
     def get_by_mxid(cls, mxid: RoomID) -> Optional['Portal']:
         return cls._select_one_or_none(cls.c.mxid == mxid)
@@ -63,4 +68,4 @@ def get_by_username(cls, username: str) -> Optional['Portal']:
 
     @classmethod
     def all(cls) -> Iterable['Portal']:
-        yield from cls._select_all()
+        yield from cls._select_all()

mautrix_telegram/example-config.yaml

@@ -218,6 +218,9 @@ bridge:
     # Whether or not created rooms should have federation enabled.
     # If false, created portal rooms will never be federated.
     federate_rooms: true
+    # The maximum number of rooms that can be bridged on this instance.
+    # When the maximum number of rooms has been reached, attempts to create new links will error. Set to -1 to disable.
+    max_portal_rooms: -1
     # Settings for converting animated stickers.
     animated_sticker:
         # Format to which animated stickers should be converted.

mautrix_telegram/portal/base.py

@@ -413,6 +413,7 @@ def find_by_username(cls, username: str) -> Optional['Portal']:
 
         return None
 
+
     @classmethod
     def get_by_tgid(cls, tgid: TelegramID, tg_receiver: Optional[TelegramID] = None,
                     peer_type: str = None) -> Optional['Portal']:
@@ -472,6 +473,13 @@ def get_by_entity(cls, entity: Union[TypeChat, TypePeer, TypeUser, TypeUserFull,
                                receiver_id if type_name == "user" else entity_id,
                                type_name if create else None)
 
+    @classmethod
+    async def reached_portal_limit(cls) -> bool:
+        limit = config.get("bridge.max_portal_rooms", -1)
+        if limit == -1:
+            return False
+        return DBPortal.count() >= limit
+
     # endregion
     # region Abstract methods (cross-called in matrix/metadata/telegram classes)
 

mautrix_telegram/portal/metadata.py

@@ -132,6 +132,8 @@ async def create_telegram_chat(self, source: 'u.User', invites: List[InputUser],
             raise ValueError("Can't create Telegram chat for portal without Matrix room.")
         elif self.tgid:
             raise ValueError("Can't create Telegram chat for portal with existing Telegram chat.")
+        elif await self.reached_portal_limit():
+            raise ValueError("Can't create Telegram chat, reached portal limit.")
 
         if len(invites) < 2:
             if self.bot is not None:
@@ -329,6 +331,10 @@ async def _create_matrix_room(self, user: 'AbstractUser', entity: Union[TypeChat
         direct = self.peer_type == "user"
         invites = invites or []
 
+        if await self.reached_portal_limit():
+            raise ValueError("Can't create Telegram chat, reached portal limit.")
+
+
         if not entity:
             entity = await self.get_entity(user)
             self.log.trace("Fetched data: %s", entity)