From f0d8f66eaaacfa75bed65bc5d0c602fbc5339c85 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 29 Mar 2023 14:37:06 +0100
Subject: [PATCH] Fix registering a device on an account with lots of devices
 (#15348)

Fixes up #15183
---
 changelog.d/15348.misc                    |  1 +
 synapse/handlers/register.py              |  2 +
 synapse/storage/databases/main/devices.py |  9 +++--
 tests/rest/client/test_register.py        | 47 +++++++++++++++++++++++
 4 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 changelog.d/15348.misc

diff --git a/changelog.d/15348.misc b/changelog.d/15348.misc
new file mode 100644
index 000000000000..f9bfc581ad53
--- /dev/null
+++ b/changelog.d/15348.misc
@@ -0,0 +1 @@
+Prune user's old devices on login if they have too many.
diff --git a/synapse/handlers/register.py b/synapse/handlers/register.py
index bb1df1e60fc1..7e9d065f50e1 100644
--- a/synapse/handlers/register.py
+++ b/synapse/handlers/register.py
@@ -946,6 +946,8 @@ async def _maybe_prune_too_many_devices(self, user_id: str) -> None:
         if not device_ids:
             return
 
+        logger.info("Pruning %d stale devices for %s", len(device_ids), user_id)
+
         # Now spawn a background loop that deletes said devices.
         async def _prune_too_many_devices_loop() -> None:
             if user_id in self._currently_pruning_devices_for_users:
diff --git a/synapse/storage/databases/main/devices.py b/synapse/storage/databases/main/devices.py
index 7647cda2c627..f61b7bc96eec 100644
--- a/synapse/storage/databases/main/devices.py
+++ b/synapse/storage/databases/main/devices.py
@@ -1638,19 +1638,22 @@ async def check_too_many_devices_for_user(self, user_id: str) -> List[str]:
             """
 
             rows = await self.db_pool.execute(
-                "check_too_many_devices_for_user_last_seen", None, sql, (user_id,)
+                "check_too_many_devices_for_user_last_seen",
+                None,
+                sql,
+                user_id,
             )
             if rows:
                 max_last_seen = max(rows[0][0], max_last_seen)
 
         # Fetch the devices to delete.
         sql = """
-            SELECT DISTINCT device_id FROM devices
+            SELECT device_id FROM devices
             LEFT JOIN e2e_device_keys_json USING (user_id, device_id)
             WHERE
                 user_id = ?
                 AND NOT hidden
-                AND last_seen < ?
+                AND last_seen <= ?
                 AND key_json IS NULL
             ORDER BY last_seen
         """
diff --git a/tests/rest/client/test_register.py b/tests/rest/client/test_register.py
index b228dba8613d..7ae84e3139aa 100644
--- a/tests/rest/client/test_register.py
+++ b/tests/rest/client/test_register.py
@@ -794,6 +794,53 @@ def test_require_approval(self) -> None:
             ApprovalNoticeMedium.NONE, channel.json_body["approval_notice_medium"]
         )
 
+    def test_check_stale_devices_get_pruned(self) -> None:
+        """Check that if a user has some stale devices we log them out when they
+        log in a new device."""
+
+        # Register some devices, but not too many that we go over the threshold
+        # where we prune more aggressively.
+        user_id = self.register_user("user", "pass")
+        for _ in range(0, 50):
+            self.login(user_id, "pass")
+
+        store = self.hs.get_datastores().main
+
+        res = self.get_success(store.get_devices_by_user(user_id))
+        self.assertEqual(len(res), 50)
+
+        # Advance time so that the above devices are considered "old".
+        self.reactor.advance(30 * 24 * 60 * 60 * 1000)
+
+        self.login(user_id, "pass")
+
+        self.reactor.pump([60] * 10)  # Ensure background job runs
+
+        # We expect all old devices to have been logged out
+        res = self.get_success(store.get_devices_by_user(user_id))
+        self.assertEqual(len(res), 1)
+
+    def test_check_recent_devices_get_pruned(self) -> None:
+        """Check that if a user has many devices we log out the last oldest
+        ones.
+
+        Note: this is similar to above, except if we lots of devices we prune
+        devices even if they're not old.
+        """
+
+        # Register a lot of devices in a short amount of time
+        user_id = self.register_user("user", "pass")
+        for _ in range(0, 100):
+            self.login(user_id, "pass")
+            self.reactor.advance(100)
+
+        store = self.hs.get_datastores().main
+
+        # We keep up to 50 devices that have been used in the last week, plus
+        # the device that was last logged in.
+        res = self.get_success(store.get_devices_by_user(user_id))
+        self.assertEqual(len(res), 51)
+
 
 class AccountValidityTestCase(unittest.HomeserverTestCase):
     servlets = [