Schedule dehydration on reload if the dehydration key is already cached locally (#29021)

* Schedule dehydration on reload

* fix test and use the right function to check dehydration is enabled

* use dehydration helper function when scheduling dehydration on restart

* fix test by passing in client object

src/MatrixClientPeg.ts

@@ -41,6 +41,7 @@ import PlatformPeg from "./PlatformPeg";
 import { formatList } from "./utils/FormattingUtils";
 import SdkConfig from "./SdkConfig";
 import { setDeviceIsolationMode } from "./settings/controllers/DeviceIsolationModeController.ts";
+import { initialiseDehydration } from "./utils/device/dehydration";
 
 export interface IMatrixClientCreds {
     homeserverUrl: string;
@@ -340,7 +341,20 @@ class MatrixClientPegClass implements IMatrixClientPeg {
 
         setDeviceIsolationMode(this.matrixClient, SettingsStore.getValue("feature_exclude_insecure_devices"));
 
-        // TODO: device dehydration and whathaveyou
+        // Start dehydration. This code is only for the case where the client
+        // gets restarted, so we only do this if we already have the dehydration
+        // key cached, and we don't have to try to rehydrate a device. If this
+        // is a new login, we will start dehydration after Secret Storage is
+        // unlocked.
+        try {
+            await initialiseDehydration({ onlyIfKeyCached: true, rehydrate: false }, this.matrixClient);
+        } catch (e) {
+            // We may get an error dehydrating, such as if cross-signing and
+            // SSSS are not set up yet.  Just log the error and continue.
+            // If SSSS gets set up later, we will re-try dehydration.
+            console.log("Error starting device dehydration", e);
+        }
+
         return;
     }
 

src/async-components/views/dialogs/security/CreateSecretStorageDialog.tsx

@@ -332,7 +332,7 @@ export default class CreateSecretStorageDialog extends React.PureComponent<IProp
                     setupNewKeyBackup: !backupInfo,
                 });
             }
-            await initialiseDehydration(true);
+            await initialiseDehydration({ createNewKey: true });
 
             this.setState({
                 phase: Phase.Stored,

src/utils/device/dehydration.ts

@@ -7,8 +7,9 @@ Please see LICENSE files in the repository root for full details.
 */
 
 import { logger } from "matrix-js-sdk/src/logger";
-import { CryptoApi } from "matrix-js-sdk/src/crypto-api";
+import { CryptoApi, StartDehydrationOpts } from "matrix-js-sdk/src/crypto-api";
 
+import type { MatrixClient } from "matrix-js-sdk/src/matrix";
 import { MatrixClientPeg } from "../../MatrixClientPeg";
 
 /**
@@ -21,14 +22,14 @@ import { MatrixClientPeg } from "../../MatrixClientPeg";
  *
  * Dehydration can currently only be enabled by setting a flag in the .well-known file.
  */
-async function deviceDehydrationEnabled(crypto: CryptoApi | undefined): Promise<boolean> {
+async function deviceDehydrationEnabled(client: MatrixClient, crypto: CryptoApi | undefined): Promise<boolean> {
     if (!crypto) {
         return false;
     }
     if (!(await crypto.isDehydrationSupported())) {
         return false;
     }
-    const wellknown = await MatrixClientPeg.safeGet().waitForClientWellKnown();
+    const wellknown = await client.waitForClientWellKnown();
     return !!wellknown?.["org.matrix.msc3814"];
 }
 
@@ -40,10 +41,11 @@ async function deviceDehydrationEnabled(crypto: CryptoApi | undefined): Promise<
  * @param createNewKey: force a new dehydration key to be created, even if one
  *   already exists.  This is used when we reset secret storage.
  */
-export async function initialiseDehydration(createNewKey: boolean = false): Promise<void> {
-    const crypto = MatrixClientPeg.safeGet().getCrypto();
-    if (await deviceDehydrationEnabled(crypto)) {
+export async function initialiseDehydration(opts: StartDehydrationOpts = {}, client?: MatrixClient): Promise<void> {
+    client = client || MatrixClientPeg.safeGet();
+    const crypto = client.getCrypto();
+    if (await deviceDehydrationEnabled(client, crypto)) {
         logger.log("Device dehydration enabled");
-        await crypto!.startDehydration(createNewKey);
+        await crypto!.startDehydration(opts);
     }
 }

test/unit-tests/MatrixClientPeg-test.ts

@@ -86,6 +86,27 @@ describe("MatrixClientPeg", () => {
             expect(mockInitRustCrypto).toHaveBeenCalledWith({ storageKey: cryptoStoreKey });
         });
 
+        it("should try to start dehydration if dehydration is enabled", async () => {
+            const mockInitRustCrypto = jest.spyOn(testPeg.safeGet(), "initRustCrypto").mockResolvedValue(undefined);
+            const mockStartDehydration = jest.fn();
+            jest.spyOn(testPeg.safeGet(), "getCrypto").mockReturnValue({
+                isDehydrationSupported: jest.fn().mockResolvedValue(true),
+                startDehydration: mockStartDehydration,
+                setDeviceIsolationMode: jest.fn(),
+            } as any);
+            jest.spyOn(testPeg.safeGet(), "waitForClientWellKnown").mockResolvedValue({
+                "m.homeserver": {
+                    base_url: "http://example.com",
+                },
+                "org.matrix.msc3814": true,
+            } as any);
+
+            const cryptoStoreKey = new Uint8Array([1, 2, 3, 4]);
+            await testPeg.start({ rustCryptoStoreKey: cryptoStoreKey });
+            expect(mockInitRustCrypto).toHaveBeenCalledWith({ storageKey: cryptoStoreKey });
+            expect(mockStartDehydration).toHaveBeenCalledWith({ onlyIfKeyCached: true, rehydrate: false });
+        });
+
         it("Should migrate existing login", async () => {
             const mockInitRustCrypto = jest.spyOn(testPeg.safeGet(), "initRustCrypto").mockResolvedValue(undefined);
 

test/unit-tests/components/structures/MatrixChat-test.tsx

@@ -128,6 +128,7 @@ describe("<MatrixChat />", () => {
         getCrypto: jest.fn().mockReturnValue({
             getVerificationRequestsToDeviceInProgress: jest.fn().mockReturnValue([]),
             isCrossSigningReady: jest.fn().mockReturnValue(false),
+            isDehydrationSupported: jest.fn().mockReturnValue(false),
             getUserDeviceInfo: jest.fn().mockReturnValue(new Map()),
             getUserVerificationStatus: jest.fn().mockResolvedValue(new UserVerificationStatus(false, false, false)),
             getVersion: jest.fn().mockReturnValue("1"),
@@ -1005,6 +1006,7 @@ describe("<MatrixChat />", () => {
                     resetKeyBackup: jest.fn(),
                     isEncryptionEnabledInRoom: jest.fn().mockResolvedValue(false),
                     checkKeyBackupAndEnable: jest.fn().mockResolvedValue(null),
+                    isDehydrationSupported: jest.fn().mockReturnValue(false),
                 };
                 loginClient.getCrypto.mockReturnValue(mockCrypto as any);
             });