Include "needsBackup" flag in inbound group session batches

matrix-org · Jan 8, 2024 · 9c62595 · 9c62595
1 parent a9c6f47
commit 9c62595
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 29 deletions.
diff --git a/spec/unit/crypto/store/CryptoStore.spec.ts b/spec/unit/crypto/store/CryptoStore.spec.ts
@@ -145,6 +145,11 @@ describe.each([
             const N_SESSIONS_PER_DEVICE = 6;
             await createSessions(N_DEVICES, N_SESSIONS_PER_DEVICE);
 
+            // Mark one of the sessions as needing backup
+            await store.doTxn("readwrite", IndexedDBCryptoStore.STORE_BACKUP, async (txn) => {
+                await store.markSessionsNeedingBackup([{ senderKey: pad43("device5"), sessionId: "session5" }], txn);
+            });
+
             const batch = await store.getEndToEndInboundGroupSessionsBatch();
             expect(batch!.length).toEqual(N_DEVICES * N_SESSIONS_PER_DEVICE);
             for (let i = 0; i < N_DEVICES; i++) {
@@ -153,6 +158,9 @@ describe.each([
 
                     expect(r.senderKey).toEqual(pad43(`device${i}`));
                     expect(r.sessionId).toEqual(`session${j}`);
+
+                    // only the last session needs backup
+                    expect(r.needsBackup).toBe(i === 5 && j === 5);
                 }
             }
         });

diff --git a/src/crypto/store/base.ts b/src/crypto/store/base.ts
@@ -182,7 +182,7 @@ export interface CryptoStore {
      * @returns A batch of Megolm Sessions, or `null` if no sessions are left.
      * @internal
      */
-    getEndToEndInboundGroupSessionsBatch(): Promise<ISession[] | null>;
+    getEndToEndInboundGroupSessionsBatch(): Promise<ISessionExtended[] | null>;
 
     /**
      * Delete a batch of Megolm sessions from the database.
@@ -223,6 +223,11 @@ export interface ISession {
     sessionData?: InboundGroupSessionData;
 }
 
+/** Extended data on a Megolm session */
+export interface ISessionExtended extends ISession {
+    needsBackup: boolean;
+}
+
 /** Data on an Olm session */
 export interface ISessionInfo {
     deviceKey?: string;

diff --git a/src/crypto/store/indexeddb-crypto-store-backend.ts b/src/crypto/store/indexeddb-crypto-store-backend.ts
@@ -21,6 +21,7 @@ import {
     IDeviceData,
     IProblem,
     ISession,
+    ISessionExtended,
     ISessionInfo,
     IWithheld,
     MigrationState,
@@ -815,29 +816,39 @@ export class Backend implements CryptoStore {
      *
      * Implementation of {@link CryptoStore.getEndToEndInboundGroupSessionsBatch}.
      */
-    public async getEndToEndInboundGroupSessionsBatch(): Promise<null | ISession[]> {
-        const result: ISession[] = [];
-        await this.doTxn("readonly", [IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS], (txn) => {
-            const objectStore = txn.objectStore(IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS);
-            const getReq = objectStore.openCursor();
-            getReq.onsuccess = function (): void {
-                try {
-                    const cursor = getReq.result;
-                    if (cursor) {
-                        result.push({
-                            senderKey: cursor.value.senderCurve25519Key,
-                            sessionId: cursor.value.sessionId,
-                            sessionData: cursor.value.session,
-                        });
-                        if (result.length < SESSION_BATCH_SIZE) {
-                            cursor.continue();
+    public async getEndToEndInboundGroupSessionsBatch(): Promise<null | ISessionExtended[]> {
+        const result: ISessionExtended[] = [];
+        await this.doTxn(
+            "readonly",
+            [IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS, IndexedDBCryptoStore.STORE_BACKUP],
+            (txn) => {
+                const sessionStore = txn.objectStore(IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS);
+                const backupStore = txn.objectStore(IndexedDBCryptoStore.STORE_BACKUP);
+
+                const getReq = sessionStore.openCursor();
+                getReq.onsuccess = function (): void {
+                    try {
+                        const cursor = getReq.result;
+                        if (cursor) {
+                            const backupGetReq = backupStore.get(cursor.key);
+                            backupGetReq.onsuccess = (): void => {
+                                result.push({
+                                    senderKey: cursor.value.senderCurve25519Key,
+                                    sessionId: cursor.value.sessionId,
+                                    sessionData: cursor.value.session,
+                                    needsBackup: backupGetReq.result !== undefined,
+                                });
+                                if (result.length < SESSION_BATCH_SIZE) {
+                                    cursor.continue();
+                                }
+                            };
                         }
+                    } catch (e) {
+                        abortWithException(txn, <Error>e);
                     }
-                } catch (e) {
-                    abortWithException(txn, <Error>e);
-                }
-            };
-        });
+                };
+            },
+        );
 
         if (result.length === 0) {
             // No sessions left.

diff --git a/src/crypto/store/indexeddb-crypto-store.ts b/src/crypto/store/indexeddb-crypto-store.ts
@@ -25,6 +25,7 @@ import {
     IDeviceData,
     IProblem,
     ISession,
+    ISessionExtended,
     ISessionInfo,
     IWithheld,
     MigrationState,
@@ -610,7 +611,7 @@ export class IndexedDBCryptoStore implements CryptoStore {
      *
      * @internal
      */
-    public getEndToEndInboundGroupSessionsBatch(): Promise<ISession[] | null> {
+    public getEndToEndInboundGroupSessionsBatch(): Promise<ISessionExtended[] | null> {
         return this.backend!.getEndToEndInboundGroupSessionsBatch();
     }
 

diff --git a/src/crypto/store/localStorage-crypto-store.ts b/src/crypto/store/localStorage-crypto-store.ts
@@ -21,6 +21,7 @@ import {
     IDeviceData,
     IProblem,
     ISession,
+    ISessionExtended,
     ISessionInfo,
     IWithheld,
     MigrationState,
@@ -357,20 +358,24 @@ export class LocalStorageCryptoStore extends MemoryCryptoStore implements Crypto
      *
      * @internal
      */
-    public async getEndToEndInboundGroupSessionsBatch(): Promise<ISession[] | null> {
-        const result: ISession[] = [];
+    public async getEndToEndInboundGroupSessionsBatch(): Promise<ISessionExtended[] | null> {
+        const sessionsNeedingBackup = getJsonItem<string[]>(this.store, KEY_SESSIONS_NEEDING_BACKUP) || {};
+        const result: ISessionExtended[] = [];
         for (let i = 0; i < this.store.length; ++i) {
             const key = this.store.key(i);
             if (key?.startsWith(KEY_INBOUND_SESSION_PREFIX)) {
+                const key2 = key.slice(KEY_INBOUND_SESSION_PREFIX.length);
+
                 // we can't use split, as the components we are trying to split out
                 // might themselves contain '/' characters. We rely on the
                 // senderKey being a (32-byte) curve25519 key, base64-encoded
                 // (hence 43 characters long).
 
                 result.push({
-                    senderKey: key.slice(KEY_INBOUND_SESSION_PREFIX.length, KEY_INBOUND_SESSION_PREFIX.length + 43),
-                    sessionId: key.slice(KEY_INBOUND_SESSION_PREFIX.length + 44),
+                    senderKey: key2.slice(0, 43),
+                    sessionId: key2.slice(44),
                     sessionData: getJsonItem(this.store, key)!,
+                    needsBackup: key2 in sessionsNeedingBackup,
                 });
 
                 if (result.length >= SESSION_BATCH_SIZE) {

diff --git a/src/crypto/store/memory-crypto-store.ts b/src/crypto/store/memory-crypto-store.ts
@@ -21,6 +21,7 @@ import {
     IDeviceData,
     IProblem,
     ISession,
+    ISessionExtended,
     ISessionInfo,
     IWithheld,
     MigrationState,
@@ -534,13 +535,14 @@ export class MemoryCryptoStore implements CryptoStore {
      *
      * @internal
      */
-    public async getEndToEndInboundGroupSessionsBatch(): Promise<null | ISession[]> {
-        const result: ISession[] = [];
+    public async getEndToEndInboundGroupSessionsBatch(): Promise<null | ISessionExtended[]> {
+        const result: ISessionExtended[] = [];
         for (const [key, session] of Object.entries(this.inboundGroupSessions)) {
             result.push({
                 senderKey: key.slice(0, 43),
                 sessionId: key.slice(44),
                 sessionData: session,
+                needsBackup: key in this.sessionsNeedingBackup,
             });
             if (result.length >= SESSION_BATCH_SIZE) {
                 return result;