element-hq · stefanceriu · Aug 12, 2023 · Jul 27, 2023
diff --git a/Riot/Model/HomeserverConfiguration/HomeserverConfigurationBuilder.swift b/Riot/Model/HomeserverConfiguration/HomeserverConfigurationBuilder.swift
@@ -50,11 +50,14 @@ final class HomeserverConfigurationBuilder: NSObject {
         } else {
             secureBackupSetupMethods = VectorWellKnownBackupSetupMethod.allCases
         }
+
+        let deviceDehydrationEnabled = wellKnown?.jsonDictionary()["org.matrix.msc3814"] as? Bool == true
 
         let encryptionConfiguration = HomeserverEncryptionConfiguration(isE2EEByDefaultEnabled: isE2EEByDefaultEnabled,
                                                                         isSecureBackupRequired: isSecureBackupRequired,
                                                                         secureBackupSetupMethods: secureBackupSetupMethods,
-                                                                        outboundKeysPreSharingMode: outboundKeysPreSharingMode)
+                                                                        outboundKeysPreSharingMode: outboundKeysPreSharingMode,
+                                                                        deviceDehydrationEnabled: deviceDehydrationEnabled)
 
         // Jitsi configuration
         let jitsiPreferredDomain: String?

diff --git a/Riot/Model/HomeserverConfiguration/HomeserverEncryptionConfiguration.swift b/Riot/Model/HomeserverConfiguration/HomeserverEncryptionConfiguration.swift
@@ -23,15 +23,18 @@ final class HomeserverEncryptionConfiguration: NSObject {
     let isSecureBackupRequired: Bool
     let secureBackupSetupMethods: [VectorWellKnownBackupSetupMethod]
     let outboundKeysPreSharingMode: MXKKeyPreSharingStrategy
+    let deviceDehydrationEnabled: Bool
 
     init(isE2EEByDefaultEnabled: Bool,
          isSecureBackupRequired: Bool,
          secureBackupSetupMethods: [VectorWellKnownBackupSetupMethod],
-         outboundKeysPreSharingMode: MXKKeyPreSharingStrategy) {
+         outboundKeysPreSharingMode: MXKKeyPreSharingStrategy,
+         deviceDehydrationEnabled: Bool) {
         self.isE2EEByDefaultEnabled = isE2EEByDefaultEnabled
         self.isSecureBackupRequired = isSecureBackupRequired
         self.outboundKeysPreSharingMode = outboundKeysPreSharingMode
         self.secureBackupSetupMethods = secureBackupSetupMethods
+        self.deviceDehydrationEnabled = deviceDehydrationEnabled
 
         super.init()
     }

diff --git a/Riot/Modules/MatrixKit/Controllers/MXKAuthenticationViewController.h b/Riot/Modules/MatrixKit/Controllers/MXKAuthenticationViewController.h
@@ -305,15 +305,5 @@
  */
 - (void)showAuthenticationFallBackView;
 
-#pragma mark - Device rehydration
-
-/**
- Call this method at an appropriate time to attempt rehydrating from an existing dehydrated device
- @param keyData Secret key data
- @param credentials Account credentials
- */
-
-- (void)attemptDeviceRehydrationWithKeyData:(NSData *)keyData credentials:(MXCredentials *)credentials;
-
 @end
 
diff --git a/Riot/Modules/MatrixKit/Controllers/MXKAuthenticationViewController.m b/Riot/Modules/MatrixKit/Controllers/MXKAuthenticationViewController.m
@@ -1495,68 +1495,6 @@ - (void)createAccountWithCredentials:(MXCredentials *)credentials
     [self _createAccountWithCredentials:credentials];
 }
 
