Add support for device dehydration v2 (Element R) (#4062)

* initial implementation of device dehydration

* add dehydrated flag for devices

* add missing dehydration.ts file, add test, add function to schedule dehydration

* add more dehydration utility functions

* stop scheduled dehydration when crypto stops

* bump matrix-crypto-sdk-wasm version, and fix tests

* adding dehydratedDevices member to mock OlmDevice isn't necessary any more

* fix yarn lock file

* more tests

* fix test

* more tests

* fix typo

* fix logic for checking if dehydration supported

* make changes from review

* add missing file

* move setup into another function

* apply changes from review

* implement simpler API

* fix type and move the code to the right spot

* apply suggestions from review

* make sure that cross-signing and secret storage are set up

spec/integ/crypto/device-dehydration.spec.ts

@@ -0,0 +1,181 @@
+/*
+Copyright 2024 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import "fake-indexeddb/auto";
+import fetchMock from "fetch-mock-jest";
+
+import { createClient, ClientEvent, MatrixClient, MatrixEvent } from "../../../src";
+import { RustCrypto } from "../../../src/rust-crypto/rust-crypto";
+import { AddSecretStorageKeyOpts } from "../../../src/secret-storage";
+import { E2EKeyReceiver } from "../../test-utils/E2EKeyReceiver";
+import { E2EKeyResponder } from "../../test-utils/E2EKeyResponder";
+
+describe("Device dehydration", () => {
+    it("should rehydrate and dehydrate a device", async () => {
+        jest.useFakeTimers({ doNotFake: ["queueMicrotask"] });
+
+        const matrixClient = createClient({
+            baseUrl: "http://test.server",
+            userId: "@alice:localhost",
+            deviceId: "aliceDevice",
+            cryptoCallbacks: {
+                getSecretStorageKey: async (keys: any, name: string) => {
+                    return [[...Object.keys(keys.keys)][0], new Uint8Array(32)];
+                },
+            },
+        });
+
+        await initializeSecretStorage(matrixClient, "@alice:localhost", "http://test.server");
+
+        // count the number of times the dehydration key gets set
+        let setDehydrationCount = 0;
+        matrixClient.on(ClientEvent.AccountData, (event: MatrixEvent) => {
+            if (event.getType() === "org.matrix.msc3814") {
+                setDehydrationCount++;
+            }
+        });
+
+        const crypto = matrixClient.getCrypto()!;
+        fetchMock.config.overwriteRoutes = true;
+
+        // start dehydration -- we start with no dehydrated device, and we
+        // store the dehydrated device that we create
+        fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+            status: 404,
+            body: {
+                errcode: "M_NOT_FOUND",
+                error: "Not found",
+            },
+        });
+        let dehydratedDeviceBody: any;
+        let dehydrationCount = 0;
+        let resolveDehydrationPromise: () => void;
+        fetchMock.put("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", (_, opts) => {
+            dehydratedDeviceBody = JSON.parse(opts.body as string);
+            dehydrationCount++;
+            if (resolveDehydrationPromise) {
+                resolveDehydrationPromise();
+            }
+            return {};
+        });
+        await crypto.startDehydration();
+
+        expect(dehydrationCount).toEqual(1);
+
+        // a week later, we should have created another dehydrated device
+        const dehydrationPromise = new Promise<void>((resolve, reject) => {
+            resolveDehydrationPromise = resolve;
+        });
+        jest.advanceTimersByTime(7 * 24 * 60 * 60 * 1000);
+        await dehydrationPromise;
+        expect(dehydrationCount).toEqual(2);
+
+        // restart dehydration -- rehydrate the device that we created above,
+        // and create a new dehydrated device.  We also set `createNewKey`, so
+        // a new dehydration key will be set
+        fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+            device_id: dehydratedDeviceBody.device_id,
+            device_data: dehydratedDeviceBody.device_data,
+        });
+        const eventsResponse = jest.fn((url, opts) => {
+            // rehydrating should make two calls to the /events endpoint.
+            // The first time will return a single event, and the second
+            // time will return no events (which will signal to the
+            // rehydration function that it can stop)
+            const body = JSON.parse(opts.body as string);
+            const nextBatch = body.next_batch ?? "0";
+            const events = nextBatch === "0" ? [{ sender: "@alice:localhost", type: "m.dummy", content: {} }] : [];
+            return {
+                events,
+                next_batch: nextBatch + "1",
+            };
+        });
+        fetchMock.post(
+            `path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device/${encodeURIComponent(dehydratedDeviceBody.device_id)}/events`,
+            eventsResponse,
+        );
+        await crypto.startDehydration(true);
+        expect(dehydrationCount).toEqual(3);
+
+        expect(setDehydrationCount).toEqual(2);
+        expect(eventsResponse.mock.calls).toHaveLength(2);
+
+        matrixClient.stopClient();
+    });
+});
+
+/** create a new secret storage and cross-signing keys */
+async function initializeSecretStorage(
+    matrixClient: MatrixClient,
+    userId: string,
+    homeserverUrl: string,
+): Promise<void> {
+    fetchMock.get("path:/_matrix/client/v3/room_keys/version", {
+        status: 404,
+        body: {
+            errcode: "M_NOT_FOUND",
+            error: "Not found",
+        },
+    });
+    const e2eKeyReceiver = new E2EKeyReceiver(homeserverUrl);
+    const e2eKeyResponder = new E2EKeyResponder(homeserverUrl);
+    e2eKeyResponder.addKeyReceiver(userId, e2eKeyReceiver);
+    fetchMock.post("path:/_matrix/client/v3/keys/device_signing/upload", {});
+    fetchMock.post("path:/_matrix/client/v3/keys/signatures/upload", {});
+    const accountData: Map<string, object> = new Map();
+    fetchMock.get("glob:http://*/_matrix/client/v3/user/*/account_data/*", (url, opts) => {
+        const name = url.split("/").pop()!;
+        const value = accountData.get(name);
+        if (value) {
+            return value;
+        } else {
+            return {
+                status: 404,
+                body: {
+                    errcode: "M_NOT_FOUND",
+                    error: "Not found",
+                },
+            };
+        }
+    });
+    fetchMock.put("glob:http://*/_matrix/client/v3/user/*/account_data/*", (url, opts) => {
+        const name = url.split("/").pop()!;
+        const value = JSON.parse(opts.body as string);
+        accountData.set(name, value);
+        matrixClient.emit(ClientEvent.AccountData, new MatrixEvent({ type: name, content: value }));
+        return {};
+    });
+
+    await matrixClient.initRustCrypto();
+    const crypto = matrixClient.getCrypto()! as RustCrypto;
+    // we need to process a sync so that the OlmMachine will upload keys
+    await crypto.preprocessToDeviceMessages([]);
+    await crypto.onSyncCompleted({});
+
+    // create initial secret storage
+    async function createSecretStorageKey() {
+        return {
+            keyInfo: {} as AddSecretStorageKeyOpts,
+            privateKey: new Uint8Array(32),
+        };
+    }
+    await matrixClient.bootstrapCrossSigning({ setupNewCrossSigning: true });
+    await matrixClient.bootstrapSecretStorage({
+        createSecretStorageKey,
+        setupNewSecretStorage: true,
+        setupNewKeyBackup: false,
+    });
+}

spec/unit/rust-crypto/OutgoingRequestProcessor.spec.ts

@@ -22,6 +22,7 @@ import {
     KeysClaimRequest,
     KeysQueryRequest,
     KeysUploadRequest,
+    PutDehydratedDeviceRequest,
     RoomMessageRequest,
     SignatureUploadRequest,
     UploadSigningKeysRequest,
@@ -233,6 +234,35 @@ describe("OutgoingRequestProcessor", () => {
         httpBackend.verifyNoOutstandingRequests();
     });
 
+    it("should handle PutDehydratedDeviceRequest", async () => {
+        // first, mock up a request as we might expect to receive it from the Rust layer ...
+        const testReq = { foo: "bar" };
+        const outgoingRequest = new PutDehydratedDeviceRequest(JSON.stringify(testReq));
+
+        // ... then poke the request into the OutgoingRequestProcessor under test
+        const reqProm = processor.makeOutgoingRequest(outgoingRequest);
+
+        // Now: check that it makes a matching HTTP request.
+        const testResponse = '{"result":1}';
+        httpBackend
+            .when("PUT", "/_matrix")
+            .check((req) => {
+                expect(req.path).toEqual(
+                    "https://example.com/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device",
+                );
+                expect(JSON.parse(req.rawData)).toEqual(testReq);
+                expect(req.headers["Accept"]).toEqual("application/json");
+                expect(req.headers["Content-Type"]).toEqual("application/json");
+            })
+            .respond(200, testResponse, true);
+
+        // PutDehydratedDeviceRequest does not need to be marked as sent, so no call to OlmMachine.markAsSent is expected.
+
+        await httpBackend.flushAllExpected();
+        await reqProm;
+        httpBackend.verifyNoOutstandingRequests();
+    });
+
     it("does not explode with unknown requests", async () => {
         const outgoingRequest = { id: "5678", type: 987 };
         const markSentCallPromise = awaitCallToMarkAsSent();

spec/unit/rust-crypto/rust-crypto.spec.ts

@@ -762,8 +762,11 @@ describe("RustCrypto", () => {
                             },
                         },
                     };
-                } else if (request instanceof RustSdkCryptoJs.UploadSigningKeysRequest) {
-                    // SigningKeysUploadRequest does not implement OutgoingRequest and does not need to be marked as sent.
+                } else if (
+                    request instanceof RustSdkCryptoJs.UploadSigningKeysRequest ||
+                    request instanceof RustSdkCryptoJs.PutDehydratedDeviceRequest
+                ) {
+                    // These request types do not implement OutgoingRequest and do not need to be marked as sent.
                     return;
                 }
                 if (request.id) {
@@ -1395,6 +1398,34 @@ describe("RustCrypto", () => {
             });
         });
     });
