Merge pull request #1614 from famedly/nico/dehydrated-devices-v2

feat: Update dehydrated devices implementation to current MSC
famedly · Nov 28, 2023 · d35872b · d35872b
2 parents 4f95817 + 52a5485
commit d35872b
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 119 deletions.
diff --git a/lib/encryption/encryption.dart b/lib/encryption/encryption.dart
@@ -68,15 +68,20 @@ class Encryption {
   }
 
   // initial login passes null to init a new olm account
-  Future<void> init(String? olmAccount,
-      {String? deviceId,
-      String? pickleKey,
-      bool isDehydratedDevice = false}) async {
+  Future<void> init(
+    String? olmAccount, {
+    String? deviceId,
+    String? pickleKey,
+    String? dehydratedDeviceAlgorithm,
+  }) async {
     ourDeviceId = deviceId ?? client.deviceID!;
+    final isDehydratedDevice = dehydratedDeviceAlgorithm != null;
     await olmManager.init(
-        olmAccount: olmAccount,
-        deviceId: isDehydratedDevice ? deviceId : ourDeviceId,
-        pickleKey: pickleKey);
+      olmAccount: olmAccount,
+      deviceId: isDehydratedDevice ? deviceId : ourDeviceId,
+      pickleKey: pickleKey,
+      dehydratedDeviceAlgorithm: dehydratedDeviceAlgorithm,
+    );
 
     if (!isDehydratedDevice) keyManager.startAutoUploadKeys();
   }

