From 3f06bbc0acaa7994b2df8c974e212f21f4a028ba Mon Sep 17 00:00:00 2001
From: Erik Johnston <erikj@element.io>
Date: Thu, 6 Jun 2024 17:10:58 +0100
Subject: [PATCH] Always return OTK counts (#17275)

Broke in https://github.com/element-hq/synapse/pull/17215
---
 changelog.d/17275.bugfix |  1 +
 synapse/handlers/sync.py | 33 +++++++++++++++++++++++++++++----
 2 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 changelog.d/17275.bugfix

diff --git a/changelog.d/17275.bugfix b/changelog.d/17275.bugfix
new file mode 100644
index 0000000000..eb522bb997
--- /dev/null
+++ b/changelog.d/17275.bugfix
@@ -0,0 +1 @@
+Fix bug where OTKs were not always included in `/sync` response when using workers.
diff --git a/synapse/handlers/sync.py b/synapse/handlers/sync.py
index e815e0ea7f..9d37e2a86f 100644
--- a/synapse/handlers/sync.py
+++ b/synapse/handlers/sync.py
@@ -285,7 +285,11 @@ class SyncResult:
         )
 
     @staticmethod
-    def empty(next_batch: StreamToken) -> "SyncResult":
+    def empty(
+        next_batch: StreamToken,
+        device_one_time_keys_count: JsonMapping,
+        device_unused_fallback_key_types: List[str],
+    ) -> "SyncResult":
         "Return a new empty result"
         return SyncResult(
             next_batch=next_batch,
@@ -297,8 +301,8 @@ class SyncResult:
             archived=[],
             to_device=[],
             device_lists=DeviceListUpdates(),
-            device_one_time_keys_count={},
-            device_unused_fallback_key_types=[],
+            device_one_time_keys_count=device_one_time_keys_count,
+            device_unused_fallback_key_types=device_unused_fallback_key_types,
         )
 
 
@@ -523,7 +527,28 @@ class SyncHandler:
                 logger.warning(
                     "Timed out waiting for worker to catch up. Returning empty response"
                 )
-                return SyncResult.empty(since_token)
+                device_id = sync_config.device_id
+                one_time_keys_count: JsonMapping = {}
+                unused_fallback_key_types: List[str] = []
+                if device_id:
+                    user_id = sync_config.user.to_string()
+                    # TODO: We should have a way to let clients differentiate between the states of:
+                    #   * no change in OTK count since the provided since token
+                    #   * the server has zero OTKs left for this device
+                    #  Spec issue: https://github.com/matrix-org/matrix-doc/issues/3298
+                    one_time_keys_count = await self.store.count_e2e_one_time_keys(
+                        user_id, device_id
+                    )
+                    unused_fallback_key_types = list(
+                        await self.store.get_e2e_unused_fallback_key_types(
+                            user_id, device_id
+                        )
+                    )
+
+                cache_context.should_cache = False  # Don't cache empty responses
+                return SyncResult.empty(
+                    since_token, one_time_keys_count, unused_fallback_key_types
+                )
 
             # If we've spent significant time waiting to catch up, take it off
             # the timeout.