-- (void)attemptDeviceRehydrationWithKeyData:(NSData *)keyData
-                                credentials:(MXCredentials *)credentials
-{
-    [self attemptDeviceRehydrationWithKeyData:keyData
-                                  credentials:credentials
-                                        retry:YES];
-}
-
-- (void)attemptDeviceRehydrationWithKeyData:(NSData *)keyData
-                                credentials:(MXCredentials *)credentials
-                                      retry:(BOOL)retry
-{
-    MXLogDebug(@"[MXKAuthenticationViewController] attemptDeviceRehydration: starting device rehydration");
-
-    if (keyData == nil)
-    {
-        MXLogError(@"[MXKAuthenticationViewController] attemptDeviceRehydration: no key provided for device rehydration");
-        [self _createAccountWithCredentials:credentials];
-        return;
-    }
-
-    MXRestClient *mxRestClient = [[MXRestClient alloc] initWithCredentials:credentials andOnUnrecognizedCertificateBlock:^BOOL(NSData *certificate) {
-        return NO;
-    } andPersistentTokenDataHandler:^(void (^handler)(NSArray<MXCredentials *> *credentials, void (^completion)(BOOL didUpdateCredentials))) {
-        [[MXKAccountManager sharedManager] readAndWriteCredentials:handler];
-    } andUnauthenticatedHandler: nil];
-
-    MXWeakify(self);
-    [[MXKAccountManager sharedManager].dehydrationService rehydrateDeviceWithMatrixRestClient:mxRestClient dehydrationKey:keyData success:^(NSString * deviceId) {
-        MXStrongifyAndReturnIfNil(self);
-
-        if (deviceId)
-        {
-            MXLogDebug(@"[MXKAuthenticationViewController] attemptDeviceRehydration: device %@ rehydrated successfully.", deviceId);
-            credentials.deviceId = deviceId;
-        }
-        else
-        {
-            MXLogDebug(@"[MXKAuthenticationViewController] attemptDeviceRehydration: device rehydration has been canceled.");
-        }
-
-        [self _createAccountWithCredentials:credentials];
-    } failure:^(NSError *error) {
-        MXStrongifyAndReturnIfNil(self);
-
-        if (retry)
-        {
-            MXLogErrorDetails(@"[MXKAuthenticationViewController] attemptDeviceRehydration: device rehydration failed due to error: Retrying", @{
-                @"error": error ?: @"unknown"
-            });
-            [self attemptDeviceRehydrationWithKeyData:keyData credentials:credentials retry:NO];
-            return;
-        }
-
-        MXLogErrorDetails(@"[MXKAuthenticationViewController] attemptDeviceRehydration: device rehydration failed due to error", @{
-            @"error": error ?: @"unknown"
-        });
-
-        [self _createAccountWithCredentials:credentials];
-    }];
-}
-
 - (void)_createAccountWithCredentials:(MXCredentials *)credentials
 {
     MXKAccount *account = [[MXKAccount alloc] initWithCredentials:credentials];

diff --git a/Riot/Modules/MatrixKit/Models/Account/MXKAccount.h b/Riot/Modules/MatrixKit/Models/Account/MXKAccount.h
@@ -360,13 +360,6 @@ typedef BOOL (^MXKAccountOnCertificateChange)(MXKAccount *mxAccount, NSData *cer
 
 #pragma mark - Sync filter
 
-/**
- Call this method at an appropriate time to attempt dehydrating to a new backup device
- */
-- (void)attemptDeviceDehydrationWithKeyData:(NSData *)keyData
-                                    success:(void (^)(void))success
-                                    failure:(void (^)(NSError *error))failure;
-
 /**
  Handle unauthenticated errors from  the server triggering hard/soft logouts as appropriate.
  */

diff --git a/Riot/Modules/MatrixKit/Models/Account/MXKAccount.m b/Riot/Modules/MatrixKit/Models/Account/MXKAccount.m
@@ -1715,70 +1715,6 @@ - (void)launchInitialServerSync
     }];
 }
 