+
+    describe("device dehydration", () => {
+        it("should detect if dehydration is supported", async () => {
+            const rustCrypto = await makeTestRustCrypto(makeMatrixHttpApi());
+            fetchMock.config.overwriteRoutes = true;
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                status: 404,
+                body: {
+                    errcode: "M_UNRECOGNIZED",
+                    error: "Unknown endpoint",
+                },
+            });
+            expect(await rustCrypto.isDehydrationSupported()).toBe(false);
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                status: 404,
+                body: {
+                    errcode: "M_NOT_FOUND",
+                    error: "Not found",
+                },
+            });
+            expect(await rustCrypto.isDehydrationSupported()).toBe(true);
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                device_id: "DEVICE_ID",
+                device_data: "data",
+            });
+            expect(await rustCrypto.isDehydrationSupported()).toBe(true);
+        });
+    });
 });
 
 /** Build a MatrixHttpApi instance */

src/crypto-api.ts

@@ -496,6 +496,42 @@ export interface CryptoApi {
      * @param version - The backup version to delete.
      */
     deleteKeyBackupVersion(version: string): Promise<void>;
+
+    ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+    //
+    // Dehydrated devices
+    //
+    ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+    /**
+     * Returns whether MSC3814 dehydrated devices are supported by the crypto
+     * backend and by the server.
+     *
+     * This should be called before calling `startDehydration`, and if this
+     * returns `false`, `startDehydration` should not be called.
+     */
+    isDehydrationSupported(): Promise<boolean>;
+
+    /**
+     * Start using device dehydration.
+     *
+     * - Rehydrates a dehydrated device, if one is available.
+     * - Creates a new dehydration key, if necessary, and stores it in Secret
+     *   Storage.
+     *   - If `createNewKey` is set to true, always creates a new key.
+     *   - If a dehydration key is not available, creates a new one.
+     * - Creates a new dehydrated device, and schedules periodically creating
+     *   new dehydrated devices.
+     *
+     * This function must not be called unless `isDehydrationSupported` returns
+     * `true`, and must not be called until after cross-signing and secret
+     * storage have been set up.
+     *
+     * @param createNewKey - whether to force creation of a new dehydration key.
+     *   This can be used, for example, if Secret Storage is being reset.  Defaults
+     *   to false.
+     */
+    startDehydration(createNewKey?: boolean): Promise<void>;
 }
 
 /** A reason code for a failure to decrypt an event. */

