Send one time key count and unused fallback keys for rust-crypto

spec/integ/crypto.spec.ts

@@ -1890,4 +1890,70 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
             expect(event.getContent().body).not.toContain("withheld");
         });
     });
+
+    describe("key upload request", () => {
+        beforeEach(() => {
+            jest.useFakeTimers();
+        });
+
+        afterEach(() => {
+            jest.useRealTimers();
+        });
+
+        function listenToUpload(): Promise<number> {
+            return new Promise((resolve) => {
+                const listener = (url: string, options: RequestInit) => {
+                    const content = JSON.parse(options.body as string);
+                    const keysCount = Object.keys(content?.one_time_keys || {}).length;
+                    if (keysCount) resolve(keysCount);
+                    return {
+                        one_time_key_counts: {
+                            signed_curve25519: keysCount ? 60 : keysCount,
+                        },
+                    };
+                };
+
+                // catch both r0 and v3 variants
+                fetchMock.post(
+                    new URL("/_matrix/client/r0/keys/upload", aliceClient.getHomeserverUrl()).toString(),
+                    listener,
+                    {
+                        overwriteRoutes: true,
+                    },
+                );
+                fetchMock.post(
+                    new URL("/_matrix/client/v3/keys/upload", aliceClient.getHomeserverUrl()).toString(),
+                    listener,
+                    {
+                        overwriteRoutes: true,
+                    },
+                );
+            });
+        }
+
+        it("should make key upload request after sync", async () => {
+            let uploadPromise = listenToUpload();
+            expectAliceKeyQuery({ device_keys: { "@alice:localhost": {} }, failures: {} });
+            await startClientAndAwaitFirstSync();
+
+            syncResponder.sendOrQueueSyncResponse(getSyncResponse([]));
+
+            await syncPromise(aliceClient);
+            expect(await uploadPromise).toBeGreaterThan(0);
+
+            uploadPromise = listenToUpload();
+            syncResponder.sendOrQueueSyncResponse({
+                next_batch: 2,
+                device_one_time_keys_count: { signed_curve25519: 0 },
+            });
+
+            // Advance local date to 2 minutes
+            // The old crypto only runs the upload every 60 seconds
+            jest.setSystemTime(Date.now() + 2 * 60 * 1000);
+
+            await syncPromise(aliceClient);
+
+            expect(await uploadPromise).toBeGreaterThan(0);
+        });
+    });
 });

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

@@ -53,7 +53,7 @@ describe("RustCrypto", () => {
         });
     });
 
