store e2ee session values as well in localStorage

element-hq · Sep 29, 2021 · 77bd0d3 · 77bd0d3
1 parent cd071e4
commit 77bd0d3
Show file tree

Hide file tree

Showing 11 changed files with 160 additions and 23 deletions.
diff --git a/src/matrix/SessionContainer.js b/src/matrix/SessionContainer.js
@@ -233,6 +233,7 @@ export class SessionContainer {
             platform: this._platform,
         });
         await this._session.load(log);
+        // TODO: check instead storage doesn't have an identity
         if (isNewLogin) {
             this._status.set(LoadStatus.SessionSetup);
             await log.wrap("createIdentity", log => this._session.createIdentity(log));

diff --git a/src/matrix/e2ee/Account.js b/src/matrix/e2ee/Account.js
@@ -15,12 +15,12 @@ limitations under the License.
 */
 
 import anotherjson from "../../../lib/another-json/index.js";
-import {SESSION_KEY_PREFIX, OLM_ALGORITHM, MEGOLM_ALGORITHM} from "./common.js";
+import {SESSION_E2EE_KEY_PREFIX, OLM_ALGORITHM, MEGOLM_ALGORITHM} from "./common.js";
 
 // use common prefix so it's easy to clear properties that are not e2ee related during session clear
-const ACCOUNT_SESSION_KEY = SESSION_KEY_PREFIX + "olmAccount";
-const DEVICE_KEY_FLAG_SESSION_KEY = SESSION_KEY_PREFIX + "areDeviceKeysUploaded";
-const SERVER_OTK_COUNT_SESSION_KEY = SESSION_KEY_PREFIX + "serverOTKCount";
+const ACCOUNT_SESSION_KEY = SESSION_E2EE_KEY_PREFIX + "olmAccount";
+const DEVICE_KEY_FLAG_SESSION_KEY = SESSION_E2EE_KEY_PREFIX + "areDeviceKeysUploaded";
+const SERVER_OTK_COUNT_SESSION_KEY = SESSION_E2EE_KEY_PREFIX + "serverOTKCount";
 
 export class Account {
     static async load({olm, pickleKey, hsApi, userId, deviceId, olmWorker, txn}) {

diff --git a/src/matrix/e2ee/common.js b/src/matrix/e2ee/common.js
@@ -20,7 +20,7 @@ import {createEnum} from "../../utils/enum.js";
 export const DecryptionSource = createEnum("Sync", "Timeline", "Retry");
 
 // use common prefix so it's easy to clear properties that are not e2ee related during session clear
-export const SESSION_KEY_PREFIX = "e2ee:";
+export const SESSION_E2EE_KEY_PREFIX = "e2ee:";
 export const OLM_ALGORITHM = "m.olm.v1.curve25519-aes-sha2";
 export const MEGOLM_ALGORITHM = "m.megolm.v1.aes-sha2";
 

diff --git a/src/matrix/storage/idb/QueryTarget.ts b/src/matrix/storage/idb/QueryTarget.ts
@@ -22,6 +22,7 @@ import {IDBKey} from "./Transaction";
 export interface ITransaction {
     idbFactory: IDBFactory;
     IDBKeyRange: typeof IDBKeyRange;
+    databaseName: string;
     addWriteError(error: StorageError, refItem: LogItem | undefined, operationName: string, keys: IDBKey[] | undefined);
 }
 
@@ -55,6 +56,10 @@ export class QueryTarget<T> {
         return this._transaction.IDBKeyRange;
     }
 
+    get databaseName(): string {
+        return this._transaction.databaseName;
+    }
+
     _openCursor(range?: IDBQuery, direction?: IDBCursorDirection): IDBRequest<IDBCursorWithValue | null> {
         if (range && direction) {
             return this._target.openCursor(range, direction);
@@ -269,6 +274,7 @@ import {QueryTargetWrapper, Store} from "./Store";
 export function tests() {
 
     class MockTransaction extends MockIDBImpl {
+        get databaseName(): string { return "mockdb"; }
         addWriteError(error: StorageError, refItem: LogItem | undefined, operationName: string, keys: IDBKey[] | undefined) {}
     }
 

diff --git a/src/matrix/storage/idb/Storage.ts b/src/matrix/storage/idb/Storage.ts
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+import {IDOMStorage} from "./types";
 import {Transaction} from "./Transaction";
 import { STORE_NAMES, StoreNames, StorageError } from "../common";
 import { reqAsPromise } from "./utils";
@@ -29,13 +30,15 @@ export class Storage {
     readonly idbFactory: IDBFactory
     readonly IDBKeyRange: typeof IDBKeyRange;
     readonly storeNames: typeof StoreNames;
+    readonly localStorage: IDOMStorage;
 
-    constructor(idbDatabase: IDBDatabase, idbFactory: IDBFactory, _IDBKeyRange: typeof IDBKeyRange, hasWebkitEarlyCloseTxnBug: boolean, logger: BaseLogger) {
+    constructor(idbDatabase: IDBDatabase, idbFactory: IDBFactory, _IDBKeyRange: typeof IDBKeyRange, hasWebkitEarlyCloseTxnBug: boolean, localStorage: IDOMStorage, logger: BaseLogger) {
         this._db = idbDatabase;
         this.idbFactory = idbFactory;
         this.IDBKeyRange = _IDBKeyRange;
         this._hasWebkitEarlyCloseTxnBug = hasWebkitEarlyCloseTxnBug;
         this.storeNames = StoreNames;
+        this.localStorage = localStorage;
         this.logger = logger;
     }
 
@@ -79,4 +82,8 @@ export class Storage {
     close(): void {
         this._db.close();
     }
+
+    get databaseName(): string {
+        return this._db.name;
+    }
 }
diff --git a/src/matrix/storage/idb/StorageFactory.ts b/src/matrix/storage/idb/StorageFactory.ts
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+import {IDOMStorage} from "./types";
 import {Storage} from "./Storage";
 import { openDatabase, reqAsPromise } from "./utils";
 import { exportSession, importSession, Export } from "./export";
@@ -23,8 +24,8 @@ import { BaseLogger } from "../../../logging/BaseLogger.js";
 import { LogItem } from "../../../logging/LogItem.js";
 
 const sessionName = (sessionId: string) => `hydrogen_session_${sessionId}`;
-const openDatabaseWithSessionId = function(sessionId: string, idbFactory: IDBFactory, log: LogItem) {
-    const create = (db, txn, oldVersion, version) => createStores(db, txn, oldVersion, version, log);
+const openDatabaseWithSessionId = function(sessionId: string, idbFactory: IDBFactory, localStorage: IDOMStorage, log: LogItem) {
+    const create = (db, txn, oldVersion, version) => createStores(db, txn, oldVersion, version, localStorage, log);
     return openDatabase(sessionName(sessionId), create, schema.length, idbFactory);
 }
 
@@ -52,12 +53,14 @@ async function requestPersistedStorage(): Promise<boolean> {
 export class StorageFactory {
     private _serviceWorkerHandler: ServiceWorkerHandler;
     private _idbFactory: IDBFactory;
-    private _IDBKeyRange: typeof IDBKeyRange
+    private _IDBKeyRange: typeof IDBKeyRange;
+    private _localStorage: IDOMStorage;
 
-    constructor(serviceWorkerHandler: ServiceWorkerHandler, idbFactory: IDBFactory = window.indexedDB, _IDBKeyRange = window.IDBKeyRange) {
+    constructor(serviceWorkerHandler: ServiceWorkerHandler, idbFactory: IDBFactory = window.indexedDB, _IDBKeyRange = window.IDBKeyRange, localStorage: IDOMStorage = window.localStorage) {
         this._serviceWorkerHandler = serviceWorkerHandler;
         this._idbFactory = idbFactory;
         this._IDBKeyRange = _IDBKeyRange;
+        this._localStorage = localStorage;
     }
 
     async create(sessionId: string, log: LogItem): Promise<Storage> {
@@ -70,8 +73,8 @@ export class StorageFactory {
         });
 
         const hasWebkitEarlyCloseTxnBug = await detectWebkitEarlyCloseTxnBug(this._idbFactory);
-        const db = await openDatabaseWithSessionId(sessionId, this._idbFactory, log);
-        return new Storage(db, this._idbFactory, this._IDBKeyRange, hasWebkitEarlyCloseTxnBug, log.logger);
+        const db = await openDatabaseWithSessionId(sessionId, this._idbFactory, this._localStorage, log);
+        return new Storage(db, this._idbFactory, this._IDBKeyRange, hasWebkitEarlyCloseTxnBug, this._localStorage, log.logger);
     }
 
     delete(sessionId: string): Promise<IDBDatabase> {
@@ -81,21 +84,22 @@ export class StorageFactory {
     }
 
     async export(sessionId: string, log: LogItem): Promise<Export> {
-        const db = await openDatabaseWithSessionId(sessionId, this._idbFactory, log);
+        const db = await openDatabaseWithSessionId(sessionId, this._idbFactory, this._localStorage, log);
         return await exportSession(db);
     }
 
     async import(sessionId: string, data: Export, log: LogItem): Promise<void> {
-        const db = await openDatabaseWithSessionId(sessionId, this._idbFactory, log);
+        const db = await openDatabaseWithSessionId(sessionId, this._idbFactory, this._localStorage, log);
         return await importSession(db, data);
     }
 }
 
-async function createStores(db: IDBDatabase, txn: IDBTransaction, oldVersion: number | null, version: number, log: LogItem): Promise<void> {
+async function createStores(db: IDBDatabase, txn: IDBTransaction, oldVersion: number | null, version: number, localStorage: IDOMStorage, log: LogItem): Promise<void> {
     const startIdx = oldVersion || 0;
     return log.wrap({l: "storage migration", oldVersion, version}, async log => {
         for(let i = startIdx; i < version; ++i) {
-            await log.wrap(`v${i + 1}`, log => schema[i](db, txn, log));
+            const migrationFunc = schema[i];
+            await log.wrap(`v${i + 1}`, log => migrationFunc(db, txn, localStorage, log));
         }
     });
 }
diff --git a/src/matrix/storage/idb/Transaction.ts b/src/matrix/storage/idb/Transaction.ts
@@ -73,6 +73,10 @@ export class Transaction {
         return this._storage.IDBKeyRange;
     }
 
+    get databaseName(): string {
+        return this._storage.databaseName;
+    }
+
     get logger(): BaseLogger {
         return this._storage.logger;
     }
@@ -94,7 +98,7 @@ export class Transaction {
     }
 
     get session(): SessionStore {
-        return this._store(StoreNames.session, idbStore => new SessionStore(idbStore));
+        return this._store(StoreNames.session, idbStore => new SessionStore(idbStore, this._storage.localStorage));
     }
 
     get roomSummary(): RoomSummaryStore {

diff --git a/src/matrix/storage/idb/schema.ts b/src/matrix/storage/idb/schema.ts
@@ -1,3 +1,4 @@
+import {IDOMStorage} from "./types";
 import {iterateCursor, NOT_DONE, reqAsPromise} from "./utils";
 import {RoomMember, EVENT_TYPE as MEMBER_EVENT_TYPE} from "../../room/members/RoomMember.js";
 import {addRoomToIdentity} from "../../e2ee/DeviceTracker.js";
@@ -7,10 +8,13 @@ import {RoomStateEntry} from "./stores/RoomStateStore";
 import {SessionStore} from "./stores/SessionStore";
 import {encodeScopeTypeKey} from "./stores/OperationStore";
 import {MAX_UNICODE} from "./stores/common";
+import {LogItem} from "../../../logging/LogItem.js";
 
+
+export type MigrationFunc = (db: IDBDatabase, txn: IDBTransaction, localStorage: IDOMStorage, log: LogItem) => Promise<void> | void;
 // FUNCTIONS SHOULD ONLY BE APPENDED!!
 // the index in the array is the database version
-export const schema = [
+export const schema: MigrationFunc[] = [
     createInitialStores,
     createMemberStore,
     migrateSession,
@@ -64,7 +68,7 @@ async function createMemberStore(db: IDBDatabase, txn: IDBTransaction): Promise<
     });
 }
 //v3
-async function migrateSession(db: IDBDatabase, txn: IDBTransaction): Promise<void> {
+async function migrateSession(db: IDBDatabase, txn: IDBTransaction, localStorage: IDOMStorage): Promise<void> {
     const session = txn.objectStore("session");
     try {
         const PRE_MIGRATION_KEY = 1;
@@ -73,7 +77,7 @@ async function migrateSession(db: IDBDatabase, txn: IDBTransaction): Promise<voi
             session.delete(PRE_MIGRATION_KEY);
             const {syncToken, syncFilterId, serverVersions} = entry.value;
             // Cast ok here because only "set" is used and we don't look into return
-            const store = new SessionStore(session as any);
+            const store = new SessionStore(session as any, localStorage);
             store.set("sync", {token: syncToken, filterId: syncFilterId});
             store.set("serverVersions", serverVersions);
         }
@@ -156,7 +160,7 @@ function createTimelineRelationsStore(db: IDBDatabase) : void {
 }
 
 //v11 doesn't change the schema, but ensures all userIdentities have all the roomIds they should (see #470)
-async function fixMissingRoomsInUserIdentities(db, txn, log) {
+async function fixMissingRoomsInUserIdentities(db: IDBDatabase, txn: IDBTransaction, localStorage: IDOMStorage, log: LogItem) {
     const roomSummaryStore = txn.objectStore("roomSummary");
     const trackedRoomIds: string[] = [];
     await iterateCursor<SummaryData>(roomSummaryStore.openCursor(), roomSummary => {

diff --git a/src/matrix/storage/idb/stores/SessionStore.ts b/src/matrix/storage/idb/stores/SessionStore.ts
@@ -14,6 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 import {Store} from "../Store";
+import {IDOMStorage} from "../types";
+import {SESSION_E2EE_KEY_PREFIX} from "../../../e2ee/common.js";
 
 export interface SessionEntry {
     key: string;
@@ -22,9 +24,11 @@ export interface SessionEntry {
 
 export class SessionStore {
     private _sessionStore: Store<SessionEntry>
+    private _localStorage: IDOMStorage;
 
-    constructor(sessionStore: Store<SessionEntry>) {
+    constructor(sessionStore: Store<SessionEntry>, localStorage: IDOMStorage) {
         this._sessionStore = sessionStore;
+        this._localStorage = localStorage;
     }
 
     async get(key: string): Promise<any> {
@@ -34,15 +38,60 @@ export class SessionStore {
         }
     }
 
+    _writeKeyToLocalStorage(key: string, value: any) {
+        // we backup to localStorage so when idb gets cleared for some reason, we don't lose our e2ee identity
+        try {
+            const lsKey = `${this._sessionStore.databaseName}.session.${key}`;
+            const lsValue = JSON.stringify(value);
+            this._localStorage.setItem(lsKey, lsValue);
+        } catch (err) {
+            console.error("could not write to localStorage", err);
+        }
+    }
+
+    writeToLocalStorage() {
+        this._sessionStore.iterateValues(undefined, (value: any, key: string) => {
+            if (key.startsWith(SESSION_E2EE_KEY_PREFIX)) {
+                this._writeKeyToLocalStorage(key, value);
+            }
+            return false;
+        });
+    }
+
+    tryRestoreFromLocalStorage(): boolean {
+        let success = false;
+        const lsPrefix = `${this._sessionStore.databaseName}.session.`;
+        const prefix = `${lsPrefix}${SESSION_E2EE_KEY_PREFIX}`;
+        for(let i = 0; i < this._localStorage.length; i += 1) {
+            const lsKey = this._localStorage.key(i)!;
+            if (lsKey.startsWith(prefix)) {
+                const value = JSON.parse(this._localStorage.getItem(lsKey)!);
+                const key = lsKey.substr(lsPrefix.length);
+                this._sessionStore.put({key, value});
+                success = true;
+            }
+        }
+        return success;
+    }
+
     set(key: string, value: any): void {
+        if (key.startsWith(SESSION_E2EE_KEY_PREFIX)) {
+            this._writeKeyToLocalStorage(key, value);
+        }
         this._sessionStore.put({key, value});
     }
 
     add(key: string, value: any): void {
+        if (key.startsWith(SESSION_E2EE_KEY_PREFIX)) {
+            this._writeKeyToLocalStorage(key, value);
+        }
         this._sessionStore.add({key, value});
     }
 
     remove(key: string): void {
+        if (key.startsWith(SESSION_E2EE_KEY_PREFIX)) {
+            this._localStorage.removeItem(this._sessionStore.databaseName + key);
+        }
         this._sessionStore.delete(key);
     }
 }
diff --git a/src/matrix/storage/idb/types.ts b/src/matrix/storage/idb/types.ts
@@ -0,0 +1,23 @@
+/*
+Copyright 2021 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.
+*/
+
+export interface IDOMStorage {
+    getItem(key: string): string | null;
+    setItem(key: string, value: string): void;
+    removeItem(key: string): void;
+    key(n: number): string | null;
+    readonly length: number;
+}
diff --git a/src/mocks/Storage.ts b/src/mocks/Storage.ts
@@ -16,12 +16,13 @@ limitations under the License.
 
 import {FDBFactory, FDBKeyRange} from "../../lib/fake-indexeddb/index.js";
 import {StorageFactory} from "../matrix/storage/idb/StorageFactory";
+import {IDOMStorage} from "../matrix/storage/idb/types";
 import {Storage} from "../matrix/storage/idb/Storage";
 import {Instance as nullLogger} from "../logging/NullLogger.js";
 import {openDatabase, CreateObjectStore} from "../matrix/storage/idb/utils";
 
 export function createMockStorage(): Promise<Storage> {
-    return new StorageFactory(null as any, new FDBFactory(), FDBKeyRange).create("1", nullLogger.item);
+    return new StorageFactory(null as any, new FDBFactory(), FDBKeyRange, new MockLocalStorage()).create("1", nullLogger.item);
 }
 
 export function createMockDatabase(name: string, createObjectStore: CreateObjectStore, impl: MockIDBImpl): Promise<IDBDatabase> {
@@ -39,3 +40,41 @@ export class MockIDBImpl {
         return FDBKeyRange;
     }
 }
+
+class MockLocalStorage implements IDOMStorage {
+    private _map: Map<string, string>;
+
+    constructor() {
+        this._map = new Map();
+    }
+
+    getItem(key: string): string | null {
+        return this._map.get(key) || null;
+    }
+
+    setItem(key: string, value: string) {
+        this._map.set(key, value);
+    }
+
+    removeItem(key: string): void {
+        this._map.delete(key);
+    }
+
+    get length(): number {
+        return this._map.size;
+    }
+
+    key(n: number): string | null {
+        const it = this._map.keys();
+        let i = -1;
+        let result;
+        while (i < n) {
+            result = it.next();
+            if (result.done) {
+                return null;
+            }
+            i += 1;
+        }
+        return result?.value || null;
+    }
+}