Improve handling of disappearing photos and files

Fixes #508

mautrix_telegram/__main__.py

@@ -86,6 +86,7 @@ class TelegramBridge(Bridge):
         Portal.init_cls(self)
         self.add_startup_actions(Puppet.init_cls(self))
         self.add_startup_actions(User.init_cls(self))
+        self.add_startup_actions(Portal.restart_scheduled_disappearing())
         if self.bot:
             self.add_startup_actions(self.bot.start())
         if self.config["bridge.resend_bridge_info"]:

mautrix_telegram/db/__init__.py

@@ -16,6 +16,7 @@
 from mautrix.util.async_db import Database
 
 from .bot_chat import BotChat
+from .disappearing_message import DisappearingMessage
 from .message import Message
 from .portal import Portal
 from .puppet import Puppet
@@ -27,7 +28,17 @@ from .user import User
 
 
 def init(db: Database) -> None:
-    for table in (Portal, Message, Reaction, User, Puppet, TelegramFile, BotChat, PgSession):
+    for table in (
+        Portal,
+        Message,
+        Reaction,
+        User,
+        Puppet,
+        TelegramFile,
+        BotChat,
+        PgSession,
+        DisappearingMessage,
+    ):
         table.db = db
 
 
@@ -42,4 +53,5 @@ __all__ = [
     "TelegramFile",
     "BotChat",
     "PgSession",
+    "DisappearingMessage",
 ]

mautrix_telegram/db/disappearing_message.py

@@ -0,0 +1,78 @@
+# mautrix-telegram - A Matrix-Telegram puppeting bridge
+# Copyright (C) 2021 Sumner Evans
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+import asyncpg
+
+from mautrix.bridge import AbstractDisappearingMessage
+from mautrix.types import EventID, RoomID
+from mautrix.util.async_db import Database
+
+fake_db = Database.create("") if TYPE_CHECKING else None
+
+
+class DisappearingMessage(AbstractDisappearingMessage):
+    db: ClassVar[Database] = fake_db
+
+    async def insert(self) -> None:
+        q = """
+            INSERT INTO disappearing_message (room_id, event_id, expiration_seconds, expiration_ts)
+            VALUES ($1, $2, $3, $4)
+        """
+        await self.db.execute(
+            q, self.room_id, self.event_id, self.expiration_seconds, self.expiration_ts
+        )
+
+    async def update(self) -> None:
+        q = "UPDATE disappearing_message SET expiration_ts=$3 WHERE room_id=$1 AND event_id=$2"
+        await self.db.execute(q, self.room_id, self.event_id, self.expiration_ts)
+
+    async def delete(self) -> None:
+        q = "DELETE from disappearing_message WHERE room_id=$1 AND event_id=$2"
+        await self.db.execute(q, self.room_id, self.event_id)
+
+    @classmethod
+    def _from_row(cls, row: asyncpg.Record) -> DisappearingMessage:
+        return cls(**row)
+
+    @classmethod
+    async def get(cls, room_id: RoomID, event_id: EventID) -> DisappearingMessage | None:
+        q = """
+            SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message
+            WHERE room_id=$1 AND mxid=$2
+        """
+        try:
+            return cls._from_row(await cls.db.fetchrow(q, room_id, event_id))
+        except Exception:
+            return None
+
+    @classmethod
+    async def get_all_scheduled(cls) -> list[DisappearingMessage]:
+        q = """
+            SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message
+            WHERE expiration_ts IS NOT NULL
+        """
+        return [cls._from_row(r) for r in await cls.db.fetch(q)]
+
+    @classmethod
+    async def get_unscheduled_for_room(cls, room_id: RoomID) -> list[DisappearingMessage]:
+        q = """
+            SELECT room_id, event_id, expiration_seconds, expiration_ts FROM disappearing_message
+            WHERE room_id = $1 AND expiration_ts IS NULL
+        """
+        return [cls._from_row(r) for r in await cls.db.fetch(q, room_id)]

mautrix_telegram/db/upgrade/__init__.py

@@ -2,4 +2,4 @@ from mautrix.util.async_db import UpgradeTable
 
 upgrade_table = UpgradeTable()
 
-from . import v01_initial_revision, v02_sponsored_events, v03_reactions
+from . import v01_initial_revision, v02_sponsored_events, v03_reactions, v04_disappearing_messages