-- (void)attemptDeviceDehydrationWithKeyData:(NSData *)keyData
-                                    success:(void (^)(void))success
-                                    failure:(void (^)(NSError *error))failure
-{
-    [self attemptDeviceDehydrationWithKeyData:keyData retry:YES success:success failure:failure];
-}
-
-- (void)attemptDeviceDehydrationWithKeyData:(NSData *)keyData
-                                      retry:(BOOL)retry
-                                    success:(void (^)(void))success
-                                    failure:(void (^)(NSError *error))failure
-{
-    if (keyData == nil)
-    {
-        MXLogWarning(@"[MXKAccount] attemptDeviceDehydrationWithRetry: no key provided for device dehydration");
-
-        if (failure)
-        {
-            failure(nil);
-        }
-
-        return;
-    }
-
-    if (![mxSession.crypto.crossSigning isKindOfClass:[MXLegacyCrossSigning class]]) {
-        MXLogFailure(@"Device dehydratation is currently only supported by legacy cross signing, add support to all implementations");
-        if (failure)
-        {
-            failure(nil);
-        }
-        return;
-    }
-    MXLegacyCrossSigning *crossSigning = (MXLegacyCrossSigning *)mxSession.crypto.crossSigning;;
-
-    MXLogDebug(@"[MXKAccount] attemptDeviceDehydrationWithRetry: starting device dehydration");
-    [[MXKAccountManager sharedManager].dehydrationService dehydrateDeviceWithMatrixRestClient:mxRestClient crossSigning:crossSigning dehydrationKey:keyData success:^(NSString *deviceId) {
-        MXLogDebug(@"[MXKAccount] attemptDeviceDehydrationWithRetry: device successfully dehydrated");
-
-        if (success)
-        {
-            success();
-        }
-    } failure:^(NSError *error) {
-        if (retry)
-        {
-            [self attemptDeviceDehydrationWithKeyData:keyData retry:NO success:success failure:failure];
-            MXLogErrorDetails(@"[MXKAccount] attemptDeviceDehydrationWithRetry: device dehydration failed due to error: Retrying.", @{
-                @"error": error ?: @"unknown"
-            });
-        }
-        else
-        {
-            MXLogErrorDetails(@"[MXKAccount] attemptDeviceDehydrationWithRetry: device dehydration failed due to error", @{
-                @"error": error ?: @"unknown"
-            });
-
-            if (failure)
-            {
-                failure(error);
-            }
-        }
-    }];
-}
-
 - (void)onMatrixSessionStateChange
 {
     // Check if pause has been requested

diff --git a/Riot/Modules/MatrixKit/Models/Account/MXKAccountManager.h b/Riot/Modules/MatrixKit/Models/Account/MXKAccountManager.h
@@ -104,8 +104,6 @@ extern NSString *const MXKAccountManagerDataType;
  */
 @property (nonatomic) BOOL isPushAvailable;
 
-@property (nonatomic, readonly) MXDehydrationService *dehydrationService;
-
 /**
  Retrieve the MXKAccounts manager.
 

diff --git a/Riot/Modules/MatrixKit/Models/Account/MXKAccountManager.m b/Riot/Modules/MatrixKit/Models/Account/MXKAccountManager.m
@@ -71,7 +71,6 @@ - (instancetype)init
     if (self)
     {
         _storeClass = [MXFileStore class];
-        _dehydrationService = [MXDehydrationService new];
         _savingAccountsEnabled = YES;
 
         // Migrate old account file to new format

diff --git a/Riot/Modules/Secrets/Recover/RecoverWithKey/SecretsRecoveryWithKeyCoordinator.swift b/Riot/Modules/Secrets/Recover/RecoverWithKey/SecretsRecoveryWithKeyCoordinator.swift
@@ -34,9 +34,9 @@ final class SecretsRecoveryWithKeyCoordinator: SecretsRecoveryWithKeyCoordinator
 
     // MARK: - Setup
 
-    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal, cancellable: Bool) {
+    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal, cancellable: Bool, dehydrationService: DehydrationService?) {
 
-        let secretsRecoveryWithKeyViewModel = SecretsRecoveryWithKeyViewModel(recoveryService: recoveryService, recoveryGoal: recoveryGoal)
+        let secretsRecoveryWithKeyViewModel = SecretsRecoveryWithKeyViewModel(recoveryService: recoveryService, recoveryGoal: recoveryGoal, dehydrationService: dehydrationService)
         let secretsRecoveryWithKeyViewController = SecretsRecoveryWithKeyViewController.instantiate(with: secretsRecoveryWithKeyViewModel, cancellable: cancellable)
         self.secretsRecoveryWithKeyViewController = secretsRecoveryWithKeyViewController
         self.secretsRecoveryWithKeyViewModel = secretsRecoveryWithKeyViewModel

diff --git a/Riot/Modules/Secrets/Recover/RecoverWithKey/SecretsRecoveryWithKeyViewModel.swift b/Riot/Modules/Secrets/Recover/RecoverWithKey/SecretsRecoveryWithKeyViewModel.swift
@@ -24,6 +24,8 @@ final class SecretsRecoveryWithKeyViewModel: SecretsRecoveryWithKeyViewModelType
 
     private let recoveryService: MXRecoveryService
 
+    private let dehydrationService: DehydrationService?
+
     // MARK: Public
 
     let recoveryGoal: SecretsRecoveryGoal
@@ -39,8 +41,9 @@ final class SecretsRecoveryWithKeyViewModel: SecretsRecoveryWithKeyViewModelType
 
     // MARK: - Setup
 
-    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal) {
+    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal, dehydrationService: DehydrationService?) {
         self.recoveryService = recoveryService
+        self.dehydrationService = dehydrationService
         self.recoveryGoal = recoveryGoal
     }
 
@@ -83,6 +86,10 @@ final class SecretsRecoveryWithKeyViewModel: SecretsRecoveryWithKeyViewModelType
                 }
                 self.update(viewState: .loaded)
                 self.coordinatorDelegate?.secretsRecoveryWithKeyViewModelDidRecover(self)
+
+                Task {
+                    await self.dehydrationService?.runDeviceDehydrationFlow(privateKeyData: privateKey)
+                }
             }, failure: { [weak self] error in
                 guard let self = self else {
                     return

diff --git a/...ules/Secrets/Recover/RecoverWithPassphrase/SecretsRecoveryWithPassphraseCoordinator.swift b/...ules/Secrets/Recover/RecoverWithPassphrase/SecretsRecoveryWithPassphraseCoordinator.swift
@@ -34,8 +34,8 @@ final class SecretsRecoveryWithPassphraseCoordinator: SecretsRecoveryWithPassphr
 
     // MARK: - Setup
 
-    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal, cancellable: Bool) {
-        let secretsRecoveryWithPassphraseViewModel = SecretsRecoveryWithPassphraseViewModel(recoveryService: recoveryService, recoveryGoal: recoveryGoal)
+    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal, cancellable: Bool, dehydrationService: DehydrationService?) {
+        let secretsRecoveryWithPassphraseViewModel = SecretsRecoveryWithPassphraseViewModel(recoveryService: recoveryService, recoveryGoal: recoveryGoal, dehydrationService: dehydrationService)
         let secretsRecoveryWithPassphraseViewController = SecretsRecoveryWithPassphraseViewController.instantiate(with: secretsRecoveryWithPassphraseViewModel, cancellable: cancellable)
         self.secretsRecoveryWithPassphraseViewController = secretsRecoveryWithPassphraseViewController
         self.secretsRecoveryWithPassphraseViewModel = secretsRecoveryWithPassphraseViewModel

diff --git a/...odules/Secrets/Recover/RecoverWithPassphrase/SecretsRecoveryWithPassphraseViewModel.swift b/...odules/Secrets/Recover/RecoverWithPassphrase/SecretsRecoveryWithPassphraseViewModel.swift
@@ -24,6 +24,8 @@ final class SecretsRecoveryWithPassphraseViewModel: SecretsRecoveryWithPassphras
 
     private let recoveryService: MXRecoveryService
 
+    private let dehydrationService: DehydrationService?
+
     // MARK: Public
 
     let recoveryGoal: SecretsRecoveryGoal
@@ -39,8 +41,9 @@ final class SecretsRecoveryWithPassphraseViewModel: SecretsRecoveryWithPassphras
 
     // MARK: - Setup
 
-    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal) {
+    init(recoveryService: MXRecoveryService, recoveryGoal: SecretsRecoveryGoal, dehydrationService: DehydrationService?) {
         self.recoveryService = recoveryService
+        self.dehydrationService = dehydrationService
         self.recoveryGoal = recoveryGoal
     }
 
@@ -103,6 +106,10 @@ final class SecretsRecoveryWithPassphraseViewModel: SecretsRecoveryWithPassphras
             }
             self.update(viewState: .loaded)
             self.coordinatorDelegate?.secretsRecoveryWithPassphraseViewModelDidRecover(self)
+
+            Task {
+                await self.dehydrationService?.runDeviceDehydrationFlow(privateKeyData: privateKey)
+            }
         }, failure: { [weak self] error in
             guard let self = self else {
                 return

diff --git a/Riot/Modules/Secrets/Recover/SecretsRecoveryCoordinator.swift b/Riot/Modules/Secrets/Recover/SecretsRecoveryCoordinator.swift
@@ -85,14 +85,28 @@ final class SecretsRecoveryCoordinator: SecretsRecoveryCoordinatorType {
 
     // MARK: - Private
 
+    private var dehydrationService: DehydrationService? {
+        if self.session.vc_homeserverConfiguration().encryption.deviceDehydrationEnabled {
+            return self.session.crypto.dehydrationService
+        }
+
+        return nil
+    }
+
     private func createRecoverFromKeyCoordinator() -> SecretsRecoveryWithKeyCoordinator {
-        let coordinator = SecretsRecoveryWithKeyCoordinator(recoveryService: self.session.crypto.recoveryService, recoveryGoal: self.recoveryGoal, cancellable: self.cancellable)
+        let coordinator = SecretsRecoveryWithKeyCoordinator(recoveryService: self.session.crypto.recoveryService,
+                                                            recoveryGoal: self.recoveryGoal,
+                                                            cancellable: self.cancellable,
+                                                            dehydrationService: dehydrationService)
         coordinator.delegate = self
         return coordinator
     }
 
     private func createRecoverFromPassphraseCoordinator() -> SecretsRecoveryWithPassphraseCoordinator {
-        let coordinator = SecretsRecoveryWithPassphraseCoordinator(recoveryService: self.session.crypto.recoveryService, recoveryGoal: self.recoveryGoal, cancellable: self.cancellable)
+        let coordinator = SecretsRecoveryWithPassphraseCoordinator(recoveryService: self.session.crypto.recoveryService,
+                                                                   recoveryGoal: self.recoveryGoal,
+                                                                   cancellable: self.cancellable,
+                                                                   dehydrationService: dehydrationService)
         coordinator.delegate = self
         return coordinator
     }

diff --git a/Riot/Modules/Secrets/Setup/RecoveryKey/SecretsSetupRecoveryKeyCoordinator.swift b/Riot/Modules/Secrets/Setup/RecoveryKey/SecretsSetupRecoveryKeyCoordinator.swift
@@ -42,8 +42,13 @@ final class SecretsSetupRecoveryKeyCoordinator: SecretsSetupRecoveryKeyCoordinat
          passphrase: String?,
          passphraseOnly: Bool,
          allowOverwrite: Bool = false,
-         cancellable: Bool) {
-        let secretsSetupRecoveryKeyViewModel = SecretsSetupRecoveryKeyViewModel(recoveryService: recoveryService, passphrase: passphrase, passphraseOnly: passphraseOnly, allowOverwrite: allowOverwrite)
+         cancellable: Bool,
+         dehydrationService: DehydrationService?) {
+        let secretsSetupRecoveryKeyViewModel = SecretsSetupRecoveryKeyViewModel(recoveryService: recoveryService,
+                                                                                passphrase: passphrase,
+                                                                                passphraseOnly: passphraseOnly,
+                                                                                allowOverwrite: allowOverwrite,
+                                                                                dehydrationService: dehydrationService)
         let secretsSetupRecoveryKeyViewController = SecretsSetupRecoveryKeyViewController.instantiate(with: secretsSetupRecoveryKeyViewModel, cancellable: cancellable)
         self.secretsSetupRecoveryKeyViewModel = secretsSetupRecoveryKeyViewModel
         self.secretsSetupRecoveryKeyViewController = secretsSetupRecoveryKeyViewController

diff --git a/Riot/Modules/Secrets/Setup/RecoveryKey/SecretsSetupRecoveryKeyViewModel.swift b/Riot/Modules/Secrets/Setup/RecoveryKey/SecretsSetupRecoveryKeyViewModel.swift
@@ -28,6 +28,7 @@ final class SecretsSetupRecoveryKeyViewModel: SecretsSetupRecoveryKeyViewModelTy
     private let passphrase: String?
     private let passphraseOnly: Bool
     private let allowOverwrite: Bool
+    private let dehydrationService: DehydrationService?
 
     // MARK: Public
 
@@ -36,11 +37,12 @@ final class SecretsSetupRecoveryKeyViewModel: SecretsSetupRecoveryKeyViewModelTy
 
     // MARK: - Setup
 
-    init(recoveryService: MXRecoveryService, passphrase: String?, passphraseOnly: Bool, allowOverwrite: Bool = false) {
+    init(recoveryService: MXRecoveryService, passphrase: String?, passphraseOnly: Bool, allowOverwrite: Bool = false, dehydrationService: DehydrationService?) {
         self.recoveryService = recoveryService
         self.passphrase = passphrase
         self.passphraseOnly = passphraseOnly
         self.allowOverwrite = allowOverwrite
+        self.dehydrationService = dehydrationService
     }
 
     // MARK: - Public
@@ -76,6 +78,10 @@ final class SecretsSetupRecoveryKeyViewModel: SecretsSetupRecoveryKeyViewModelTy
 
         self.recoveryService.createRecovery(forSecrets: nil, withPassphrase: self.passphrase, createServicesBackups: true, success: { secretStorageKeyCreationInfo in
             self.update(viewState: .recoveryCreated(secretStorageKeyCreationInfo.recoveryKey))
+
+            Task {
+                await self.dehydrationService?.runDeviceDehydrationFlow(privateKeyData: secretStorageKeyCreationInfo.privateKey)
+            }
         }, failure: { error in
             self.update(viewState: .error(error))
         })