lock on senderKey while enc/decrypting olm sessions

src/matrix/Session.js

@@ -23,6 +23,8 @@ import {DeviceMessageHandler} from "./DeviceMessageHandler.js";
 import {Decryption as OlmDecryption} from "./e2ee/olm/Decryption.js";
 import {Decryption as MegOlmDecryption} from "./e2ee/megolm/Decryption.js";
 import {DeviceTracker} from "./e2ee/DeviceTracker.js";
+import {LockMap} from "../utils/LockMap.js";
+
 const PICKLE_KEY = "DEFAULT_KEY";
 
 export class Session {
@@ -54,6 +56,7 @@ export class Session {
 
     // called once this._e2eeAccount is assigned
     _setupEncryption() {
+        const senderKeyLock = new LockMap();
         const olmDecryption = new OlmDecryption({
             account: this._e2eeAccount,
             pickleKey: PICKLE_KEY,
@@ -61,6 +64,7 @@ export class Session {
             ownUserId: this._user.id,
             storage: this._storage,
             olm: this._olm,
+            senderKeyLock
         });
         const megolmDecryption = new MegOlmDecryption({pickleKey: PICKLE_KEY, olm: this._olm});
         this._deviceMessageHandler.enableEncryption({olmDecryption, megolmDecryption});

src/matrix/e2ee/olm/Decryption.js

@@ -31,14 +31,14 @@ function sortSessions(sessions) {
 }
 
 export class Decryption {
-    constructor({account, pickleKey, now, ownUserId, storage, olm}) {
+    constructor({account, pickleKey, now, ownUserId, storage, olm, senderKeyLock}) {
         this._account = account;
         this._pickleKey = pickleKey;
         this._now = now;
         this._ownUserId = ownUserId;
         this._storage = storage;
         this._olm = olm;
-        this._createOutboundSessionPromise = null;
+        this._senderKeyLock = senderKeyLock;
     }
 
     // we need decryptAll because there is some parallelization we can do for decrypting different sender keys at once
@@ -53,6 +53,12 @@ export class Decryption {
     async decryptAll(events) {
         const eventsPerSenderKey = groupBy(events, event => event.content?.["sender_key"]);
         const timestamp = this._now();
+        // take a lock on all senderKeys so encryption or other calls to decryptAll (should not happen)
+        // don't modify the sessions at the same time
+        const locks = await Promise.all(Array.from(eventsPerSenderKey.keys()).map(senderKey => {
+            return this._senderKeyLock.takeLock(senderKey);
+        }));
+        try {
             const readSessionsTxn = await this._storage.readTxn([this._storage.storeNames.olmSessions]);
             // decrypt events for different sender keys in parallel
             const results = await Promise.all(Array.from(eventsPerSenderKey.entries()).map(([senderKey, events]) => {
@@ -61,7 +67,14 @@ export class Decryption {
             const payloads = results.reduce((all, r) => all.concat(r.payloads), []);
             const errors = results.reduce((all, r) => all.concat(r.errors), []);
             const senderKeyDecryptions = results.map(r => r.senderKeyDecryption);
-        return new DecryptionChanges(senderKeyDecryptions, payloads, errors);
+            return new DecryptionChanges(senderKeyDecryptions, payloads, errors, locks);
+        } catch (err) {
+            // make sure the locks are release if something throws
+            for (const lock of locks) {
+                lock.release();
+            }
+            throw err;
+        }
     }
 
     async _decryptAllForSenderKey(senderKey, events, timestamp, readSessionsTxn) {
@@ -235,11 +248,12 @@ class SenderKeyDecryption {
 }
 
 class DecryptionChanges {
-    constructor(senderKeyDecryptions, payloads, errors, account) {
+    constructor(senderKeyDecryptions, payloads, errors, account, locks) {
         this._senderKeyDecryptions = senderKeyDecryptions;
         this._account = account;    
         this.payloads = payloads;
         this.errors = errors;
+        this._locks = locks;
     }
 
     get hasNewSessions() {
@@ -247,6 +261,7 @@ class DecryptionChanges {
     }
 
     write(txn) {
+        try {
             for (const senderKeyDecryption of this._senderKeyDecryptions) {
                 for (const session of senderKeyDecryption.getModifiedSessions()) {
                     txn.olmSessions.set(session.data);
@@ -268,5 +283,10 @@ class DecryptionChanges {
                     }
                 }
             }
+        } finally {
+            for (const lock of this._locks) {
+                lock.release();
+            }
+        }
     }
 }

src/matrix/e2ee/olm/Encryption.js

@@ -31,7 +31,7 @@ function findFirstSessionId(sessionIds) {
 const OTK_ALGORITHM = "signed_curve25519";
 
 export class Encryption {
-    constructor({account, olm, olmUtil, userId, storage, now, pickleKey}) {
+    constructor({account, olm, olmUtil, userId, storage, now, pickleKey, senderKeyLock}) {
         this._account = account;
         this._olm = olm;
         this._olmUtil = olmUtil;
@@ -39,12 +39,20 @@ export class Encryption {
         this._storage = storage;
         this._now = now;
         this._pickleKey = pickleKey;
+        this._senderKeyLock = senderKeyLock;
     }
 
     async encrypt(type, content, devices, hsApi) {
+        // TODO: see if we can only hold some of the locks until after the /keys/claim call (if needed) 
+        // take a lock on all senderKeys so decryption and other calls to encrypt (should not happen)
+        // don't modify the sessions at the same time
+        const locks = await Promise.all(devices.map(device => {
+            return this._senderKeyLock.takeLock(device.curve25519Key);
+        }));
+        try {
             const {
                 devicesWithoutSession,
-            existingEncryptionTargets
+                existingEncryptionTargets,
             } = await this._findExistingSessions(devices);
         
             const timestamp = this._now(); 
@@ -56,9 +64,6 @@ export class Encryption {
                         devicesWithoutSession, hsApi, timestamp);
                     encryptionTargets = encryptionTargets.concat(newEncryptionTargets);
                 }
-            // TODO: if we read and write in two different txns,
-            // is there a chance we overwrite a session modified by the decryption during sync?
-            // I think so. We'll have to have a lock while sending ...
                 await this._loadSessions(existingEncryptionTargets);
                 encryptionTargets = encryptionTargets.concat(existingEncryptionTargets);
                 const messages = encryptionTargets.map(target => {
@@ -72,6 +77,11 @@ export class Encryption {
                     target.dispose();
                 }
             }
+        } finally {
+            for (const lock of locks) {
+                lock.release();
+            }
+        }
     }
 
     async _findExistingSessions(devices) {

src/utils/Lock.js

@@ -0,0 +1,86 @@
+/*
+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.
+*/
+
+export class Lock {
+    constructor() {
+        this._promise = null;
+        this._resolve = null;
+    }
+
+    take() {
+        if (!this._promise) {
+            this._promise = new Promise(resolve => {
+                this._resolve = resolve;
+            });
+            return true;
+        }
+        return false;
+    }
+
+    get isTaken() {
+        return !!this._promise;
+    }
+
+    release() {
+        if (this._resolve) {
+            this._promise = null;
+            const resolve = this._resolve;
+            this._resolve = null;
+            resolve();
+        }
+    }
+
+    released() {
+        return this._promise;
+    }
+}
+
+export function tests() {
+    return {
+        "taking a lock twice returns false": assert => {
+            const lock = new Lock();
+            assert.equal(lock.take(), true);
+            assert.equal(lock.isTaken, true);
+            assert.equal(lock.take(), false);
+        },
+        "can take a released lock again": assert => {
+            const lock = new Lock();
+            lock.take();
+            lock.release();
+            assert.equal(lock.isTaken, false);
+            assert.equal(lock.take(), true);
+        },
+        "2 waiting for lock, only first one gets it": async assert => {
+            const lock = new Lock();
+            lock.take();
+
+            let first = false;
+            lock.released().then(() => first = lock.take());
+            let second = false;
+            lock.released().then(() => second = lock.take());
+            const promise = lock.released();
+            lock.release();
+            await promise;
+            assert.equal(first, true);
+            assert.equal(second, false);
+        },
+        "await non-taken lock": async assert => {
+            const lock = new Lock();
+            await lock.released();
+            assert(true);
+        }
+    }
+}

src/utils/LockMap.js

@@ -0,0 +1,93 @@
+/*
+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 {Lock} from "./Lock.js";
+
+export class LockMap {
+    constructor() {
+        this._map = new Map();
+    }
+
+    async takeLock(key) {
+        let lock = this._map.get(key);
+        if (lock) {
+            while (!lock.take()) {
+                await lock.released();
+            }
+        } else {
+            lock = new Lock();
+            lock.take();
+            this._map.set(key, lock);
+        }
+        // don't leave old locks lying around
+        lock.released().then(() => {
+            // give others a chance to take the lock first
+            Promise.resolve().then(() => {
+                if (!lock.isTaken) {
+                    this._map.delete(key);
+                }
+            });
+        });
+        return lock;
+    }
+}
+
+export function tests() {
+    return {
+        "taking a lock on the same key blocks": async assert => {
+            const lockMap = new LockMap();
+            const lock = await lockMap.takeLock("foo");
+            let second = false;
+            const prom = lockMap.takeLock("foo").then(() => {
+                second = true;
+            });
+            assert.equal(second, false);
+            // do a delay to make sure prom does not resolve on its own
+            await Promise.resolve();
+            lock.release();
+            await prom;
+            assert.equal(second, true);
+        },
+        "lock is not cleaned up with second request": async assert => {
+            const lockMap = new LockMap();
+            const lock = await lockMap.takeLock("foo");
+            let ranSecond = false;
+            const prom = lockMap.takeLock("foo").then(returnedLock => {
+                ranSecond = true;
+                assert.equal(returnedLock.isTaken, true);
+                // peek into internals, naughty
+                assert.equal(lockMap._map.get("foo"), returnedLock);
+            });
+            lock.release();
+            await prom;
+            // double delay to make sure cleanup logic ran
+            await Promise.resolve();
+            await Promise.resolve();
+            assert.equal(ranSecond, true);
+        },
+        "lock is cleaned up without other request": async assert => {
+            const lockMap = new LockMap();
+            const lock = await lockMap.takeLock("foo");
+            await Promise.resolve();
+            lock.release();
+            // double delay to make sure cleanup logic ran
+            await Promise.resolve();
+            await Promise.resolve();
+            assert.equal(lockMap._map.has("foo"), false);
+        },
+        
+    };
+}