mautrix_telegram/db/upgrade/v04_disappearing_messages.py

@@ -0,0 +1,32 @@
+# mautrix-telegram - A Matrix-Telegram puppeting bridge
+# Copyright (C) 2022 Tulir Asokan
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from asyncpg import Connection
+
+from . import upgrade_table
+
+
+@upgrade_table.register(description="Add support for disappearing messages")
+async def upgrade_v4(conn: Connection) -> None:
+    await conn.execute(
+        """CREATE TABLE disappearing_message (
+            room_id             TEXT,
+            event_id            TEXT,
+            expiration_seconds  BIGINT,
+            expiration_ts       BIGINT,
+
+            PRIMARY KEY (room_id, event_id)
+        )"""
+    )

mautrix_telegram/portal.py

@@ -195,6 +195,7 @@ from mautrix.util.simple_template import SimpleTemplate
 from . import abstract_user as au, formatter, portal_util as putil, puppet as p, user as u, util
 from .config import Config
 from .db import (
+    DisappearingMessage,
     Message as DBMessage,
     Portal as DBPortal,
     Reaction as DBReaction,
@@ -243,6 +244,7 @@ class DocAttrs(NamedTuple):
 class Portal(DBPortal, BasePortal):
     bot: "Bot"
     config: Config
+    disappearing_msg_class = DisappearingMessage
 
     # Instance cache
     by_mxid: dict[RoomID, Portal] = {}
@@ -321,6 +323,7 @@ class Portal(DBPortal, BasePortal):
             photo_id=photo_id,
             local_config=local_config or {},
         )
+        BasePortal.__init__(self)
         self.log = self.log.getChild(self.tgid_log if self.tgid else self.mxid)
         self._main_intent = None
         self.deleted = False
@@ -2106,15 +2109,6 @@ class Portal(DBPortal, BasePortal):
             return f"https://t.me/c/{self.tgid}/{evt.id}"
         return None
 
-    async def _expire_telegram_photo(self, intent: IntentAPI, event_id: EventID, ttl: int) -> None:
-        try:
-            content = TextMessageEventContent(msgtype=MessageType.NOTICE, body="Photo has expired")
-            content.set_edit(event_id)
-            await asyncio.sleep(ttl)
-            await self._send_message(intent, content)
-        except Exception:
-            self.log.warning("Failed to expire Telegram photo %s", event_id, exc_info=True)
-
     async def _handle_telegram_photo(
         self, source: au.AbstractUser, intent: IntentAPI, evt: Message, relates_to: RelatesTo
     ) -> EventID | None:
@@ -2171,13 +2165,15 @@ class Portal(DBPortal, BasePortal):
             content.url = file.mxc
         result = await self._send_message(intent, content, timestamp=evt.date)
         if media.ttl_seconds:
-            asyncio.create_task(self._expire_telegram_photo(intent, result, media.ttl_seconds))
+            await DisappearingMessage(self.mxid, result, media.ttl_seconds).insert()
         if evt.message:
             caption_content = await formatter.telegram_to_matrix(
                 evt, source, self.main_intent, no_reply_fallback=True
             )
             caption_content.external_url = content.external_url
             result = await self._send_message(intent, caption_content, timestamp=evt.date)
+            if media.ttl_seconds:
+                await DisappearingMessage(self.mxid, result, media.ttl_seconds).insert()
         return result
 
     @staticmethod
@@ -2355,12 +2351,16 @@ class Portal(DBPortal, BasePortal):
         else:
             content.url = file.mxc
         res = await self._send_message(intent, content, event_type=event_type, timestamp=evt.date)
+        if evt.media.ttl_seconds:
+            await DisappearingMessage(self.mxid, res, evt.media.ttl_seconds).insert()
         if evt.message:
             caption_content = await formatter.telegram_to_matrix(
                 evt, source, self.main_intent, no_reply_fallback=True
             )
             caption_content.external_url = content.external_url
             res = await self._send_message(intent, caption_content, timestamp=evt.date)
+            if evt.media.ttl_seconds:
+                await DisappearingMessage(self.mxid, res, evt.media.ttl_seconds).insert()
         return res
 
     def _location_message_to_content(