-    describe("to-device messages", () => {
+    describe("call preprocess methods", () => {
         let rustCrypto: RustCrypto;
 
         beforeEach(async () => {
@@ -92,6 +92,16 @@ describe("RustCrypto", () => {
             const res = await rustCrypto.preprocessToDeviceMessages(inputs);
             expect(res).toEqual(inputs);
         });
+
+        it("should pass through one time key counts", async () => {
+            const oneTimeKeyCounts = new Map<string, number>([["signed_curve25519", 50]]);
+            await expect(rustCrypto.preprocessOneTimeKeyCounts(oneTimeKeyCounts)).resolves.not.toBeDefined();
+        });
+
+        it("should pass through unused fallback keys", async () => {
+            const unusedFallbackKeys = new Set(["signed_curve25519"]);
+            await expect(rustCrypto.preprocessUnusedFallbackKeys(unusedFallbackKeys)).resolves.not.toBeDefined();
+        });
     });
 
     describe("outgoing requests", () => {

src/common-crypto/CryptoBackend.ts

@@ -105,6 +105,34 @@ export interface SyncCryptoCallbacks {
      */
     preprocessToDeviceMessages(events: IToDeviceEvent[]): Promise<IToDeviceEvent[]>;
 
+    /**
+     * Called by the /sync loop whenever there are incoming to-device messages.
+     *
+     * The implementation may preprocess the received messages (eg, decrypt them) and return an
+     * updated list of messages for dispatch to the rest of the system.
+     *
+     * Note that, unlike {@link ClientEvent.ToDeviceEvent} events, this is called on the raw to-device
+     * messages, rather than the results of any decryption attempts.
+     *
+     * @param oneTimeKeysCounts - the received one time key counts
+     * @returns A list of preprocessed to-device messages.
+     */
+    preprocessOneTimeKeyCounts(oneTimeKeysCounts: Map<string, number>): Promise<void>;
+
+    /**
+     * Called by the /sync loop whenever there are incoming to-device messages.
+     *
+     * The implementation may preprocess the received messages (eg, decrypt them) and return an
+     * updated list of messages for dispatch to the rest of the system.
+     *
+     * Note that, unlike {@link ClientEvent.ToDeviceEvent} events, this is called on the raw to-device
+     * messages, rather than the results of any decryption attempts.
+     *
+     * @param unusedFallbackKeys - the received unused fallback keys
+     * @returns A list of preprocessed to-device messages.
+     */
+    preprocessUnusedFallbackKeys(unusedFallbackKeys: Set<string>): Promise<void>;
+
     /**
      * Called by the /sync loop whenever an m.room.encryption event is received.
      *

src/crypto/index.ts

@@ -3207,6 +3207,17 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
         });
     }
 
+    public preprocessOneTimeKeyCounts(oneTimeKeysCounts: Map<string, number>): Promise<void> {
+        const currentCount = oneTimeKeysCounts.get("signed_curve25519") || 0;
+        this.updateOneTimeKeyCount(currentCount);
+        return Promise.resolve();
+    }
+
+    public preprocessUnusedFallbackKeys(unusedFallbackKeys: Set<string>): Promise<void> {
+        this.setNeedsNewFallback(!unusedFallbackKeys.has("signed_curve25519"));
+        return Promise.resolve();
+    }
+
     private onToDeviceEvent = (event: MatrixEvent): void => {
         try {
             logger.log(

src/rust-crypto/rust-crypto.ts

@@ -180,25 +180,62 @@ export class RustCrypto implements CryptoBackend {
     //
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-    /** called by the sync loop to preprocess incoming to-device messages
-     *
+    /**
+     * Apply sync changes to the olm machine
      * @param events - the received to-device messages
+     * @param oneTimeKeysCounts - the received one time key counts
+     * @param unusedFallbackKeys - the received unused fallback keys
      * @returns A list of preprocessed to-device messages.
      */
-    public async preprocessToDeviceMessages(events: IToDeviceEvent[]): Promise<IToDeviceEvent[]> {
-        // send the received to-device messages into receiveSyncChanges. We have no info on device-list changes,
-        // one-time-keys, or fallback keys, so just pass empty data.
+    private async receiveSyncChanges({
+        events,
+        oneTimeKeysCounts = new Map<string, number>(),
+        unusedFallbackKeys = new Set<string>(),
+    }: {
+        events?: IToDeviceEvent[];
+        oneTimeKeysCounts?: Map<string, number>;
+        unusedFallbackKeys?: Set<string>;
+    }): Promise<IToDeviceEvent[]> {
         const result = await this.olmMachine.receiveSyncChanges(
-            JSON.stringify(events),
+            events ? JSON.stringify(events) : "[]",
             new RustSdkCryptoJs.DeviceLists(),
-            new Map(),
-            new Set(),
+            oneTimeKeysCounts,
+            unusedFallbackKeys,
         );
 
         // receiveSyncChanges returns a JSON-encoded list of decrypted to-device messages.
         return JSON.parse(result);
     }
 
+    /** called by the sync loop to preprocess incoming to-device messages
+     *
+     * @param events - the received to-device messages
+     * @returns A list of preprocessed to-device messages.
+     */
+    public preprocessToDeviceMessages(events: IToDeviceEvent[]): Promise<IToDeviceEvent[]> {
+        // send the received to-device messages into receiveSyncChanges. We have no info on device-list changes,
+        // one-time-keys, or fallback keys, so just pass empty data.
+        return this.receiveSyncChanges({ events });
+    }
+
+    /** called by the sync loop to preprocess one time key counts
+     *
+     * @param oneTimeKeysCounts - the received one time key counts
+     * @returns A list of preprocessed to-device messages.
+     */
+    public async preprocessOneTimeKeyCounts(oneTimeKeysCounts: Map<string, number>): Promise<void> {
+        await this.receiveSyncChanges({ oneTimeKeysCounts });
+    }
+
+    /** called by the sync loop to preprocess unused fallback keys
+     *
+     * @param unusedFallbackKeys - the received unused fallback keys
+     * @returns A list of preprocessed to-device messages.
+     */
+    public async preprocessUnusedFallbackKeys(unusedFallbackKeys: Set<string>): Promise<void> {
+        await this.receiveSyncChanges({ unusedFallbackKeys });
+    }
+
     /** called by the sync loop on m.room.encrypted events
      *
      * @param room - in which the event was received

src/sync.ts

@@ -1523,22 +1523,17 @@ export class SyncApi {
         }
 
         // Handle one_time_keys_count
-        if (this.syncOpts.crypto && data.device_one_time_keys_count) {
-            const currentCount = data.device_one_time_keys_count.signed_curve25519 || 0;
-            this.syncOpts.crypto.updateOneTimeKeyCount(currentCount);
+        if (data.device_one_time_keys_count) {
+            const map = new Map<string, number>(Object.entries(data.device_one_time_keys_count));
+            this.syncOpts.cryptoCallbacks?.preprocessOneTimeKeyCounts(map);
         }
-        if (
-            this.syncOpts.crypto &&
-            (data.device_unused_fallback_key_types || data["org.matrix.msc2732.device_unused_fallback_key_types"])
-        ) {
+        if (data.device_unused_fallback_key_types || data["org.matrix.msc2732.device_unused_fallback_key_types"]) {
             // The presence of device_unused_fallback_key_types indicates that the
             // server supports fallback keys. If there's no unused
             // signed_curve25519 fallback key we need a new one.
             const unusedFallbackKeys =
                 data.device_unused_fallback_key_types || data["org.matrix.msc2732.device_unused_fallback_key_types"];
-            this.syncOpts.crypto.setNeedsNewFallback(
-                Array.isArray(unusedFallbackKeys) && !unusedFallbackKeys.includes("signed_curve25519"),
-            );
+            this.syncOpts.cryptoCallbacks?.preprocessUnusedFallbackKeys(new Set<string>(unusedFallbackKeys || null));
         }
     }