pantalaimon: Add a config option to drop old room keys

This patch adds a config option that enables a mode wher a proxy will
only keep the latest room key from a sender in a certain room. This is
useful for bots that use pantalaimon since they are mostly only
interested in the latest messages and don't care about room history.

contrib/pantalaimon.conf

@@ -12,3 +12,4 @@ Proxy = http://localhost:8080
 SSL = False
 IgnoreVerification = False
 UseKeyring = True
+DropOldKeys = False

docs/man/pantalaimon.5

@@ -51,6 +51,11 @@ This option configures if a proxy instance should use the OS keyring to store
 its own access tokens. The access tokens are required for the daemon to resume
 operation. If this is set to "No", access tokens are stored in the pantalaimon
 database in plaintext. Defaults to "Yes".
+.It Cm DropOldKeys
+This option configures if a proxy instance should only keep the latest version
+of a room key from a certain user around. This effectively means that only newly
+incoming messages will be decryptable, the proxy will be unable to decrypt the
+room history.  Defaults to "No".
 .It Cm SearchRequests
 This option configures if the proxy should make additional HTTP requests to the
 server when clients use the search API endpoint. Some data that is required to

docs/man/pantalaimon.5.md

@@ -62,6 +62,13 @@ The following keys are optional in the proxy instance sections:
 > operation. If this is set to "No", access tokens are stored in the pantalaimon
 > database in plaintext. Defaults to "Yes".
 
+**DropOldKeys**
+
+> This option configures if a proxy instance should only keep the latest version
+> of a room key from a certain user around. This effectively means that only newly
+> incoming messages will be decryptable, the proxy will be unable to decrypt the
+> room history.  Defaults to "No".
+
 Aditional to the homeserver section a special section with the name
 **Default**
 can be used to configure the following values for all homeservers:
@@ -150,4 +157,4 @@ pantalaimon(8)
 was written by
 Damir Jelić <[poljar@termina.org.uk](mailto:poljar@termina.org.uk)>.
 
-Linux 5.1.3-arch2-1-ARCH - May 8, 2019
+Linux 5.11.16-arch1-1 - May 8, 2019

pantalaimon/config.py

@@ -39,6 +39,7 @@ class PanConfigParser(configparser.ConfigParser):
                 "IndexingBatchSize": "100",
                 "HistoryFetchDelay": "3000",
                 "DebugEncryption": "False",
+                "DropOldKeys": "False",
             },
             converters={
                 "address": parse_address,
@@ -121,6 +122,8 @@ class ServerConfig:
             the room history.
         history_fetch_delay (int): The delay between room history fetching
             requests in seconds.
+        drop_old_keys (bool): Should Pantalaimon only keep the most recent
+            decryption key around.
     """
 
     name = attr.ib(type=str)
@@ -137,6 +140,7 @@ class ServerConfig:
     index_encrypted_only = attr.ib(type=bool, default=True)
     indexing_batch_size = attr.ib(type=int, default=100)
     history_fetch_delay = attr.ib(type=int, default=3)
+    drop_old_keys = attr.ib(type=bool, default=False)
 
 
 @attr.s
@@ -229,6 +233,7 @@ class PanConfig:
                         f"already defined before."
                     )
                 listen_set.add(listen_tuple)
+                drop_old_keys = section.getboolean("DropOldKeys")
 
                 server_conf = ServerConfig(
                     section_name,
@@ -243,6 +248,7 @@ class PanConfig:
                     index_encrypted_only,
                     indexing_batch_size,
                     history_fetch_delay / 1000,
+                    drop_old_keys,
                 )
 
                 self.servers[section_name] = server_conf

pantalaimon/main.py

@@ -29,6 +29,7 @@ from logbook import StderrHandler
 from pantalaimon.config import PanConfig, PanConfigError, parse_log_level
 from pantalaimon.daemon import ProxyDaemon
 from pantalaimon.log import logger
+from pantalaimon.store import KeyDroppingSqliteStore
 from pantalaimon.thread_messages import DaemonResponse
 from pantalaimon.ui import UI_ENABLED
 
@@ -47,6 +48,8 @@ def create_dirs(data_dir, conf_dir):
 
 async def init(data_dir, server_conf, send_queue, recv_queue):
     """Initialize the proxy and the http server."""
+    store_class = KeyDroppingSqliteStore if server_conf.drop_old_keys else None
+
     proxy = ProxyDaemon(
         server_conf.name,
         server_conf.homeserver,
@@ -56,6 +59,7 @@ async def init(data_dir, server_conf, send_queue, recv_queue):
         recv_queue=recv_queue.async_q if recv_queue else None,
         proxy=server_conf.proxy.geturl() if server_conf.proxy else None,
         ssl=None if server_conf.ssl is True else False,
+        client_store_class=store_class,
     )
 
     # 100 MB max POST size
@@ -101,7 +105,6 @@ async def init(data_dir, server_conf, send_queue, recv_queue):
                 r"/_matrix/client/r0/profile/{userId}/avatar_url",
                 proxy.profile,
             ),
-
         ]
     )
     app.router.add_route("*", "/" + "{proxyPath:.*}", proxy.router)

pantalaimon/store.py

@@ -18,10 +18,12 @@ from collections import defaultdict
 from typing import Any, Dict, List, Optional, Tuple
 
 import attr
-from nio.crypto import TrustState
+from nio.crypto import TrustState, GroupSessionStore
 from nio.store import (
     Accounts,
+    MegolmInboundSessions,
     DeviceKeys,
+    SqliteStore,
     DeviceTrustState,
     use_database,
     use_database_atomic,
@@ -29,7 +31,6 @@ from nio.store import (
 from peewee import SQL, DoesNotExist, ForeignKeyField, Model, SqliteDatabase, TextField
 from cachetools import LRUCache
 
-
 MAX_LOADED_MEDIA = 10000
 MAX_LOADED_UPLOAD = 10000
 
@@ -452,3 +453,47 @@ class PanStore:
             store[account.user_id] = device_store
 
         return store
+
+
+class KeyDroppingSqliteStore(SqliteStore):
+    @use_database
+    def save_inbound_group_session(self, session):
+        """Save the provided Megolm inbound group session to the database.
+
+        Args:
+            session (InboundGroupSession): The session to save.
+        """
+        account = self._get_account()
+        assert account
+
+        MegolmInboundSessions.delete().where(
+            MegolmInboundSessions.sender_key == session.sender_key,
+            MegolmInboundSessions.account == account,
+            MegolmInboundSessions.room_id == session.room_id,
+        ).execute()
+
+        super().save_inbound_group_session(session)
+
+    @use_database
+    def load_inbound_group_sessions(self):
+        store = super().load_inbound_group_sessions()
+
+        return KeyDroppingGroupSessionStore.from_group_session_store(store)
+
+
+class KeyDroppingGroupSessionStore(GroupSessionStore):
+    def from_group_session_store(store):
+        new_store = KeyDroppingGroupSessionStore()
+        new_store._entries = store._entries
+
+        return new_store
+
+    def add(self, session) -> bool:
+        room_id = session.room_id
+        sender_key = session.sender_key
+        if session in self._entries[room_id][sender_key].values():
+            return False
+
+        self._entries[room_id][sender_key].clear()
+        self._entries[room_id][sender_key][session.id] = session
+        return True