From d3b535f07018d40721fe668b6bd793ab0ced2a29 Mon Sep 17 00:00:00 2001
From: Will Hunt <will@half-shot.uk>
Date: Tue, 22 Sep 2020 15:17:41 +0100
Subject: [PATCH] Add option to forget sessions when inserting a new session

---
 pantalaimon/client.py |  6 ++----
 pantalaimon/config.py |  7 +++++--
 pantalaimon/daemon.py |  1 +
 pantalaimon/store.py  | 20 ++++++++++++++++++++
 4 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/pantalaimon/client.py b/pantalaimon/client.py
index b2d2a72..7d9d7a1 100644
--- a/pantalaimon/client.py
+++ b/pantalaimon/client.py
@@ -18,7 +18,6 @@ from collections import defaultdict
 from pprint import pformat
 from typing import Any, Dict, Optional
 from urllib.parse import urlparse
-
 from aiohttp.client_exceptions import ClientConnectionError
 from jsonschema import Draft4Validator, FormatChecker, validators
 from playhouse.sqliteq import SqliteQueueDatabase
@@ -50,10 +49,9 @@ from nio import (
 )
 from nio.crypto import Sas
 from nio.store import SqliteStore
-
 from pantalaimon.index import INDEXING_ENABLED
 from pantalaimon.log import logger
-from pantalaimon.store import FetchTask, MediaInfo
+from pantalaimon.store import FetchTask, MediaInfo, PanSqliteStore
 from pantalaimon.thread_messages import (
     DaemonResponse,
     InviteSasSignal,
@@ -162,7 +160,7 @@ class PanClient(AsyncClient):
         media_info=None,
     ):
         config = config or AsyncClientConfig(
-            store=store_class or SqliteStore, store_name="pan.db"
+            store=store_class or PanSqliteStore, store_name="pan.db"
         )
         super().__init__(homeserver, user_id, device_id, store_path, config, ssl, proxy)
 
diff --git a/pantalaimon/config.py b/pantalaimon/config.py
index 2a01714..65dc62b 100644
--- a/pantalaimon/config.py
+++ b/pantalaimon/config.py
@@ -121,6 +121,8 @@ class ServerConfig:
             the room history.
         history_fetch_delay (int): The delay between room history fetching
             requests in seconds.
+        store_forgetful (bool): Enable or disable discarding of previous sessions
+            from the store.
     """
 
     name = attr.ib(type=str)
@@ -137,7 +139,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)
-
+    store_forgetful = attr.ib(type=bool, default=False)
 
 @attr.s
 class PanConfig:
@@ -201,7 +203,7 @@ class PanConfig:
                 proxy = section.geturl("Proxy")
                 search_requests = section.getboolean("SearchRequests")
                 index_encrypted_only = section.getboolean("IndexEncryptedOnly")
-
+                store_forgetful = config["Default"].getboolean("StoreForgetful")
                 indexing_batch_size = section.getint("IndexingBatchSize")
 
                 if not 1 < indexing_batch_size <= 1000:
@@ -243,6 +245,7 @@ class PanConfig:
                     index_encrypted_only,
                     indexing_batch_size,
                     history_fetch_delay / 1000,
+                    store_forgetful,
                 )
 
                 self.servers[section_name] = server_conf
diff --git a/pantalaimon/daemon.py b/pantalaimon/daemon.py
index 9efee6b..720ca81 100755
--- a/pantalaimon/daemon.py
+++ b/pantalaimon/daemon.py
@@ -150,6 +150,7 @@ class ProxyDaemon:
             pan_client.user_id = user_id
             pan_client.access_token = token
             pan_client.load_store()
+            pan_client.store.forgetful = self.conf.store_forgetful
             self.pan_clients[user_id] = pan_client
 
             loop.create_task(
diff --git a/pantalaimon/store.py b/pantalaimon/store.py
index 910df64..4a82d06 100644
--- a/pantalaimon/store.py
+++ b/pantalaimon/store.py
@@ -25,6 +25,9 @@ from nio.store import (
     DeviceTrustState,
     use_database,
     use_database_atomic,
+    SqliteStore,
+    MegolmInboundSessions,
+    ForwardedChains
 )
 from peewee import SQL, DoesNotExist, ForeignKeyField, Model, SqliteDatabase, TextField
 from cachetools import LRUCache
@@ -372,3 +375,20 @@ class PanStore:
             store[account.user_id] = device_store
 
         return store
+
+class PanSqliteStore(SqliteStore):
+    forgetful = False
+
+    @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.
+        """
+        # Delete previous sessions
+        if self.forgetful:
+            MegolmInboundSessions.delete().where(
+                (MegolmInboundSessions.sender_key == session.sender_key) |
+                (MegolmInboundSessions.room_id == session.room_id)
+            ).execute()
+        super().save_inbound_group_session(session)
\ No newline at end of file