From 14a45be321e10878773ac0c39608c0c60d09cd92 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20H=C3=B6rist?= <philipp@hoerist.com>
Date: Thu, 21 Feb 2019 00:02:45 +0100
Subject: [PATCH] [omemo] Set devices inactive after 300 unacknowledged
 messages

---
 omemo/backend/devices.py          | 15 +++++++++++++++
 omemo/backend/liteaxolotlstore.py | 21 ++++++++++++++++-----
 omemo/backend/state.py            |  7 +++++++
 omemo/backend/util.py             |  1 +
 4 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/omemo/backend/devices.py b/omemo/backend/devices.py
index ca57aaf..5f35e6b 100644
--- a/omemo/backend/devices.py
+++ b/omemo/backend/devices.py
@@ -19,6 +19,8 @@ from collections import defaultdict
 
 from gajim.common import app
 
+from omemo.backend.util import UNACKNOWLEDGED_COUNT
+
 log = logging.getLogger('gajim.plugin_system.omemo')
 
 
@@ -39,6 +41,15 @@ class DeviceManager:
             self.add_device(jid, device)
 
     def update_devicelist(self, jid, devicelist):
+        for device in list(devicelist):
+            if device == self.own_device:
+                continue
+            count = self._storage.getUnacknowledgedCount(jid, device)
+            if count > UNACKNOWLEDGED_COUNT:
+                log.warning('Ignore device because of %s unacknowledged'
+                            ' messages: %s %s', count, jid, device)
+                devicelist.remove(device)
+
         self.__device_store[jid] = set(devicelist)
         log.info('Saved devices for %s', jid)
         self._storage.setActiveState(jid, devicelist)
@@ -60,6 +71,10 @@ class DeviceManager:
     def add_device(self, jid, device):
         self.__device_store[jid].add(device)
 
+    def remove_device(self, jid, device):
+        self.__device_store[jid].discard(device)
+        self._storage.setInactive(jid, device)
+
     def get_devices(self, jid, without_self=False):
         devices = set(self.__device_store[jid])
         if without_self:
diff --git a/omemo/backend/liteaxolotlstore.py b/omemo/backend/liteaxolotlstore.py
index 004d3ed..f6995f5 100644
--- a/omemo/backend/liteaxolotlstore.py
+++ b/omemo/backend/liteaxolotlstore.py
@@ -320,16 +320,22 @@ class LiteAxolotlStore(AxolotlStore):
                        ', '.join(['?'] * len(recipientIds)))
         return self._con.execute(query, recipientIds).fetchall()
 
-    def setActiveState(self, jid, deviceList):
+    def setActiveState(self, jid, devicelist):
         query = '''UPDATE sessions SET active = 1
                    WHERE recipient_id = ? AND device_id IN ({})'''.format(
-                       ', '.join(['?'] * len(deviceList)))
-        self._con.execute(query, (jid,) + tuple(deviceList))
+                       ', '.join(['?'] * len(devicelist)))
+        self._con.execute(query, (jid,) + tuple(devicelist))
 
         query = '''UPDATE sessions SET active = 0
                    WHERE recipient_id = ? AND device_id NOT IN ({})'''.format(
-                       ', '.join(['?'] * len(deviceList)))
-        self._con.execute(query, (jid,) + tuple(deviceList))
+                       ', '.join(['?'] * len(devicelist)))
+        self._con.execute(query, (jid,) + tuple(devicelist))
+        self._con.commit()
+
+    def setInactive(self, jid, device_id):
+        query = '''UPDATE sessions SET active = 0
+                   WHERE recipient_id = ? AND device_id = ?'''
+        self._con.execute(query, (jid, device_id))
         self._con.commit()
 
     def getInactiveSessionsKeys(self, recipientId):
@@ -515,3 +521,8 @@ class LiteAxolotlStore(AxolotlStore):
                    WHERE recipient_id = ? AND public_key = ?'''
         self._con.execute(query, (timestamp, recipient_id, identity_key))
         self._con.commit()
+
+    def getUnacknowledgedCount(self, recipient_id, device_id):
+        record = self.loadSession(recipient_id, device_id)
+        state = record.getSessionState()
+        return state.getSenderChainKey().getIndex()
diff --git a/omemo/backend/state.py b/omemo/backend/state.py
index c4d6105..5e4237a 100644
--- a/omemo/backend/state.py
+++ b/omemo/backend/state.py
@@ -42,6 +42,7 @@ from omemo.backend.util import DEFAULT_PREKEY_AMOUNT
 from omemo.backend.util import MIN_PREKEY_AMOUNT
 from omemo.backend.util import SPK_CYCLE_TIME
 from omemo.backend.util import SPK_ARCHIVE_TIME
+from omemo.backend.util import UNACKNOWLEDGED_COUNT
 
 
 log = logging.getLogger('gajim.plugin_system.omemo')
@@ -172,6 +173,12 @@ class OmemoState(DeviceManager):
         whisper_messages = defaultdict(dict)
 
         for jid_, device in devices_for_encryption:
+            count = self._storage.getUnacknowledgedCount(jid_, device)
+            if count >= UNACKNOWLEDGED_COUNT:
+                log.warning('Set device inactive %s because of %s '
+                            'unacknowledged messages', device, count)
+                self.remove_device(jid_, device)
+
             try:
                 whisper_messages[jid_][device] = self._get_whisper_message(
                     jid_, device, result.key)
diff --git a/omemo/backend/util.py b/omemo/backend/util.py
index dd3ea7a..b31539c 100644
--- a/omemo/backend/util.py
+++ b/omemo/backend/util.py
@@ -22,6 +22,7 @@ DEFAULT_PREKEY_AMOUNT = 100
 MIN_PREKEY_AMOUNT = 80
 SPK_ARCHIVE_TIME = 86400 * 15  # 15 Days
 SPK_CYCLE_TIME = 86400         # 24 Hours
+UNACKNOWLEDGED_COUNT = 300
 
 
 class Trust(IntEnum):