Fetch encryption info from server if not cached

Author: tulir

mautrix/crypto/base.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022 Tulir Asokan
+# Copyright (c) 2023 Tulir Asokan
 #
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
@@ -12,13 +12,17 @@
 
 import olm
 
+from mautrix.errors import MForbidden, MNotFound
 from mautrix.types import (
     DeviceID,
     EncryptionKeyAlgorithm,
+    EventType,
     IdentityKey,
     KeyID,
     RequestedKeyInfo,
+    RoomEncryptionStateEventContent,
     RoomID,
+    RoomKeyEventContent,
     SessionID,
     SigningKey,
     TrustState,
@@ -82,6 +86,34 @@ def _mark_session_received(self, session_id: SessionID) -> None:
         except KeyError:
             return
 
+    async def _fill_encryption_info(self, evt: RoomKeyEventContent) -> None:
+        encryption_info = await self.state_store.get_encryption_info(evt.room_id)
+        if not encryption_info:
+            self.log.warning(
+                f"Encryption info for {evt.room_id} not found in state store, fetching from server"
+            )
+            try:
+                encryption_info = await self.client.get_state_event(
+                    evt.room_id, EventType.ROOM_ENCRYPTION
+                )
+            except (MNotFound, MForbidden) as e:
+                self.log.warning(
+                    f"Failed to get encryption info for {evt.room_id} from server: {e},"
+                    " using defaults"
+                )
+                encryption_info = RoomEncryptionStateEventContent()
+            if not encryption_info:
+                self.log.warning(
+                    f"Didn't find encryption info for {evt.room_id} on server either,"
+                    " using defaults"
+                )
+                encryption_info = RoomEncryptionStateEventContent()
+
+        if not evt.beeper_max_age_ms:
+            evt.beeper_max_age_ms = encryption_info.rotation_period_ms
+        if not evt.beeper_max_messages:
+            evt.beeper_max_messages = encryption_info.rotation_period_msgs
+
 
 canonical_json = functools.partial(
     json.dumps, ensure_ascii=False, separators=(",", ":"), sort_keys=True

mautrix/crypto/key_request.py

@@ -1,10 +1,9 @@
-# Copyright (c) 2022 Tulir Asokan
+# Copyright (c) 2023 Tulir Asokan
 #
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from typing import Dict, List, Optional, Union
-from datetime import timedelta
 import asyncio
 import uuid
 
@@ -122,12 +121,7 @@ async def _receive_forwarded_room_key(self, evt: DecryptedOlmEvent) -> None:
             )
             return
         if not key.beeper_max_messages or not key.beeper_max_age_ms:
-            encryption_info = await self.state_store.get_encryption_info(key.room_id)
-            if encryption_info:
-                if not key.beeper_max_age_ms:
-                    key.beeper_max_age_ms = encryption_info.rotation_period_ms
-                if not key.beeper_max_messages:
-                    key.beeper_max_messages = encryption_info.rotation_period_msgs
+            await self._fill_encryption_info(key)
         key.forwarding_key_chain.append(evt.sender_key)
         sess = InboundGroupSession.import_session(
             key.session_key,

mautrix/crypto/machine.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022 Tulir Asokan
+# Copyright (c) 2023 Tulir Asokan
 #
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
@@ -223,12 +223,7 @@ async def _receive_room_key(self, evt: DecryptedOlmEvent) -> None:
         if evt.content.algorithm != EncryptionAlgorithm.MEGOLM_V1 or not evt.keys.ed25519:
             return
         if not evt.content.beeper_max_messages or not evt.content.beeper_max_age_ms:
-            encryption_info = await self.state_store.get_encryption_info(evt.content.room_id)
-            if encryption_info:
-                if not evt.content.beeper_max_age_ms:
-                    evt.content.beeper_max_age_ms = encryption_info.rotation_period_ms
-                if not evt.content.beeper_max_messages:
-                    evt.content.beeper_max_messages = encryption_info.rotation_period_msgs
+            await self._fill_encryption_info(evt.content)
         if self.delete_previous_keys_on_receive and not evt.content.beeper_is_scheduled:
             removed_ids = await self.crypto_store.redact_group_sessions(
                 evt.content.room_id, evt.sender_key, reason="received new key from device"