From 290aaad63a9dbb7d3acd845f08edc9c2b5d55143 Mon Sep 17 00:00:00 2001
From: Bruno Windels <bruno@windels.cloud>
Date: Tue, 25 Jan 2022 18:46:51 +0100
Subject: [PATCH] add sessionsNeedingBackup store

---
 src/matrix/storage/common.ts                  |  1 +
 src/matrix/storage/idb/Transaction.ts         |  5 ++
 src/matrix/storage/idb/schema.ts              |  8 ++-
 .../idb/stores/SessionNeedingBackupStore.ts   | 56 +++++++++++++++++++
 4 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 src/matrix/storage/idb/stores/SessionNeedingBackupStore.ts

diff --git a/src/matrix/storage/common.ts b/src/matrix/storage/common.ts
index 23bb0d31..3929c061 100644
--- a/src/matrix/storage/common.ts
+++ b/src/matrix/storage/common.ts
@@ -33,6 +33,7 @@ export enum StoreNames {
     groupSessionDecryptions = "groupSessionDecryptions",
     operations = "operations",
     accountData = "accountData",
+    sessionsNeedingBackup = "sessionsNeedingBackup",
 }
 
 export const STORE_NAMES: Readonly<StoreNames[]> = Object.values(StoreNames);
diff --git a/src/matrix/storage/idb/Transaction.ts b/src/matrix/storage/idb/Transaction.ts
index 80894105..9b765310 100644
--- a/src/matrix/storage/idb/Transaction.ts
+++ b/src/matrix/storage/idb/Transaction.ts
@@ -36,6 +36,7 @@ import {OutboundGroupSessionStore} from "./stores/OutboundGroupSessionStore";
 import {GroupSessionDecryptionStore} from "./stores/GroupSessionDecryptionStore";
 import {OperationStore} from "./stores/OperationStore";
 import {AccountDataStore} from "./stores/AccountDataStore";
+import {SessionNeedingBackupStore} from "./stores/SessionNeedingBackupStore";
 import type {ILogger, ILogItem} from "../../../logging/types";
 
 export type IDBKey = IDBValidKey | IDBKeyRange;
@@ -151,6 +152,10 @@ export class Transaction {
     get inboundGroupSessions(): InboundGroupSessionStore {
         return this._store(StoreNames.inboundGroupSessions, idbStore => new InboundGroupSessionStore(idbStore));
     }
+
+    get sessionsNeedingBackup(): SessionNeedingBackupStore {
+        return this._store(StoreNames.sessionsNeedingBackup, idbStore => new SessionNeedingBackupStore(idbStore));
+    }
     
     get outboundGroupSessions(): OutboundGroupSessionStore {
         return this._store(StoreNames.outboundGroupSessions, idbStore => new OutboundGroupSessionStore(idbStore));
diff --git a/src/matrix/storage/idb/schema.ts b/src/matrix/storage/idb/schema.ts
index ad3e5896..44a04bba 100644
--- a/src/matrix/storage/idb/schema.ts
+++ b/src/matrix/storage/idb/schema.ts
@@ -31,7 +31,8 @@ export const schema: MigrationFunc[] = [
     fixMissingRoomsInUserIdentities,
     changeSSSSKeyPrefix,
     backupAndRestoreE2EEAccountToLocalStorage,
-    clearAllStores
+    clearAllStores,
+    createSessionsNeedingBackup
 ];
 // TODO: how to deal with git merge conflicts of this array?
 
@@ -270,3 +271,8 @@ async function clearAllStores(db: IDBDatabase, txn: IDBTransaction) {
         }
     }
 }
+
+// v15 adds the sessionsNeedingBackup store, for session backup
+function createSessionsNeedingBackup(db: IDBDatabase): void {
+    db.createObjectStore("sessionsNeedingBackup", {keyPath: "key"});
+}
diff --git a/src/matrix/storage/idb/stores/SessionNeedingBackupStore.ts b/src/matrix/storage/idb/stores/SessionNeedingBackupStore.ts
new file mode 100644
index 00000000..27104191
--- /dev/null
+++ b/src/matrix/storage/idb/stores/SessionNeedingBackupStore.ts
@@ -0,0 +1,56 @@
+/*
+Copyright 2020 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 type {Store} from "../Store";
+
+export type BackupEntry = {
+    roomId: string;
+    senderKey: string;
+    sessionId: string;
+};
+
+type StorageEntry = {
+    key: string
+};
+
+function encodeKey(roomId: string, senderKey: string, sessionId: string): string {
+    return `${roomId}|${senderKey}|${sessionId}`;
+}
+
+function decodeKey(key: string): BackupEntry {
+    const [roomId, senderKey, sessionId] = key.split("|");
+    return {roomId, senderKey, sessionId};
+}
+
+export class SessionNeedingBackupStore {
+    constructor(private store: Store<StorageEntry>) {}
+
+    async getFirstEntries(amount: number): Promise<BackupEntry[]> {
+        const storageEntries = await this.store.selectLimit(undefined, amount);
+        return storageEntries.map(s => decodeKey(s.key));
+    }
+
+    set(roomId: string, senderKey: string, sessionId: string): void {
+        const storageEntry : StorageEntry = {
+            key: encodeKey(roomId, senderKey, sessionId),
+        };
+        this.store.put(storageEntry);
+    }
+    
+    remove(roomId: string, senderKey: string, sessionId: string): void {
+        this.store.delete(encodeKey(roomId, senderKey, sessionId));
+    }
+}