diff --git a/mautrix_telegram/commands/portal/bridge.py b/mautrix_telegram/commands/portal/bridge.py index 833bdc4a..a9136f28 100644 --- a/mautrix_telegram/commands/portal/bridge.py +++ b/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", diff --git a/mautrix_telegram/config.py b/mautrix_telegram/config.py index 3613a374..bfba842a 100644 --- a/mautrix_telegram/config.py +++ b/mautrix_telegram/config.py @@ -137,6 +137,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") diff --git a/mautrix_telegram/db/portal.py b/mautrix_telegram/db/portal.py index 344509e6..4c66f1cb 100644 --- a/mautrix_telegram/db/portal.py +++ b/mautrix_telegram/db/portal.py @@ -15,7 +15,7 @@ # along with this program. If not, see . from typing import Optional, Iterable -from sqlalchemy import Column, Integer, String, Boolean, Text, func, sql +from sqlalchemy import Column, Integer, 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() \ No newline at end of file diff --git a/mautrix_telegram/example-config.yaml b/mautrix_telegram/example-config.yaml index 4401afd8..9a3684c2 100644 --- a/mautrix_telegram/example-config.yaml +++ b/mautrix_telegram/example-config.yaml @@ -208,6 +208,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. diff --git a/mautrix_telegram/portal/base.py b/mautrix_telegram/portal/base.py index d8233e4a..5bc8814d 100644 --- a/mautrix_telegram/portal/base.py +++ b/mautrix_telegram/portal/base.py @@ -397,6 +397,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']: @@ -456,6 +457,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) diff --git a/mautrix_telegram/portal/metadata.py b/mautrix_telegram/portal/metadata.py index 02a12384..22f34728 100644 --- a/mautrix_telegram/portal/metadata.py +++ b/mautrix_telegram/portal/metadata.py @@ -121,6 +121,8 @@ async def create_telegram_chat(self, source: 'u.User', supergroup: bool = False) 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.") invites = await self._get_telegram_users_in_matrix_room() if len(invites) < 2: @@ -303,6 +305,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)