Use more dataclasses for uploading e2ee device keys

Author: tulir

mautrix/client/api/modules/crypto.py

@@ -5,13 +5,15 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Union
 
 from mautrix.api import Method, Path
 from mautrix.errors import MatrixResponseError
 from mautrix.types import (
+    JSON,
     ClaimKeysResponse,
     DeviceID,
+    DeviceKeys,
     EncryptionKeyAlgorithm,
     EventType,
     QueryKeysResponse,
@@ -82,7 +84,7 @@ async def send_to_one_device(
     async def upload_keys(
         self,
         one_time_keys: dict[str, Any] | None = None,
-        device_keys: dict[str, Any] | None = None,
+        device_keys: DeviceKeys | dict[str, Any] | None = None,
     ) -> dict[EncryptionKeyAlgorithm, int]:
         """
         Publishes end-to-end encryption keys for the device.
@@ -102,8 +104,12 @@ async def upload_keys(
         """
         data = {}
         if device_keys:
+            if isinstance(device_keys, Serializable):
+                device_keys = device_keys.serialize()
             data["device_keys"] = device_keys
         if one_time_keys:
+            if isinstance(one_time_keys, Serializable):
+                one_time_keys = one_time_keys.serialize()
             data["one_time_keys"] = one_time_keys
         resp = await self.api.request(Method.POST, Path.v3.keys.upload, data)
         try:

mautrix/crypto/account.py

@@ -10,15 +10,17 @@
 
 from mautrix.types import (
     DeviceID,
+    DeviceKeys,
     EncryptionAlgorithm,
     EncryptionKeyAlgorithm,
     IdentityKey,
+    KeyID,
     SigningKey,
     UserID,
 )
 
-from . import base
 from .sessions import Session
+from .signature import sign_olm
 
 
 class OlmAccount(olm.Account):
@@ -74,19 +76,18 @@ def new_outbound_session(self, target_key: IdentityKey, one_time_key: IdentityKe
             session.pickle("roundtrip"), passphrase="roundtrip", creation_time=datetime.now()
         )
 
-    def get_device_keys(self, user_id: UserID, device_id: DeviceID) -> Dict[str, Any]:
-        device_keys = {
-            "user_id": user_id,
-            "device_id": device_id,
-            "algorithms": [EncryptionAlgorithm.OLM_V1.value, EncryptionAlgorithm.MEGOLM_V1.value],
-            "keys": {
-                f"{algorithm}:{device_id}": key for algorithm, key in self.identity_keys.items()
+    def get_device_keys(self, user_id: UserID, device_id: DeviceID) -> DeviceKeys:
+        device_keys = DeviceKeys(
+            user_id=user_id,
+            device_id=device_id,
+            algorithms=[EncryptionAlgorithm.OLM_V1, EncryptionAlgorithm.MEGOLM_V1],
+            keys={
+                KeyID(algorithm=EncryptionKeyAlgorithm(algorithm), key_id=key): key
+                for algorithm, key in self.identity_keys.items()
             },
-        }
-        signature = self.sign(base.canonical_json(device_keys))
-        device_keys["signatures"] = {
-            user_id: {f"{EncryptionKeyAlgorithm.ED25519}:{device_id}": signature}
-        }
+            signatures={},
+        )
+        device_keys.signatures[user_id] = {KeyID.ed25519(device_id): sign_olm(device_keys, self)}
         return device_keys
 
     def get_one_time_keys(
@@ -97,12 +98,12 @@ def get_one_time_keys(
             self.generate_one_time_keys(new_count)
         keys = {}
         for key_id, key in self.one_time_keys.get("curve25519", {}).items():
-            signature = self.sign(base.canonical_json({"key": key}))
-            keys[f"{EncryptionKeyAlgorithm.SIGNED_CURVE25519}:{key_id}"] = {
+            keys[str(KeyID.signed_curve25519(IdentityKey(key_id)))] = {
                 "key": key,
                 "signatures": {
-                    user_id: {f"{EncryptionKeyAlgorithm.ED25519}:{device_id}": signature}
+                    user_id: {
+                        str(KeyID.ed25519(device_id)): sign_olm({"key": key}, self),
+                    }
                 },
             }
-        self.mark_keys_as_published()
         return keys

mautrix/crypto/machine.py

@@ -313,6 +313,7 @@ async def _share_keys(self, current_otk_count: int | None) -> None:
         self.log.debug(f"Uploading {len(one_time_keys)} one-time keys")
         resp = await self.client.upload_keys(one_time_keys=one_time_keys, device_keys=device_keys)
         self.account.shared = True
+        self.account.mark_keys_as_published()
         self._last_key_share = time.monotonic()
         await self.crypto_store.put_account(self.account)
         self.log.debug(f"Shared keys and saved account, new keys: {resp}")

mautrix/crypto/signature.py

@@ -7,9 +7,19 @@
 import functools
 import json
 
+import olm
 import unpaddedbase64
 
-from mautrix.types import DeviceID, EncryptionKeyAlgorithm, KeyID, SigningKey, UserID
+from mautrix.types import (
+    JSON,
+    DeviceID,
+    EncryptionKeyAlgorithm,
+    KeyID,
+    Serializable,
+    Signature,
+    SigningKey,
+    UserID,
+)
 
 try:
     from Crypto.PublicKey import ECC
@@ -28,6 +38,14 @@ class SignedObject(TypedDict):
     unsigned: Any
 
 
+def sign_olm(data: dict[str, JSON] | Serializable, key: olm.PkSigning | olm.Account) -> Signature:
+    if isinstance(data, Serializable):
+        data = data.serialize()
+    data.pop("signatures", None)
+    data.pop("unsigned", None)
+    return Signature(key.sign(canonical_json(data)))
+
+
 def verify_signature_json(
     data: "SignedObject", user_id: UserID, key_name: DeviceID | str, key: SigningKey
 ) -> bool:

mautrix/types/event/encrypted.py

@@ -9,7 +9,7 @@
 
 from attr import dataclass
 
-from ..primitive import JSON, DeviceID, IdentityKey, SessionID
+from ..primitive import JSON, DeviceID, IdentityKey, SessionID, SigningKey
 from ..util import ExtensibleEnum, Obj, Serializable, SerializableAttrs, deserializer, field
 from .base import BaseRoomEvent, BaseUnsigned
 from .message import RelatesTo
@@ -43,6 +43,18 @@ def deserialize(cls, raw: JSON) -> "KeyID":
     def __str__(self) -> str:
         return f"{self.algorithm.value}:{self.key_id}"
 
+    @classmethod
+    def ed25519(cls, key_id: SigningKey | DeviceID) -> "KeyID":
+        return cls(EncryptionKeyAlgorithm.ED25519, key_id)
+
+    @classmethod
+    def curve25519(cls, key_id: IdentityKey) -> "KeyID":
+        return cls(EncryptionKeyAlgorithm.CURVE25519, key_id)
+
+    @classmethod
+    def signed_curve25519(cls, key_id: IdentityKey) -> "KeyID":
+        return cls(EncryptionKeyAlgorithm.SIGNED_CURVE25519, key_id)
+
 
 class OlmMsgType(Serializable, IntEnum):
     PREKEY = 0