src/crypto/index.ts

@@ -4287,6 +4287,21 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
     public getRoomEncryption(roomId: string): IRoomEncryption | null {
         return this.roomList.getRoomEncryption(roomId);
     }
+
+    /**
+     * Returns whether dehydrated devices are supported by the crypto backend
+     * and by the server.
+     */
+    public async isDehydrationSupported(): Promise<boolean> {
+        return false;
+    }
+
+    /**
+     * Stub function -- dehydration is not implemented here, so throw error
+     */
+    public async startDehydration(createNewKey?: boolean): Promise<void> {
+        throw new Error("Not implemented");
+    }
 }
 
 /**

src/models/device.ts

@@ -51,6 +51,9 @@ export class Device {
     /** display name of the device */
     public readonly displayName?: string;
 
+    /** whether the device is a dehydrated device */
+    public readonly dehydrated: boolean = false;
+
     public constructor(opts: DeviceParameters) {
         this.deviceId = opts.deviceId;
         this.userId = opts.userId;
@@ -59,6 +62,7 @@ export class Device {
         this.verified = opts.verified || DeviceVerification.Unverified;
         this.signatures = opts.signatures || new Map();
         this.displayName = opts.displayName;
+        this.dehydrated = !!opts.dehydrated;
     }
 
     /**