From ae5cc172907a6b6e39f7e4b9267521e2bec6c2e9 Mon Sep 17 00:00:00 2001
From: Bruno Windels <bruno@windels.cloud>
Date: Mon, 31 Jan 2022 14:36:04 +0100
Subject: [PATCH] mark all inbound sessions to be backed up again when changing
 version

---
 src/matrix/Session.js                         | 21 ++++++++++++++-----
 src/matrix/e2ee/megolm/keybackup/KeyBackup.ts |  4 ++++
 src/matrix/ssss/index.ts                      |  6 +++++-
 .../idb/stores/InboundGroupSessionStore.ts    | 12 +++++++++++
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/src/matrix/Session.js b/src/matrix/Session.js
index 3d1530c9..2a1618a2 100644
--- a/src/matrix/Session.js
+++ b/src/matrix/Session.js
@@ -203,21 +203,32 @@ export class Session {
                 this._storage.storeNames.accountData,
             ]);
             if (await this._createKeyBackup(key, readTxn, log)) {
-                await this._writeSSSSKey(key);
+                // only after having read a secret, write the key
+                // as we only find out if it was good if the MAC verification succeeds
+                await this._writeSSSSKey(key, log);
                 this._keyBackup.get().flush(log);
                 return key;
             }
         });
     }
 
-    async _writeSSSSKey(key) {
-        // only after having read a secret, write the key
-        // as we only find out if it was good if the MAC verification succeeds
+    async _writeSSSSKey(key, log) {
+        // we're going to write the 4S key, and also the backup version.
+        // this way, we can detect when we enter a key for a new backup version
+        // and mark all inbound sessions to be backed up again
+        const backupVersion = this._keyBackup.get()?.version;
         const writeTxn = await this._storage.readWriteTxn([
             this._storage.storeNames.session,
+            this._storage.storeNames.inboundGroupSessions,
         ]);
         try {
-            ssssWriteKey(key, writeTxn);
+            const previousBackupVersion = await ssssWriteKey(key, backupVersion, writeTxn);
+            log.set("previousBackupVersion", previousBackupVersion);
+            log.set("backupVersion", backupVersion);
+            if (typeof previousBackupVersion === "number" && previousBackupVersion !== backupVersion) {
+                const amountMarked = await this._keyBackup.markAllForBackup(writeTxn);
+                log.set("amountMarkedForBackup", amountMarked);
+            }
         } catch (err) {
             writeTxn.abort();
             throw err;
diff --git a/src/matrix/e2ee/megolm/keybackup/KeyBackup.ts b/src/matrix/e2ee/megolm/keybackup/KeyBackup.ts
index e9a8756d..8f9611dd 100644
--- a/src/matrix/e2ee/megolm/keybackup/KeyBackup.ts
+++ b/src/matrix/e2ee/megolm/keybackup/KeyBackup.ts
@@ -62,6 +62,10 @@ export class KeyBackup {
         }
     }
 
+    markAllForBackup(txn: Transaction): Promise<number> {
+        return txn.inboundGroupSessions.markAllAsNotBackedUp();
+    }
+
     flush(log: ILogItem): void {
         if (!this.operationInProgress.get()) {
             log.wrapDetached("flush key backup", async log => {
diff --git a/src/matrix/ssss/index.ts b/src/matrix/ssss/index.ts
index 37f47963..fd4c2245 100644
--- a/src/matrix/ssss/index.ts
+++ b/src/matrix/ssss/index.ts
@@ -27,6 +27,7 @@ import type * as OlmNamespace from "@matrix-org/olm"
 type Olm = typeof OlmNamespace;
 
 const SSSS_KEY = `${SESSION_E2EE_KEY_PREFIX}ssssKey`;
+const BACKUPVERSION_KEY = `${SESSION_E2EE_KEY_PREFIX}keyBackupVersion`;
 
 export enum KeyType {
     "RecoveryKey",
@@ -49,8 +50,11 @@ async function readDefaultKeyDescription(storage: Storage): Promise<KeyDescripti
     return new KeyDescription(id, keyAccountData.content as KeyDescriptionData);
 }
 
-export async function writeKey(key: Key, txn: Transaction): Promise<void> {
+export async function writeKey(key: Key, keyBackupVersion: number, txn: Transaction): Promise<number | undefined> {
+    const existingVersion: number | undefined = await txn.session.get(BACKUPVERSION_KEY);
+    txn.session.set(BACKUPVERSION_KEY, keyBackupVersion);
     txn.session.set(SSSS_KEY, {id: key.id, binaryKey: key.binaryKey});
+    return existingVersion;
 }
 
 export async function readKey(txn: Transaction): Promise<Key | undefined> {
diff --git a/src/matrix/storage/idb/stores/InboundGroupSessionStore.ts b/src/matrix/storage/idb/stores/InboundGroupSessionStore.ts
index 4d3497a4..b78c817e 100644
--- a/src/matrix/storage/idb/stores/InboundGroupSessionStore.ts
+++ b/src/matrix/storage/idb/stores/InboundGroupSessionStore.ts
@@ -91,4 +91,16 @@ export class InboundGroupSessionStore {
             this._store.put(entry);
         }
     }
+
+    async markAllAsNotBackedUp(): Promise<number> {
+        const backedUpKey = this._store.IDBKeyRange.only(BackupStatus.BackedUp);
+        let count = 0;
+        await this._store.index("byBackup").iterateValues(backedUpKey, (val: InboundGroupSessionEntry, key: IDBValidKey, cur: IDBCursorWithValue) => {
+            val.backup = BackupStatus.NotBackedUp;
+            cur.update(val);
+            count += 1;
+            return false;
+        });
+        return count;
+    }
 }