diff --git a/lib/encryption/olm_manager.dart b/lib/encryption/olm_manager.dart
@@ -57,21 +57,25 @@ class OlmManager {
   final Map<String, List<OlmSession>> _olmSessions = {};
 
   // NOTE(Nico): On initial login we pass null to create a new account
-  Future<void> init(
-      {String? olmAccount,
-      required String? deviceId,
-      String? pickleKey}) async {
+  Future<void> init({
+    String? olmAccount,
+    required String? deviceId,
+    String? pickleKey,
+    String? dehydratedDeviceAlgorithm,
+  }) async {
     ourDeviceId = deviceId;
     if (olmAccount == null) {
       try {
         await olm.init();
         _olmAccount = olm.Account();
         _olmAccount!.create();
         if (!await uploadKeys(
-            uploadDeviceKeys: true,
-            updateDatabase: false,
-            // dehydrated devices don't have a device id when created, so skip upload in that case.
-            skipAllUploads: deviceId == null)) {
+          uploadDeviceKeys: true,
+          updateDatabase: false,
+          dehydratedDeviceAlgorithm: dehydratedDeviceAlgorithm,
+          dehydratedDevicePickleKey:
+              dehydratedDeviceAlgorithm != null ? pickleKey : null,
+        )) {
           throw ('Upload key failed');
         }
       } catch (_) {
@@ -131,7 +135,8 @@ class OlmManager {
     int? oldKeyCount = 0,
     bool updateDatabase = true,
     bool? unusedFallbackKey = false,
-    bool skipAllUploads = false,
+    String? dehydratedDeviceAlgorithm,
+    String? dehydratedDevicePickleKey,
     int retry = 1,
   }) async {
     final olmAccount = _olmAccount;
@@ -179,11 +184,6 @@ class OlmManager {
         await encryption.olmDatabase?.updateClientKeys(pickledOlmAccount!);
       }
 
-      if (skipAllUploads) {
-        _uploadKeysLock = false;
-        return true;
-      }
-
       // and now generate the payload to upload
       var deviceKeys = <String, dynamic>{
         'user_id': client.userID,
@@ -239,23 +239,36 @@ class OlmManager {
 
       // Workaround: Make sure we stop if we got logged out in the meantime.
       if (!client.isLogged()) return true;
-      final currentUpload = this.currentUpload =
-          CancelableOperation.fromFuture(ourDeviceId == client.deviceID
-              ? client.uploadKeys(
-                  deviceKeys: uploadDeviceKeys
-                      ? MatrixDeviceKeys.fromJson(deviceKeys)
-                      : null,
-                  oneTimeKeys: signedOneTimeKeys,
-                  fallbackKeys: signedFallbackKeys,
-                )
-              : client.uploadKeysForDevice(
-                  ourDeviceId!,
-                  deviceKeys: uploadDeviceKeys
-                      ? MatrixDeviceKeys.fromJson(deviceKeys)
-                      : null,
-                  oneTimeKeys: signedOneTimeKeys,
-                  fallbackKeys: signedFallbackKeys,
-                ));
+
+      if (ourDeviceId != client.deviceID) {
+        if (dehydratedDeviceAlgorithm == null ||
+            dehydratedDevicePickleKey == null) {
+          throw Exception(
+              'You need to provide both the pickle key and the algorithm to use dehydrated devices!');
+        }
+
+        await client.uploadDehydratedDevice(
+          deviceId: ourDeviceId!,
+          initialDeviceDisplayName: 'Dehydrated Device',
+          deviceKeys:
+              uploadDeviceKeys ? MatrixDeviceKeys.fromJson(deviceKeys) : null,
+          oneTimeKeys: signedOneTimeKeys,
+          fallbackKeys: signedFallbackKeys,
+          deviceData: {
+            'algorithm': dehydratedDeviceAlgorithm,
+            'device': encryption.olmManager
+                .pickleOlmAccountWithKey(dehydratedDevicePickleKey),
+          },
+        );
+        return true;
+      }
+      final currentUpload =
+          this.currentUpload = CancelableOperation.fromFuture(client.uploadKeys(
+        deviceKeys:
+            uploadDeviceKeys ? MatrixDeviceKeys.fromJson(deviceKeys) : null,
+        oneTimeKeys: signedOneTimeKeys,
+        fallbackKeys: signedFallbackKeys,
+      ));
       final response = await currentUpload.valueOrCancellation();
       if (response == null) {
         _uploadKeysLock = false;
@@ -276,8 +289,8 @@ class OlmManager {
       // we failed to upload the keys. If we only tried to upload one time keys, try to recover by removing them and generating new ones.
       if (!uploadDeviceKeys &&
           unusedFallbackKey != false &&
-          !skipAllUploads &&
           retry > 0 &&
+          dehydratedDeviceAlgorithm != null &&
           signedOneTimeKeys.isNotEmpty &&
           exception.error == MatrixError.M_UNKNOWN) {
         Logs().w('Rotating otks because upload failed', exception);
@@ -302,7 +315,6 @@ class OlmManager {
             oldKeyCount: oldKeyCount,
             updateDatabase: updateDatabase,
             unusedFallbackKey: unusedFallbackKey,
-            skipAllUploads: skipAllUploads,
             retry: retry - 1);
       }
     } finally {

diff --git a/lib/msc_extensions/msc_3814_dehydrated_devices/api.dart b/lib/msc_extensions/msc_3814_dehydrated_devices/api.dart
@@ -29,41 +29,29 @@ import 'package:matrix/msc_extensions/msc_3814_dehydrated_devices/model/dehydrat
 /// Endpoints related to MSC3814, dehydrated devices v2 aka shrivelled sessions
 /// https://github.com/matrix-org/matrix-spec-proposals/pull/3814
 extension DehydratedDeviceMatrixApi on MatrixApi {
-  /// Publishes end-to-end encryption keys for the specified device.
-  /// https://github.com/matrix-org/matrix-spec-proposals/pull/3814
-  Future<Map<String, int>> uploadKeysForDevice(String device,
-      {MatrixDeviceKeys? deviceKeys,
-      Map<String, dynamic>? oneTimeKeys,
-      Map<String, dynamic>? fallbackKeys}) async {
-    final response = await request(
-      RequestType.POST,
-      '/client/v3/keys/upload/${Uri.encodeComponent(device)}',
-      data: {
-        if (deviceKeys != null) 'device_keys': deviceKeys.toJson(),
-        if (oneTimeKeys != null) 'one_time_keys': oneTimeKeys,
-        if (fallbackKeys != null) ...{
-          'fallback_keys': fallbackKeys,
-          'org.matrix.msc2732.fallback_keys': fallbackKeys,
-        },
-      },
-    );
-    return Map<String, int>.from(
-        response.tryGetMap<String, Object?>('one_time_key_counts') ??
-            <String, int>{});
-  }
-
   /// uploads a dehydrated device.
   /// https://github.com/matrix-org/matrix-spec-proposals/pull/3814
-  Future<String> uploadDehydratedDevice(
-      {String? initialDeviceDisplayName,
-      Map<String, dynamic>? deviceData}) async {
+  Future<String> uploadDehydratedDevice({
+    required String deviceId,
+    String? initialDeviceDisplayName,
+    Map<String, dynamic>? deviceData,
+    MatrixDeviceKeys? deviceKeys,
+    Map<String, dynamic>? oneTimeKeys,
+    Map<String, dynamic>? fallbackKeys,
+  }) async {
     final response = await request(
       RequestType.PUT,
       '/client/unstable/org.matrix.msc3814.v1/dehydrated_device',
       data: {
+        'device_id': deviceId,
         if (initialDeviceDisplayName != null)
           'initial_device_display_name': initialDeviceDisplayName,
         if (deviceData != null) 'device_data': deviceData,
+        if (deviceKeys != null) 'device_keys': deviceKeys.toJson(),
+        if (oneTimeKeys != null) 'one_time_keys': oneTimeKeys,
+        if (fallbackKeys != null) ...{
+          'fallback_keys': fallbackKeys,
+        },
       },
     );
     return response['device_id'] as String;
@@ -82,15 +70,15 @@ extension DehydratedDeviceMatrixApi on MatrixApi {
   /// fetch events sent to a dehydrated device.
   /// https://github.com/matrix-org/matrix-spec-proposals/pull/3814
   Future<DehydratedDeviceEvents> getDehydratedDeviceEvents(String deviceId,
-      {String? from, int limit = 100}) async {
-    final response = await request(
-      RequestType.GET,
-      '/client/unstable/org.matrix.msc3814.v1/dehydrated_device/$deviceId/events',
-      query: {
-        if (from != null) 'from': from,
-        'limit': limit.toString(),
-      },
-    );
+      {String? nextBatch, int limit = 100}) async {
+    final response = await request(RequestType.POST,
+        '/client/unstable/org.matrix.msc3814.v1/dehydrated_device/$deviceId/events',
+        query: {
+          'limit': limit.toString(),
+        },
+        data: {
+          if (nextBatch != null) 'next_batch': nextBatch,
+        });
     return DehydratedDeviceEvents.fromJson(response);
   }
 }
diff --git a/lib/msc_extensions/msc_3814_dehydrated_devices/msc_3814_dehydrated_devices.dart b/lib/msc_extensions/msc_3814_dehydrated_devices/msc_3814_dehydrated_devices.dart
@@ -1,6 +1,7 @@
 library msc_3814_dehydrated_devices;
 
 import 'dart:convert';
+import 'dart:math';
 
 import 'package:matrix/encryption.dart';
 import 'package:matrix/matrix.dart';
@@ -78,10 +79,12 @@ extension DehydratedDeviceHandler on Client {
       // We need to be careful to not use the client.deviceId here and such.
       final encryption = Encryption(client: this);
       try {
-        await encryption.init(pickledDevice,
-            deviceId: device.deviceId,
-            pickleKey: pickleDeviceKey,
-            isDehydratedDevice: true);
+        await encryption.init(
+          pickledDevice,
+          deviceId: device.deviceId,
+          pickleKey: pickleDeviceKey,
+          dehydratedDeviceAlgorithm: _dehydratedDeviceAlgorithm,
+        );
 
         if (dehydratedDeviceIdentity.curve25519Key != encryption.identityKey ||
             dehydratedDeviceIdentity.ed25519Key != encryption.fingerprintKey) {
@@ -97,7 +100,7 @@ extension DehydratedDeviceHandler on Client {
 
         do {
           events = await getDehydratedDeviceEvents(device.deviceId,
-              from: events?.nextBatch);
+              nextBatch: events?.nextBatch);
 
           for (final e in events.events ?? []) {
             // We are only interested in roomkeys, which ALWAYS need to be encrypted.
@@ -111,6 +114,12 @@ extension DehydratedDeviceHandler on Client {
           }
         } while (events.events?.isNotEmpty == true);
 
+        // make sure the sessions we just received get uploaded before we upload a new device (which deletes the old device).
+        await this
+            .encryption
+            ?.keyManager
+            .uploadInboundGroupSessions(skipIfInProgress: false);
+
         await _uploadNewDevice(secureStorage);
       } finally {
         await encryption.dispose();
@@ -136,18 +145,22 @@ extension DehydratedDeviceHandler on Client {
             _ssssSecretNameForDehydratedDevice, pickleDeviceKey);
       }
 
+      const chars =
+          'AaBbCcDdEeFfGgHhIiJjKkLlMmNnOoPpQqRrSsTtUuVvWwXxYyZz1234567890';
+      final rnd = Random();
+
+      final deviceIdSuffix = String.fromCharCodes(Iterable.generate(
+          10, (_) => chars.codeUnitAt(rnd.nextInt(chars.length))));
+      final String device = 'FAM$deviceIdSuffix';
+
       // Generate a new olm account for the dehydrated device.
-      await encryption.init(null,
-          deviceId: null, isDehydratedDevice: true, pickleKey: pickleDeviceKey);
-      String device;
       try {
-        device = await uploadDehydratedDevice(
-            initialDeviceDisplayName: 'Dehydrated Device',
-            deviceData: {
-              'algorithm': _dehydratedDeviceAlgorithm,
-              'device': encryption.olmManager
-                  .pickleOlmAccountWithKey(pickleDeviceKey),
-            });
+        await encryption.init(
+          null,
+          deviceId: device,
+          pickleKey: pickleDeviceKey,
+          dehydratedDeviceAlgorithm: _dehydratedDeviceAlgorithm,
+        );
       } on MatrixException catch (_) {
         // dehydrated devices unsupported, do noting.
         Logs().i('Dehydrated devices unsupported, skipping upload.');
@@ -158,11 +171,6 @@ extension DehydratedDeviceHandler on Client {
       encryption.ourDeviceId = device;
       encryption.olmManager.ourDeviceId = device;
 
-      await encryption.olmManager.uploadKeys(
-          uploadDeviceKeys: true,
-          updateDatabase: false,
-          unusedFallbackKey: true);
-
       // cross sign the device from our currently signed in device
       await updateUserDeviceKeys(additionalUsers: {userID!});
       final keysToSign = <SignableKey>[