Introduce in-memory transactions for sessions

libtextsecure/storage/unprocessed.js

@@ -21,12 +21,6 @@
     get(id) {
       return textsecure.storage.protocol.getUnprocessedById(id);
     },
-    add(data) {
-      return textsecure.storage.protocol.addUnprocessed(data);
-    },
-    batchAdd(array) {
-      return textsecure.storage.protocol.addMultipleUnprocessed(array);
-    },
     updateAttempts(id, attempts) {
       return textsecure.storage.protocol.updateUnprocessedAttempts(
         id,

package.json

@@ -180,7 +180,8 @@
     "@types/backbone": "1.4.3",
     "@types/better-sqlite3": "5.4.1",
     "@types/blueimp-load-image": "5.14.1",
-    "@types/chai": "4.1.2",
+    "@types/chai": "4.2.18",
+    "@types/chai-as-promised": "7.1.4",
     "@types/classnames": "2.2.3",
     "@types/config": "0.0.34",
     "@types/dashdash": "1.14.0",
@@ -232,7 +233,8 @@
     "babel-core": "7.0.0-bridge.0",
     "babel-loader": "8.0.6",
     "babel-plugin-lodash": "3.3.4",
-    "chai": "4.1.2",
+    "chai": "4.3.4",
+    "chai-as-promised": "7.1.1",
     "core-js": "2.6.9",
     "cross-env": "5.2.0",
     "css-loader": "3.2.0",

ts/LibSignalStores.ts

@@ -24,8 +24,13 @@ import {
 } from '@signalapp/signal-client';
 import { freezePreKey, freezeSignedPreKey } from './SignalProtocolStore';
 
+import { UnprocessedType } from './textsecure/Types.d';
+
 import { typedArrayToArrayBuffer } from './Crypto';
 
+import { assert } from './util/assert';
+import { Lock } from './util/Lock';
+
 function encodedNameFromAddress(address: ProtocolAddress): string {
   const name = address.name();
   const deviceId = address.deviceId();
@@ -33,25 +38,75 @@ function encodedNameFromAddress(address: ProtocolAddress): string {
   return encodedName;
 }
 
+export type SessionsOptions = {
+  readonly transactionOnly?: boolean;
+};
+
 export class Sessions extends SessionStore {
+  private readonly lock = new Lock();
+
+  private inTransaction = false;
+
+  constructor(private readonly options: SessionsOptions = {}) {
+    super();
+  }
+
+  public async transaction<T>(fn: () => Promise<T>): Promise<T> {
+    assert(!this.inTransaction, 'Already in transaction');
+    this.inTransaction = true;
+
+    try {
+      return await window.textsecure.storage.protocol.sessionTransaction(
+        'Sessions.transaction',
+        fn,
+        this.lock
+      );
+    } finally {
+      this.inTransaction = false;
+    }
+  }
+
+  public async addUnprocessed(array: Array<UnprocessedType>): Promise<void> {
+    await window.textsecure.storage.protocol.addMultipleUnprocessed(array, {
+      lock: this.lock,
+    });
+  }
+
+  // SessionStore overrides
+
   async saveSession(
     address: ProtocolAddress,
     record: SessionRecord
   ): Promise<void> {
+    this.checkInTransaction();
+
     await window.textsecure.storage.protocol.storeSession(
       encodedNameFromAddress(address),
-      record
+      record,
+      { lock: this.lock }
     );
   }
 
   async getSession(name: ProtocolAddress): Promise<SessionRecord | null> {
+    this.checkInTransaction();
+
     const encodedName = encodedNameFromAddress(name);
     const record = await window.textsecure.storage.protocol.loadSession(
-      encodedName
+      encodedName,
+      { lock: this.lock }
     );
 
     return record || null;
   }
+
+  // Private
+
+  private checkInTransaction(): void {
+    assert(
+      this.inTransaction || !this.options.transactionOnly,
+      'Accessing session store outside of transaction'
+    );
+  }
 }
 
 export class IdentityKeys extends IdentityKeyStore {

ts/SignalProtocolStore.ts

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
 /* eslint-disable class-methods-use-this */
+/* eslint-disable no-restricted-syntax */
 
 import PQueue from 'p-queue';
 import { isNumber } from 'lodash';
@@ -22,7 +23,9 @@ import {
   fromEncodedBinaryToArrayBuffer,
   typedArrayToArrayBuffer,
 } from './Crypto';
+import { assert } from './util/assert';
 import { isNotNil } from './util/isNotNil';
+import { Lock } from './util/Lock';
 import { isMoreRecentThan } from './util/timestamp';
 import {
   sessionRecordToProtobuf,
@@ -102,9 +105,22 @@ type CacheEntryType<DBType, HydratedType> =
     }
   | { hydrated: true; fromDB: DBType; item: HydratedType };
 
+type MapFields =
+  | 'identityKeys'
+  | 'preKeys'
+  | 'senderKeys'
+  | 'sessions'
+  | 'signedPreKeys';
+
+export type SessionTransactionOptions = {
+  readonly lock?: Lock;
+};
+
+const GLOBAL_LOCK = new Lock();
+
 async function _fillCaches<ID, T extends HasIdType<ID>, HydratedType>(
   object: SignalProtocolStore,
-  field: keyof SignalProtocolStore,
+  field: MapFields,
   itemsPromise: Promise<Array<T>>
 ): Promise<void> {
   const items = await itemsPromise;
@@ -182,6 +198,8 @@ const EventsMixin = (function EventsMixin(this: unknown) {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
 } as any) as typeof window.Backbone.EventsMixin;
 
+type SessionCacheEntry = CacheEntryType<SessionType, SessionRecord>;
+
 export class SignalProtocolStore extends EventsMixin {
   // Enums used across the app
 
@@ -197,7 +215,15 @@ export class SignalProtocolStore extends EventsMixin {
 
   senderKeys?: Map<string, CacheEntryType<SenderKeyType, SenderKeyRecord>>;
 
-  sessions?: Map<string, CacheEntryType<SessionType, SessionRecord>>;
+  sessions?: Map<string, SessionCacheEntry>;
+
+  sessionLock?: Lock;
+
+  sessionLockQueue: Array<() => void> = [];
+
+  pendingSessions = new Map<string, SessionCacheEntry>();
+
+  pendingUnprocessed = new Map<string, UnprocessedType>();
 
   preKeys?: Map<number, CacheEntryType<PreKeyType, PreKeyRecord>>;
 
@@ -562,43 +588,154 @@ export class SignalProtocolStore extends EventsMixin {
 
   // Sessions
 
-  async loadSession(
+  // Re-entrant session transaction routine. Only one session transaction could
-    encodedAddress: string
+  // be running at the same time.
-  ): Promise<SessionRecord | undefined> {
+  //
-    if (!this.sessions) {
+  // While in transaction:
-      throw new Error('loadSession: this.sessions not yet cached!');
+  //
+  // - `storeSession()` adds the updated session to the `pendingSessions`
+  // - `loadSession()` looks up the session first in `pendingSessions` and only
+  //   then in the main `sessions` store
+  //
+  // When transaction ends:
+  //
+  // - successfully: pending session stores are batched into the database
+  // - with an error: pending session stores are reverted
+  async sessionTransaction<T>(
+    name: string,
+    body: () => Promise<T>,
+    lock: Lock = GLOBAL_LOCK
+  ): Promise<T> {
+    // Allow re-entering from LibSignalStores
+    const isNested = this.sessionLock === lock;
+    if (this.sessionLock && !isNested) {
+      window.log.info(`sessionTransaction(${name}): sessions locked, waiting`);
+      await new Promise<void>(resolve => this.sessionLockQueue.push(resolve));
     }
 
-    if (encodedAddress === null || encodedAddress === undefined) {
+    if (!isNested) {
-      throw new Error('loadSession: encodedAddress was undefined/null');
+      if (lock !== GLOBAL_LOCK) {
+        window.log.info(`sessionTransaction(${name}): enter`);
+      }
+      this.sessionLock = lock;
     }
 
+    let result: T;
     try {
-      const id = await normalizeEncodedAddress(encodedAddress);
+      result = await body();
-      const entry = this.sessions.get(id);
-
-      if (!entry) {
-        return undefined;
-      }
-
-      if (entry.hydrated) {
-        return entry.item;
-      }
-
-      const item = await this._maybeMigrateSession(entry.fromDB);
-      this.sessions.set(id, {
-        hydrated: true,
-        item,
-        fromDB: entry.fromDB,
-      });
-      return item;
     } catch (error) {
-      const errorString = error && error.stack ? error.stack : error;
+      if (!isNested) {
-      window.log.error(
+        await this.revertSessions(name, error);
-        `loadSession: failed to load session ${encodedAddress}: ${errorString}`
+        this.releaseSessionLock();
-      );
+      }
-      return undefined;
+      throw error;
     }
+
+    if (!isNested) {
+      await this.commitSessions(name);
+      this.releaseSessionLock();
+    }
+
+    return result;
+  }
+
+  private async commitSessions(name: string): Promise<void> {
+    const { pendingSessions, pendingUnprocessed } = this;
+
+    if (pendingSessions.size === 0 && pendingUnprocessed.size === 0) {
+      return;
+    }
+
+    window.log.info(
+      `commitSessions(${name}): pending sessions ${pendingSessions.size} ` +
+        `pending unprocessed ${pendingUnprocessed.size}`
+    );
+
+    this.pendingSessions = new Map();
+    this.pendingUnprocessed = new Map();
+
+    // Commit both unprocessed and sessions in the same database transaction
+    // to unroll both on error.
+    await window.Signal.Data.commitSessionsAndUnprocessed({
+      sessions: Array.from(pendingSessions.values()).map(
+        ({ fromDB }) => fromDB
+      ),
+      unprocessed: Array.from(pendingUnprocessed.values()),
+    });
+
+    const { sessions } = this;
+    assert(sessions !== undefined, "Can't commit unhydrated storage");
+
+    // Apply changes to in-memory storage after successful DB write.
+    pendingSessions.forEach((value, key) => {
+      sessions.set(key, value);
+    });
+  }
+
+  private async revertSessions(name: string, error: Error): Promise<void> {
+    window.log.info(
+      `revertSessions(${name}): pending size ${this.pendingSessions.size}`,
+      error && error.stack
+    );
+    this.pendingSessions.clear();
+    this.pendingUnprocessed.clear();
+  }
+
+  private releaseSessionLock(): void {
+    this.sessionLock = undefined;
+    const next = this.sessionLockQueue.shift();
+    if (next) {
+      next();
+    }
+  }
+
+  async loadSession(
+    encodedAddress: string,
+    { lock }: SessionTransactionOptions = {}
+  ): Promise<SessionRecord | undefined> {
+    return this.sessionTransaction(
+      'loadSession',
+      async () => {
+        if (!this.sessions) {
+          throw new Error('loadSession: this.sessions not yet cached!');
+        }
+
+        if (encodedAddress === null || encodedAddress === undefined) {
+          throw new Error('loadSession: encodedAddress was undefined/null');
+        }
+
+        try {
+          const id = await normalizeEncodedAddress(encodedAddress);
+          const map = this.pendingSessions.has(id)
+            ? this.pendingSessions
+            : this.sessions;
+          const entry = map.get(id);
+
+          if (!entry) {
+            return undefined;
+          }
+
+          if (entry.hydrated) {
+            return entry.item;
+          }
+
+          const item = await this._maybeMigrateSession(entry.fromDB);
+          map.set(id, {
+            hydrated: true,
+            item,
+            fromDB: entry.fromDB,
+          });
+          return item;
+        } catch (error) {
+          const errorString = error && error.stack ? error.stack : error;
+          window.log.error(
+            `loadSession: failed to load session ${encodedAddress}: ${errorString}`
+          );
+          return undefined;
+        }
+      },
+      lock
+    );
   }
 
   private async _maybeMigrateSession(
@@ -643,139 +780,155 @@ export class SignalProtocolStore extends EventsMixin {
 
   async storeSession(
     encodedAddress: string,
-    record: SessionRecord
+    record: SessionRecord,
+    { lock }: SessionTransactionOptions = {}
   ): Promise<void> {
-    if (!this.sessions) {
+    await this.sessionTransaction(
-      throw new Error('storeSession: this.sessions not yet cached!');
+      'storeSession',
-    }
+      async () => {
+        if (!this.sessions) {
+          throw new Error('storeSession: this.sessions not yet cached!');
+        }
 
-    if (encodedAddress === null || encodedAddress === undefined) {
+        if (encodedAddress === null || encodedAddress === undefined) {
-      throw new Error('storeSession: encodedAddress was undefined/null');
+          throw new Error('storeSession: encodedAddress was undefined/null');
-    }
+        }
-    const unencoded = window.textsecure.utils.unencodeNumber(encodedAddress);
+        const unencoded = window.textsecure.utils.unencodeNumber(
-    const deviceId = parseInt(unencoded[1], 10);
+          encodedAddress
+        );
+        const deviceId = parseInt(unencoded[1], 10);
 
-    try {
+        try {
-      const id = await normalizeEncodedAddress(encodedAddress);
+          const id = await normalizeEncodedAddress(encodedAddress);
-      const fromDB = {
+          const fromDB = {
-        id,
+            id,
-        version: 2,
+            version: 2,
-        conversationId: window.textsecure.utils.unencodeNumber(id)[0],
+            conversationId: window.textsecure.utils.unencodeNumber(id)[0],
-        deviceId,
+            deviceId,
-        record: record.serialize().toString('base64'),
+            record: record.serialize().toString('base64'),
-      };
+          };
 
-      await window.Signal.Data.createOrUpdateSession(fromDB);
+          const newSession = {
-      this.sessions.set(id, {
+            hydrated: true,
-        hydrated: true,
+            fromDB,
-        fromDB,
+            item: record,
-        item: record,
+          };
-      });
+
-    } catch (error) {
+          this.pendingSessions.set(id, newSession);
-      const errorString = error && error.stack ? error.stack : error;
+        } catch (error) {
-      window.log.error(
+          const errorString = error && error.stack ? error.stack : error;
-        `storeSession: Save failed fo ${encodedAddress}: ${errorString}`
+          window.log.error(
-      );
+            `storeSession: Save failed fo ${encodedAddress}: ${errorString}`
-      throw error;
+          );
-    }
+          throw error;
+        }
+      },
+      lock
+    );
   }
 
   async getDeviceIds(identifier: string): Promise<Array<number>> {
-    if (!this.sessions) {
+    return this.sessionTransaction('getDeviceIds', async () => {
-      throw new Error('getDeviceIds: this.sessions not yet cached!');
+      if (!this.sessions) {
-    }
+        throw new Error('getDeviceIds: this.sessions not yet cached!');
-    if (identifier === null || identifier === undefined) {
+      }
-      throw new Error('getDeviceIds: identifier was undefined/null');
+      if (identifier === null || identifier === undefined) {
-    }
+        throw new Error('getDeviceIds: identifier was undefined/null');
-
-    try {
-      const id = window.ConversationController.getConversationId(identifier);
-      if (!id) {
-        throw new Error(
-          `getDeviceIds: No conversationId found for identifier ${identifier}`
-        );
       }
 
-      const allSessions = Array.from(this.sessions.values());
+      try {
-      const entries = allSessions.filter(
+        const id = window.ConversationController.getConversationId(identifier);
-        session => session.fromDB.conversationId === id
+        if (!id) {
-      );
+          throw new Error(
-      const openIds = await Promise.all(
+            `getDeviceIds: No conversationId found for identifier ${identifier}`
-        entries.map(async entry => {
+          );
-          if (entry.hydrated) {
+        }
-            const record = entry.item;
+
+        const allSessions = this._getAllSessions();
+        const entries = allSessions.filter(
+          session => session.fromDB.conversationId === id
+        );
+        const openIds = await Promise.all(
+          entries.map(async entry => {
+            if (entry.hydrated) {
+              const record = entry.item;
+              if (record.hasCurrentState()) {
+                return entry.fromDB.deviceId;
+              }
+
+              return undefined;
+            }
+
+            const record = await this._maybeMigrateSession(entry.fromDB);
             if (record.hasCurrentState()) {
               return entry.fromDB.deviceId;
             }
 
             return undefined;
-          }
+          })
+        );
 
-          const record = await this._maybeMigrateSession(entry.fromDB);
+        return openIds.filter(isNotNil);
-          if (record.hasCurrentState()) {
+      } catch (error) {
-            return entry.fromDB.deviceId;
+        window.log.error(
-          }
+          `getDeviceIds: Failed to get device ids for identifier ${identifier}`,
+          error && error.stack ? error.stack : error
+        );
+      }
 
-          return undefined;
+      return [];
-        })
+    });
-      );
-
-      return openIds.filter(isNotNil);
-    } catch (error) {
-      window.log.error(
-        `getDeviceIds: Failed to get device ids for identifier ${identifier}`,
-        error && error.stack ? error.stack : error
-      );
-    }
-
-    return [];
   }
 
   async removeSession(encodedAddress: string): Promise<void> {
-    if (!this.sessions) {
+    return this.sessionTransaction('removeSession', async () => {
-      throw new Error('removeSession: this.sessions not yet cached!');
+      if (!this.sessions) {
-    }
+        throw new Error('removeSession: this.sessions not yet cached!');
+      }
 
-    window.log.info('removeSession: deleting session for', encodedAddress);
+      window.log.info('removeSession: deleting session for', encodedAddress);
-    try {
+      try {
-      const id = await normalizeEncodedAddress(encodedAddress);
+        const id = await normalizeEncodedAddress(encodedAddress);
-      await window.Signal.Data.removeSessionById(id);
+        await window.Signal.Data.removeSessionById(id);
-      this.sessions.delete(id);
+        this.sessions.delete(id);
-    } catch (e) {
+        this.pendingSessions.delete(id);
-      window.log.error(
+      } catch (e) {
-        `removeSession: Failed to delete session for ${encodedAddress}`
+        window.log.error(
-      );
+          `removeSession: Failed to delete session for ${encodedAddress}`
-    }
+        );
+      }
+    });
   }
 
   async removeAllSessions(identifier: string): Promise<void> {
-    if (!this.sessions) {
+    return this.sessionTransaction('removeAllSessions', async () => {
-      throw new Error('removeAllSessions: this.sessions not yet cached!');
+      if (!this.sessions) {
-    }
+        throw new Error('removeAllSessions: this.sessions not yet cached!');
-
-    if (identifier === null || identifier === undefined) {
-      throw new Error('removeAllSessions: identifier was undefined/null');
-    }
-
-    window.log.info('removeAllSessions: deleting sessions for', identifier);
-
-    const id = window.ConversationController.getConversationId(identifier);
-
-    const entries = Array.from(this.sessions.values());
-
-    for (let i = 0, max = entries.length; i < max; i += 1) {
-      const entry = entries[i];
-      if (entry.fromDB.conversationId === id) {
-        this.sessions.delete(entry.fromDB.id);
       }
-    }
 
-    await window.Signal.Data.removeSessionsByConversation(identifier);
+      if (identifier === null || identifier === undefined) {
+        throw new Error('removeAllSessions: identifier was undefined/null');
+      }
+
+      window.log.info('removeAllSessions: deleting sessions for', identifier);
+
+      const id = window.ConversationController.getConversationId(identifier);
+
+      const entries = Array.from(this.sessions.values());
+
+      for (let i = 0, max = entries.length; i < max; i += 1) {
+        const entry = entries[i];
+        if (entry.fromDB.conversationId === id) {
+          this.sessions.delete(entry.fromDB.id);
+          this.pendingSessions.delete(entry.fromDB.id);
+        }
+      }
+
+      await window.Signal.Data.removeSessionsByConversation(identifier);
+    });
   }
 
-  private async _archiveSession(
+  private async _archiveSession(entry?: SessionCacheEntry) {
-    entry?: CacheEntryType<SessionType, SessionRecord>
-  ) {
     if (!entry) {
       return;
     }
@@ -796,74 +949,87 @@ export class SignalProtocolStore extends EventsMixin {
   }
 
   async archiveSession(encodedAddress: string): Promise<void> {
-    if (!this.sessions) {
+    return this.sessionTransaction('archiveSession', async () => {
-      throw new Error('archiveSession: this.sessions not yet cached!');
+      if (!this.sessions) {
-    }
+        throw new Error('archiveSession: this.sessions not yet cached!');
+      }
 
-    window.log.info(`archiveSession: session for ${encodedAddress}`);
+      window.log.info(`archiveSession: session for ${encodedAddress}`);
 
-    const id = await normalizeEncodedAddress(encodedAddress);
+      const id = await normalizeEncodedAddress(encodedAddress);
-    const entry = this.sessions.get(id);
 
-    await this._archiveSession(entry);
+      const entry = this.pendingSessions.get(id) || this.sessions.get(id);
+
+      await this._archiveSession(entry);
+    });
   }
 
   async archiveSiblingSessions(encodedAddress: string): Promise<void> {
-    if (!this.sessions) {
+    return this.sessionTransaction('archiveSiblingSessions', async () => {
-      throw new Error('archiveSiblingSessions: this.sessions not yet cached!');
+      if (!this.sessions) {
-    }
+        throw new Error(
+          'archiveSiblingSessions: this.sessions not yet cached!'
+        );
+      }
 
-    window.log.info(
+      window.log.info(
-      'archiveSiblingSessions: archiving sibling sessions for',
+        'archiveSiblingSessions: archiving sibling sessions for',
-      encodedAddress
+        encodedAddress
-    );
+      );
 
-    const id = await normalizeEncodedAddress(encodedAddress);
+      const id = await normalizeEncodedAddress(encodedAddress);
-    const [identifier, deviceId] = window.textsecure.utils.unencodeNumber(id);
+      const [identifier, deviceId] = window.textsecure.utils.unencodeNumber(id);
-    const deviceIdNumber = parseInt(deviceId, 10);
+      const deviceIdNumber = parseInt(deviceId, 10);
 
-    const allEntries = Array.from(this.sessions.values());
+      const allEntries = this._getAllSessions();
-    const entries = allEntries.filter(
+      const entries = allEntries.filter(
-      entry =>
+        entry =>
-        entry.fromDB.conversationId === identifier &&
+          entry.fromDB.conversationId === identifier &&
-        entry.fromDB.deviceId !== deviceIdNumber
+          entry.fromDB.deviceId !== deviceIdNumber
-    );
+      );
 
-    await Promise.all(
+      await Promise.all(
-      entries.map(async entry => {
+        entries.map(async entry => {
-        await this._archiveSession(entry);
+          await this._archiveSession(entry);
-      })
+        })
-    );
+      );
+    });
   }
 
   async archiveAllSessions(identifier: string): Promise<void> {
-    if (!this.sessions) {
+    return this.sessionTransaction('archiveAllSessions', async () => {
-      throw new Error('archiveAllSessions: this.sessions not yet cached!');
+      if (!this.sessions) {
-    }
+        throw new Error('archiveAllSessions: this.sessions not yet cached!');
+      }
 
-    window.log.info(
+      window.log.info(
-      'archiveAllSessions: archiving all sessions for',
+        'archiveAllSessions: archiving all sessions for',
-      identifier
+        identifier
-    );
+      );
 
-    const id = window.ConversationController.getConversationId(identifier);
+      const id = window.ConversationController.getConversationId(identifier);
-    const allEntries = Array.from(this.sessions.values());
-    const entries = allEntries.filter(
-      entry => entry.fromDB.conversationId === id
-    );
 
-    await Promise.all(
+      const allEntries = this._getAllSessions();
-      entries.map(async entry => {
+      const entries = allEntries.filter(
-        await this._archiveSession(entry);
+        entry => entry.fromDB.conversationId === id
-      })
+      );
-    );
+
+      await Promise.all(
+        entries.map(async entry => {
+          await this._archiveSession(entry);
+        })
+      );
+    });
   }
 
   async clearSessionStore(): Promise<void> {
-    if (this.sessions) {
+    return this.sessionTransaction('clearSessionStore', async () => {
-      this.sessions.clear();
+      if (this.sessions) {
-    }
+        this.sessions.clear();
-    window.Signal.Data.removeAllSessions();
+      }
+      this.pendingSessions.clear();
+      await window.Signal.Data.removeAllSessions();
+    });
   }
 
   // Identity Keys
@@ -1403,56 +1569,92 @@ export class SignalProtocolStore extends EventsMixin {
 
   // Not yet processed messages - for resiliency
   getUnprocessedCount(): Promise<number> {
-    return window.Signal.Data.getUnprocessedCount();
+    return this.sessionTransaction('getUnprocessedCount', async () => {
+      this._checkNoPendingUnprocessed();
+      return window.Signal.Data.getUnprocessedCount();
+    });
   }
 
   getAllUnprocessed(): Promise<Array<UnprocessedType>> {
-    return window.Signal.Data.getAllUnprocessed();
+    return this.sessionTransaction('getAllUnprocessed', async () => {
+      this._checkNoPendingUnprocessed();
+      return window.Signal.Data.getAllUnprocessed();
+    });
   }
 
   getUnprocessedById(id: string): Promise<UnprocessedType | undefined> {
-    return window.Signal.Data.getUnprocessedById(id);
+    return this.sessionTransaction('getUnprocessedById', async () => {
-  }
+      this._checkNoPendingUnprocessed();
-
+      return window.Signal.Data.getUnprocessedById(id);
-  addUnprocessed(data: UnprocessedType): Promise<string> {
-    // We need to pass forceSave because the data has an id already, which will cause
-    //   an update instead of an insert.
-    return window.Signal.Data.saveUnprocessed(data, {
-      forceSave: true,
     });
   }
 
-  addMultipleUnprocessed(array: Array<UnprocessedType>): Promise<void> {
+  addUnprocessed(
-    // We need to pass forceSave because the data has an id already, which will cause
+    data: UnprocessedType,
-    //   an update instead of an insert.
+    { lock }: SessionTransactionOptions = {}
-    return window.Signal.Data.saveUnprocesseds(array, {
+  ): Promise<void> {
-      forceSave: true,
+    return this.sessionTransaction(
-    });
+      'addUnprocessed',
+      async () => {
+        this.pendingUnprocessed.set(data.id, data);
+      },
+      lock
+    );
+  }
+
+  addMultipleUnprocessed(
+    array: Array<UnprocessedType>,
+    { lock }: SessionTransactionOptions = {}
+  ): Promise<void> {
+    return this.sessionTransaction(
+      'addMultipleUnprocessed',
+      async () => {
+        for (const elem of array) {
+          this.pendingUnprocessed.set(elem.id, elem);
+        }
+      },
+      lock
+    );
   }
 
   updateUnprocessedAttempts(id: string, attempts: number): Promise<void> {
-    return window.Signal.Data.updateUnprocessedAttempts(id, attempts);
+    return this.sessionTransaction('updateUnprocessedAttempts', async () => {
+      this._checkNoPendingUnprocessed();
+      await window.Signal.Data.updateUnprocessedAttempts(id, attempts);
+    });
   }
 
   updateUnprocessedWithData(
     id: string,
     data: UnprocessedUpdateType
   ): Promise<void> {
-    return window.Signal.Data.updateUnprocessedWithData(id, data);
+    return this.sessionTransaction('updateUnprocessedWithData', async () => {
+      this._checkNoPendingUnprocessed();
+      await window.Signal.Data.updateUnprocessedWithData(id, data);
+    });
   }
 
   updateUnprocessedsWithData(
     items: Array<{ id: string; data: UnprocessedUpdateType }>
   ): Promise<void> {
-    return window.Signal.Data.updateUnprocessedsWithData(items);
+    return this.sessionTransaction('updateUnprocessedsWithData', async () => {
+      this._checkNoPendingUnprocessed();
+      await window.Signal.Data.updateUnprocessedsWithData(items);
+    });
   }
 
   removeUnprocessed(idOrArray: string | Array<string>): Promise<void> {
-    return window.Signal.Data.removeUnprocessed(idOrArray);
+    return this.sessionTransaction('removeUnprocessed', async () => {
+      this._checkNoPendingUnprocessed();
+      await window.Signal.Data.removeUnprocessed(idOrArray);
+    });
   }
 
   removeAllUnprocessed(): Promise<void> {
-    return window.Signal.Data.removeAllUnprocessed();
+    return this.sessionTransaction('removeAllUnprocessed', async () => {
+      this._checkNoPendingUnprocessed();
+      await window.Signal.Data.removeAllUnprocessed();
+    });
   }
 
   async removeAllData(): Promise<void> {
@@ -1473,6 +1675,30 @@ export class SignalProtocolStore extends EventsMixin {
     window.storage.reset();
     await window.storage.fetch();
   }
+
+  private _getAllSessions(): Array<SessionCacheEntry> {
+    const union = new Map<string, SessionCacheEntry>();
+
+    this.sessions?.forEach((value, key) => {
+      union.set(key, value);
+    });
+    this.pendingSessions.forEach((value, key) => {
+      union.set(key, value);
+    });
+
+    return Array.from(union.values());
+  }
+
+  private _checkNoPendingUnprocessed(): void {
+    assert(
+      !this.sessionLock || this.sessionLock === GLOBAL_LOCK,
+      "Can't use this function with a global lock"
+    );
+    assert(
+      this.pendingUnprocessed.size === 0,
+      'Missing support for pending unprocessed'
+    );
+  }
 }
 
 window.SignalProtocolStore = SignalProtocolStore;

ts/sql/Client.ts

@@ -142,6 +142,7 @@ const dataInterface: ClientInterface = {
 
   createOrUpdateSession,
   createOrUpdateSessions,
+  commitSessionsAndUnprocessed,
   getSessionById,
   getSessionsById,
   bulkAddSessions,
@@ -767,6 +768,12 @@ async function createOrUpdateSession(data: SessionType) {
 async function createOrUpdateSessions(array: Array<SessionType>) {
   await channels.createOrUpdateSessions(array);
 }
+async function commitSessionsAndUnprocessed(options: {
+  sessions: Array<SessionType>;
+  unprocessed: Array<UnprocessedType>;
+}) {
+  await channels.commitSessionsAndUnprocessed(options);
+}
 async function getSessionById(id: string) {
   const session = await channels.getSessionById(id);
 
@@ -1353,22 +1360,14 @@ async function getUnprocessedById(id: string) {
   return channels.getUnprocessedById(id);
 }
 
-async function saveUnprocessed(
+async function saveUnprocessed(data: UnprocessedType) {
-  data: UnprocessedType,
+  const id = await channels.saveUnprocessed(_cleanData(data));
-  { forceSave }: { forceSave?: boolean } = {}
-) {
-  const id = await channels.saveUnprocessed(_cleanData(data), { forceSave });
 
   return id;
 }
 
-async function saveUnprocesseds(
+async function saveUnprocesseds(arrayOfUnprocessed: Array<UnprocessedType>) {
-  arrayOfUnprocessed: Array<UnprocessedType>,
+  await channels.saveUnprocesseds(_cleanData(arrayOfUnprocessed));
-  { forceSave }: { forceSave?: boolean } = {}
-) {
-  await channels.saveUnprocesseds(_cleanData(arrayOfUnprocessed), {
-    forceSave,
-  });
 }
 
 async function updateUnprocessedAttempts(id: string, attempts: number) {

ts/sql/Interface.ts

@@ -131,11 +131,11 @@ export type UnprocessedType = {
   timestamp: number;
   version: number;
   attempts: number;
-  envelope: string;
+  envelope?: string;
 
   source?: string;
   sourceUuid?: string;
-  sourceDevice?: string;
+  sourceDevice?: number;
   serverTimestamp?: number;
   decrypted?: string;
 };
@@ -188,6 +188,10 @@ export type DataInterface = {
 
   createOrUpdateSession: (data: SessionType) => Promise<void>;
   createOrUpdateSessions: (array: Array<SessionType>) => Promise<void>;
+  commitSessionsAndUnprocessed(options: {
+    sessions: Array<SessionType>;
+    unprocessed: Array<UnprocessedType>;
+  }): Promise<void>;
   getSessionById: (id: string) => Promise<SessionType | undefined>;
   getSessionsById: (conversationId: string) => Promise<Array<SessionType>>;
   bulkAddSessions: (array: Array<SessionType>) => Promise<void>;

ts/sql/Server.ts

@@ -133,6 +133,7 @@ const dataInterface: ServerInterface = {
 
   createOrUpdateSession,
   createOrUpdateSessions,
+  commitSessionsAndUnprocessed,
   getSessionById,
   getSessionsById,
   bulkAddSessions,
@@ -2217,6 +2218,26 @@ async function createOrUpdateSessions(
   })();
 }
 
+async function commitSessionsAndUnprocessed({
+  sessions,
+  unprocessed,
+}: {
+  sessions: Array<SessionType>;
+  unprocessed: Array<UnprocessedType>;
+}): Promise<void> {
+  const db = getInstance();
+
+  db.transaction(() => {
+    for (const item of sessions) {
+      createOrUpdateSession(item);
+    }
+
+    for (const item of unprocessed) {
+      saveUnprocessedSync(item);
+    }
+  })();
+}
+
 async function getSessionById(id: string): Promise<SessionType | undefined> {
   return getById(SESSIONS_TABLE, id);
 }
@@ -3948,82 +3969,79 @@ async function getTapToViewMessagesNeedingErase(): Promise<Array<MessageType>> {
   return rows.map(row => jsonToObject(row.json));
 }
 
-function saveUnprocessedSync(
+function saveUnprocessedSync(data: UnprocessedType): string {
-  data: UnprocessedType,
-  { forceSave }: { forceSave?: boolean } = {}
-): string {
   const db = getInstance();
-  const { id, timestamp, version, attempts, envelope } = data;
+  const {
+    id,
+    timestamp,
+    version,
+    attempts,
+    envelope,
+    source,
+    sourceUuid,
+    sourceDevice,
+    serverTimestamp,
+    decrypted,
+  } = data;
   if (!id) {
     throw new Error('saveUnprocessed: id was falsey');
   }
 
-  if (forceSave) {
-    prepare(
-      db,
-      `
-      INSERT INTO unprocessed (
-        id,
-        timestamp,
-        version,
-        attempts,
-        envelope
-      ) values (
-        $id,
-        $timestamp,
-        $version,
-        $attempts,
-        $envelope
-      );
-      `
-    ).run({
-      id,
-      timestamp,
-      version,
-      attempts,
-      envelope,
-    });
-
-    return id;
-  }
-
   prepare(
     db,
     `
-    UPDATE unprocessed SET
+    INSERT OR REPLACE INTO unprocessed (
-      timestamp = $timestamp,
+      id,
-      version = $version,
+      timestamp,
-      attempts = $attempts,
+      version,
-      envelope = $envelope
+      attempts,
-    WHERE id = $id;
+      envelope,
+      source,
+      sourceUuid,
+      sourceDevice,
+      serverTimestamp,
+      decrypted
+    ) values (
+      $id,
+      $timestamp,
+      $version,
+      $attempts,
+      $envelope,
+      $source,
+      $sourceUuid,
+      $sourceDevice,
+      $serverTimestamp,
+      $decrypted
+    );
     `
   ).run({
     id,
     timestamp,
     version,
     attempts,
-    envelope,
+    envelope: envelope || null,
+    source: source || null,
+    sourceUuid: sourceUuid || null,
+    sourceDevice: sourceDevice || null,
+    serverTimestamp: serverTimestamp || null,
+    decrypted: decrypted || null,
   });
 
   return id;
 }
 
-async function saveUnprocessed(
+async function saveUnprocessed(data: UnprocessedType): Promise<string> {
-  data: UnprocessedType,
+  return saveUnprocessedSync(data);
-  options: { forceSave?: boolean } = {}
-): Promise<string> {
-  return saveUnprocessedSync(data, options);
 }
 
 async function saveUnprocesseds(
-  arrayOfUnprocessed: Array<UnprocessedType>,
+  arrayOfUnprocessed: Array<UnprocessedType>
-  { forceSave }: { forceSave?: boolean } = {}
 ): Promise<void> {
   const db = getInstance();
 
   db.transaction(() => {
     for (const unprocessed of arrayOfUnprocessed) {
-      assertSync(saveUnprocessedSync(unprocessed, { forceSave }));
+      assertSync(saveUnprocessedSync(unprocessed));
     }
   })();
 }

ts/test-electron/SignalProtocolStore_test.ts

@@ -3,7 +3,8 @@
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import { assert } from 'chai';
+import chai, { assert } from 'chai';
+import chaiAsPromised from 'chai-as-promised';
 import {
   Direction,
   SenderKeyRecord,
@@ -12,12 +13,15 @@ import {
 
 import { signal } from '../protobuf/compiled';
 import { sessionStructureToArrayBuffer } from '../util/sessionTranslation';
+import { Lock } from '../util/Lock';
 
 import { getRandomBytes, constantTimeEqual } from '../Crypto';
 import { clampPrivateKey, setPublicKeyTypeByte } from '../Curve';
 import { SignalProtocolStore } from '../SignalProtocolStore';
 import { IdentityKeyType, KeyPairType } from '../textsecure/Types.d';
 
+chai.use(chaiAsPromised);
+
 const {
   RecordStructure,
   SessionStructure,
@@ -1233,7 +1237,7 @@ describe('SignalProtocolStore', () => {
       await store.removeAllSessions(number);
 
       const records = await Promise.all(
-        devices.map(store.loadSession.bind(store))
+        devices.map(device => store.loadSession(device))
       );
 
       for (let i = 0, max = records.length; i < max; i += 1) {
@@ -1276,6 +1280,96 @@ describe('SignalProtocolStore', () => {
     });
   });
 
+  describe('sessionTransaction', () => {
+    beforeEach(async () => {
+      await store.removeAllUnprocessed();
+      await store.removeAllSessions(number);
+    });
+
+    it('commits session stores and unprocessed on success', async () => {
+      const id = `${number}.1`;
+      const testRecord = getSessionRecord();
+
+      await store.sessionTransaction('test', async () => {
+        await store.storeSession(id, testRecord);
+
+        await store.addUnprocessed({
+          id: '2-two',
+          envelope: 'second',
+          timestamp: 2,
+          version: 2,
+          attempts: 0,
+        });
+        assert.equal(await store.loadSession(id), testRecord);
+      });
+
+      assert.equal(await store.loadSession(id), testRecord);
+
+      const allUnprocessed = await store.getAllUnprocessed();
+      assert.deepEqual(
+        allUnprocessed.map(({ envelope }) => envelope),
+        ['second']
+      );
+    });
+
+    it('reverts session stores and unprocessed on error', async () => {
+      const id = `${number}.1`;
+      const testRecord = getSessionRecord();
+      const failedRecord = getSessionRecord();
+
+      await store.storeSession(id, testRecord);
+      assert.equal(await store.loadSession(id), testRecord);
+
+      await assert.isRejected(
+        store.sessionTransaction('test', async () => {
+          await store.storeSession(id, failedRecord);
+          assert.equal(await store.loadSession(id), failedRecord);
+
+          await store.addUnprocessed({
+            id: '2-two',
+            envelope: 'second',
+            timestamp: 2,
+            version: 2,
+            attempts: 0,
+          });
+
+          throw new Error('Failure');
+        }),
+        'Failure'
+      );
+
+      assert.equal(await store.loadSession(id), testRecord);
+      assert.deepEqual(await store.getAllUnprocessed(), []);
+    });
+
+    it('can be re-entered', async () => {
+      const id = `${number}.1`;
+      const testRecord = getSessionRecord();
+
+      const lock = new Lock();
+
+      await store.sessionTransaction(
+        'test',
+        async () => {
+          await store.sessionTransaction(
+            'nested',
+            async () => {
+              await store.storeSession(id, testRecord, { lock });
+
+              assert.equal(await store.loadSession(id, { lock }), testRecord);
+            },
+            lock
+          );
+
+          assert.equal(await store.loadSession(id, { lock }), testRecord);
+        },
+        lock
+      );
+
+      assert.equal(await store.loadSession(id), testRecord);
+    });
+  });
+
   describe('Not yet processed messages', () => {
     beforeEach(async () => {
       await store.removeAllUnprocessed();

ts/textsecure.d.ts

@@ -61,15 +61,11 @@ export type TextSecureType = {
       setUuidAndDeviceId: (uuid: string, deviceId: number) => Promise<void>;
     };
     unprocessed: {
-      batchAdd: (dataArray: Array<UnprocessedType>) => Promise<void>;
       remove: (id: string | Array<string>) => Promise<void>;
       getCount: () => Promise<number>;
       removeAll: () => Promise<void>;
       getAll: () => Promise<Array<UnprocessedType>>;
       updateAttempts: (id: string, attempts: number) => Promise<void>;
-      addDecryptedDataToList: (
-        array: Array<Partial<UnprocessedType>>
-      ) => Promise<void>;
     };
     get: (key: string, defaultValue?: any) => any;
     put: (key: string, value: any) => Promise<void>;

ts/textsecure/MessageReceiver.ts

@@ -8,6 +8,7 @@
 /* eslint-disable camelcase */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable max-classes-per-file */
+/* eslint-disable no-restricted-syntax */
 
 import { isNumber, map, omit, noop } from 'lodash';
 import PQueue from 'p-queue';
@@ -121,9 +122,10 @@ type CacheAddItemType = {
   request: IncomingWebSocketRequest;
 };
 
-type CacheUpdateItemType = {
+type DecryptedEnvelope = {
-  id: string;
+  readonly plaintext: ArrayBuffer;
-  data: Partial<UnprocessedType>;
+  readonly data: UnprocessedType;
+  readonly envelope: EnvelopeClass;
 };
 
 class MessageReceiverInner extends EventTarget {
@@ -135,8 +137,6 @@ class MessageReceiverInner extends EventTarget {
 
   cacheRemoveBatcher: BatcherType<string>;
 
-  cacheUpdateBatcher: BatcherType<CacheUpdateItemType>;
-
   calledClose?: boolean;
 
   count: number;
@@ -233,12 +233,6 @@ class MessageReceiverInner extends EventTarget {
         this.cacheAndQueueBatch(items);
       },
     });
-    this.cacheUpdateBatcher = createBatcher<CacheUpdateItemType>({
-      name: 'MessageReceiver.cacheUpdateBatcher',
-      wait: 75,
-      maxSize: 30,
-      processBatch: this.cacheUpdateBatch.bind(this),
-    });
     this.cacheRemoveBatcher = createBatcher<string>({
       name: 'MessageReceiver.cacheRemoveBatcher',
       wait: 75,
@@ -314,7 +308,6 @@ class MessageReceiverInner extends EventTarget {
   unregisterBatchers() {
     window.log.info('MessageReceiver: unregister batchers');
     this.cacheAddBatcher.unregister();
-    this.cacheUpdateBatcher.unregister();
     this.cacheRemoveBatcher.unregister();
   }
 
@@ -514,7 +507,7 @@ class MessageReceiverInner extends EventTarget {
     return messageAgeSec;
   }
 
-  async addToQueue(task: () => Promise<void>) {
+  async addToQueue<T>(task: () => Promise<T>): Promise<T> {
     this.count += 1;
 
     const promise = this.pendingQueue.add(task);
@@ -538,7 +531,6 @@ class MessageReceiverInner extends EventTarget {
     const emitEmpty = async () => {
       await Promise.all([
         this.cacheAddBatcher.flushAndWait(),
-        this.cacheUpdateBatcher.flushAndWait(),
         this.cacheRemoveBatcher.flushAndWait(),
       ]);
 
@@ -655,7 +647,7 @@ class MessageReceiverInner extends EventTarget {
         }
         this.queueDecryptedEnvelope(envelope, payloadPlaintext);
       } else {
-        this.queueEnvelope(envelope);
+        this.queueEnvelope(new Sessions(), envelope);
       }
     } catch (error) {
       window.log.error(
@@ -755,21 +747,88 @@ class MessageReceiverInner extends EventTarget {
 
   async cacheAndQueueBatch(items: Array<CacheAddItemType>) {
     window.log.info('MessageReceiver.cacheAndQueueBatch', items.length);
-    const dataArray = items.map(item => item.data);
+
+    const decrypted: Array<DecryptedEnvelope> = [];
+
     try {
-      await window.textsecure.storage.unprocessed.batchAdd(dataArray);
+      const sessionStore = new Sessions({
-      items.forEach(item => {
+        transactionOnly: true,
+      });
+      const failed: Array<UnprocessedType> = [];
+
+      // Below we:
+      //
+      // 1. Enter session transaction
+      // 2. Decrypt all batched envelopes
+      // 3. Persist both decrypted envelopes and envelopes that we failed to
+      //    decrypt (for future retries, see `attempts` field)
+      // 4. Leave session transaction and commit all pending session updates
+      // 5. Acknowledge envelopes (can't fail)
+      // 6. Finally process decrypted envelopes
+      await sessionStore.transaction(async () => {
+        await Promise.all<void>(
+          items.map(async ({ data, envelope }) => {
+            try {
+              const plaintext = await this.queueEnvelope(
+                sessionStore,
+                envelope
+              );
+              if (plaintext) {
+                decrypted.push({ plaintext, data, envelope });
+              }
+            } catch (error) {
+              failed.push(data);
+              window.log.error(
+                'cacheAndQueue error when processing the envelope',
+                error && error.stack ? error.stack : error
+              );
+            }
+          })
+        );
+
+        window.log.info(
+          'MessageReceiver.cacheAndQueueBatch storing ' +
+            `${decrypted.length} decrypted envelopes`
+        );
+
+        // Store both decrypted and failed unprocessed envelopes
+        const unprocesseds: Array<UnprocessedType> = decrypted.map(
+          ({ envelope, data, plaintext }) => {
+            return {
+              ...data,
+
+              // We have sucessfully decrypted the message so don't bother with
+              // storing the envelope.
+              envelope: '',
+
+              source: envelope.source,
+              sourceUuid: envelope.sourceUuid,
+              sourceDevice: envelope.sourceDevice,
+              serverTimestamp: envelope.serverTimestamp,
+              decrypted: MessageReceiverInner.arrayBufferToStringBase64(
+                plaintext
+              ),
+            };
+          }
+        );
+
+        await sessionStore.addUnprocessed(unprocesseds.concat(failed));
+      });
+
+      window.log.info(
+        'MessageReceiver.cacheAndQueueBatch acknowledging receipt'
+      );
+
+      // Acknowledge all envelopes
+      for (const { request } of items) {
         try {
-          item.request.respond(200, 'OK');
+          request.respond(200, 'OK');
         } catch (error) {
           window.log.error(
             'cacheAndQueueBatch: Failed to send 200 to server; still queuing envelope'
           );
         }
-        this.queueEnvelope(item.envelope);
+      }
-      });
-
-      this.maybeScheduleRetryTimeout();
     } catch (error) {
       window.log.error(
         'cacheAndQueue error trying to add messages to cache:',
@@ -779,7 +838,25 @@ class MessageReceiverInner extends EventTarget {
       items.forEach(item => {
         item.request.respond(500, 'Failed to cache message');
       });
+      return;
     }
+
+    await Promise.all(
+      decrypted.map(async ({ envelope, plaintext }) => {
+        try {
+          await this.queueDecryptedEnvelope(envelope, plaintext);
+        } catch (error) {
+          window.log.error(
+            'cacheAndQueue error when processing decrypted envelope',
+            error && error.stack ? error.stack : error
+          );
+        }
+      })
+    );
+
+    window.log.info('MessageReceiver.cacheAndQueueBatch fully processed');
+
+    this.maybeScheduleRetryTimeout();
   }
 
   cacheAndQueue(
@@ -802,23 +879,6 @@ class MessageReceiverInner extends EventTarget {
     });
   }
 
-  async cacheUpdateBatch(items: Array<Partial<UnprocessedType>>) {
-    window.log.info('MessageReceiver.cacheUpdateBatch', items.length);
-    await window.textsecure.storage.unprocessed.addDecryptedDataToList(items);
-  }
-
-  updateCache(envelope: EnvelopeClass, plaintext: ArrayBuffer) {
-    const { id } = envelope;
-    const data = {
-      source: envelope.source,
-      sourceUuid: envelope.sourceUuid,
-      sourceDevice: envelope.sourceDevice,
-      serverTimestamp: envelope.serverTimestamp,
-      decrypted: MessageReceiverInner.arrayBufferToStringBase64(plaintext),
-    };
-    this.cacheUpdateBatcher.add({ id, data });
-  }
-
   async cacheRemoveBatch(items: Array<string>) {
     await window.textsecure.storage.unprocessed.remove(items);
   }
@@ -851,18 +911,22 @@ class MessageReceiverInner extends EventTarget {
     });
   }
 
-  async queueEnvelope(envelope: EnvelopeClass) {
+  async queueEnvelope(
+    sessionStore: Sessions,
+    envelope: EnvelopeClass
+  ): Promise<ArrayBuffer | undefined> {
     const id = this.getEnvelopeId(envelope);
     window.log.info('queueing envelope', id);
 
-    const task = this.handleEnvelope.bind(this, envelope);
+    const task = this.decryptEnvelope.bind(this, sessionStore, envelope);
     const taskWithTimeout = window.textsecure.createTaskWithTimeout(
       task,
       `queueEnvelope ${id}`
     );
-    const promise = this.addToQueue(taskWithTimeout);
 
-    return promise.catch(error => {
+    try {
+      return await this.addToQueue(taskWithTimeout);
+    } catch (error) {
       const args = [
         'queueEnvelope error handling envelope',
         this.getEnvelopeId(envelope),
@@ -875,12 +939,11 @@ class MessageReceiverInner extends EventTarget {
       } else {
         window.log.error(...args);
       }
-    });
+      return undefined;
+    }
   }
 
-  // Same as handleEnvelope, just without the decryption step. Necessary for handling
+  // Called after `decryptEnvelope` decrypted the message.
-  //   messages which were successfully decrypted, but application logic didn't finish
-  //   processing.
   async handleDecryptedEnvelope(
     envelope: EnvelopeClass,
     plaintext: ArrayBuffer
@@ -906,21 +969,26 @@ class MessageReceiverInner extends EventTarget {
     throw new Error('Received message with no content and no legacyMessage');
   }
 
-  async handleEnvelope(envelope: EnvelopeClass) {
+  async decryptEnvelope(
+    sessionStore: Sessions,
+    envelope: EnvelopeClass
+  ): Promise<ArrayBuffer | undefined> {
     if (this.stoppingProcessing) {
-      return Promise.resolve();
+      return undefined;
     }
 
     if (envelope.type === window.textsecure.protobuf.Envelope.Type.RECEIPT) {
-      return this.onDeliveryReceipt(envelope);
+      await this.onDeliveryReceipt(envelope);
+      return undefined;
     }
 
     if (envelope.content) {
-      return this.handleContentMessage(envelope);
+      return this.decryptContentMessage(sessionStore, envelope);
     }
     if (envelope.legacyMessage) {
-      return this.handleLegacyMessage(envelope);
+      return this.decryptLegacyMessage(sessionStore, envelope);
     }
+
     this.removeFromCache(envelope);
     throw new Error('Received message with no content and no legacyMessage');
   }
@@ -935,7 +1003,7 @@ class MessageReceiverInner extends EventTarget {
     return -1;
   }
 
-  async onDeliveryReceipt(envelope: EnvelopeClass) {
+  async onDeliveryReceipt(envelope: EnvelopeClass): Promise<void> {
     return new Promise((resolve, reject) => {
       const ev = new Event('delivery');
       ev.confirm = this.removeFromCache.bind(this, envelope);
@@ -968,6 +1036,7 @@ class MessageReceiverInner extends EventTarget {
   }
 
   async decrypt(
+    sessionStore: Sessions,
     envelope: EnvelopeClass,
     ciphertext: ByteBufferClass
   ): Promise<ArrayBuffer | null> {
@@ -990,7 +1059,6 @@ class MessageReceiverInner extends EventTarget {
       throw new Error('MessageReceiver.decrypt: Failed to fetch local UUID');
     }
 
-    const sessionStore = new Sessions();
     const identityKeyStore = new IdentityKeys();
     const preKeyStore = new PreKeys();
     const signedPreKeyStore = new SignedPreKeys();
@@ -1216,15 +1284,6 @@ class MessageReceiverInner extends EventTarget {
             return null;
           }
 
-          // Note: this is an out of band update; there are cases where the item in the
-          //   cache has already been deleted by the time this runs. That's okay.
-          try {
-            this.updateCache(envelope, plaintext);
-          } catch (error) {
-            const errorString = error && error.stack ? error.stack : error;
-            window.log.error(`decrypt: updateCache failed: ${errorString}`);
-          }
-
           return plaintext;
         }
       )
@@ -1540,18 +1599,25 @@ class MessageReceiverInner extends EventTarget {
     );
   }
 
-  async handleLegacyMessage(envelope: EnvelopeClass) {
+  async decryptLegacyMessage(
+    sessionStore: Sessions,
+    envelope: EnvelopeClass
+  ): Promise<ArrayBuffer | undefined> {
     window.log.info(
-      'MessageReceiver.handleLegacyMessage',
+      'MessageReceiver.decryptLegacyMessage',
       this.getEnvelopeId(envelope)
     );
-    return this.decrypt(envelope, envelope.legacyMessage).then(plaintext => {
+    const plaintext = await this.decrypt(
-      if (!plaintext) {
+      sessionStore,
-        window.log.warn('handleLegacyMessage: plaintext was falsey');
+      envelope,
-        return null;
+      envelope.legacyMessage
-      }
+    );
-      return this.innerHandleLegacyMessage(envelope, plaintext);
+    if (!plaintext) {
-    });
+      window.log.warn('decryptLegacyMessage: plaintext was falsey');
+      return undefined;
+    }
+
+    return plaintext;
   }
 
   async innerHandleLegacyMessage(
@@ -1562,18 +1628,25 @@ class MessageReceiverInner extends EventTarget {
     return this.handleDataMessage(envelope, message);
   }
 
-  async handleContentMessage(envelope: EnvelopeClass) {
+  async decryptContentMessage(
+    sessionStore: Sessions,
+    envelope: EnvelopeClass
+  ): Promise<ArrayBuffer | undefined> {
     window.log.info(
-      'MessageReceiver.handleContentMessage',
+      'MessageReceiver.decryptContentMessage',
       this.getEnvelopeId(envelope)
     );
-    return this.decrypt(envelope, envelope.content).then(plaintext => {
+    const plaintext = await this.decrypt(
-      if (!plaintext) {
+      sessionStore,
-        window.log.warn('handleContentMessage: plaintext was falsey');
+      envelope,
-        return null;
+      envelope.content
-      }
+    );
-      return this.innerHandleContentMessage(envelope, plaintext);
+    if (!plaintext) {
-    });
+      window.log.warn('decryptContentMessage: plaintext was falsey');
+      return undefined;
+    }
+
+    return plaintext;
   }
 
   async innerHandleContentMessage(

ts/util/Lock.ts

@@ -0,0 +1,4 @@
+// Copyright 2021 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+export class Lock {}

yarn.lock

@@ -2200,10 +2200,17 @@
     "@types/connect" "*"
     "@types/node" "*"
 
-"@types/chai@4.1.2":
+"@types/chai-as-promised@7.1.4":
-  version "4.1.2"
+  version "7.1.4"
-  resolved "https://registry.yarnpkg.com/@types/chai/-/chai-4.1.2.tgz#f1af664769cfb50af805431c407425ed619daa21"
+  resolved "https://registry.yarnpkg.com/@types/chai-as-promised/-/chai-as-promised-7.1.4.tgz#caf64e76fb056b8c8ced4b761ed499272b737601"
-  integrity sha512-D8uQwKYUw2KESkorZ27ykzXgvkDJYXVEihGklgfp5I4HUP8D6IxtcdLTMB1emjQiWzV7WZ5ihm1cxIzVwjoleQ==
+  integrity sha512-1y3L1cHePcIm5vXkh1DSGf/zQq5n5xDKG1fpCvf18+uOkpce0Z1ozNFPkyWsVswK7ntN1sZBw3oU6gmN+pDUcA==
+  dependencies:
+    "@types/chai" "*"
+
+"@types/chai@*", "@types/chai@4.2.18":
+  version "4.2.18"
+  resolved "https://registry.yarnpkg.com/@types/chai/-/chai-4.2.18.tgz#0c8e298dbff8205e2266606c1ea5fbdba29b46e4"
+  integrity sha512-rS27+EkB/RE1Iz3u0XtVL5q36MGDWbgYe7zWiodyKNUnthxY0rukK5V36eiUCtCisB7NN8zKYH6DO2M37qxFEQ==
 
 "@types/classnames@2.2.3":
   version "2.2.3"
@@ -3869,9 +3876,10 @@ assert@^1.1.1:
   dependencies:
     util "0.10.3"
 
-assertion-error@^1.0.1:
+assertion-error@^1.1.0:
-  version "1.0.2"
+  version "1.1.0"
-  resolved "https://registry.yarnpkg.com/assertion-error/-/assertion-error-1.0.2.tgz#13ca515d86206da0bac66e834dd397d87581094c"
+  resolved "https://registry.yarnpkg.com/assertion-error/-/assertion-error-1.1.0.tgz#e60b6b0e8f301bd97e5375215bda406c85118c0b"
+  integrity sha512-jgsaNduz+ndvGyFt3uSuWqvy4lCnIJiovtouQN5JZHOKCS2QuhEdbcQHFhVksz2N2U9hXJo8odG7ETyWlEeuDw==
 
 assign-symbols@^1.0.0:
   version "1.0.0"
@@ -5087,17 +5095,24 @@ catharsis@^0.8.10:
   dependencies:
     lodash "^4.17.14"
 
-chai@4.1.2:
+chai-as-promised@7.1.1:
-  version "4.1.2"
+  version "7.1.1"
-  resolved "https://registry.yarnpkg.com/chai/-/chai-4.1.2.tgz#0f64584ba642f0f2ace2806279f4f06ca23ad73c"
+  resolved "https://registry.yarnpkg.com/chai-as-promised/-/chai-as-promised-7.1.1.tgz#08645d825deb8696ee61725dbf590c012eb00ca0"
-  integrity sha1-D2RYS6ZC8PKs4oBiefTwbKI61zw=
+  integrity sha512-azL6xMoi+uxu6z4rhWQ1jbdUhOMhis2PvscD/xjLqNMkv3BPPp2JyyuTHOrf9BOosGpNQ11v6BKv/g57RXbiaA==
   dependencies:
-    assertion-error "^1.0.1"
+    check-error "^1.0.2"
-    check-error "^1.0.1"
+
-    deep-eql "^3.0.0"
+chai@4.3.4:
+  version "4.3.4"
+  resolved "https://registry.yarnpkg.com/chai/-/chai-4.3.4.tgz#b55e655b31e1eac7099be4c08c21964fce2e6c49"
+  integrity sha512-yS5H68VYOCtN1cjfwumDSuzn/9c+yza4f3reKXlE5rUg7SFcCEy90gJvydNgOYtblyf4Zi6jIWRnXOgErta0KA==
+  dependencies:
+    assertion-error "^1.1.0"
+    check-error "^1.0.2"
+    deep-eql "^3.0.1"
     get-func-name "^2.0.0"
-    pathval "^1.0.0"
+    pathval "^1.1.1"
-    type-detect "^4.0.0"
+    type-detect "^4.0.5"
 
 chainsaw@~0.1.0:
   version "0.1.0"
@@ -5168,9 +5183,10 @@ chardet@^0.7.0:
   resolved "https://registry.yarnpkg.com/chardet/-/chardet-0.7.0.tgz#90094849f0937f2eedc2425d0d28a9e5f0cbad9e"
   integrity sha512-mT8iDcrh03qDGRRmoA2hmBJnxpllMR+0/0qlzjqZES6NdiWDcZkCNAk4rPFZ9Q85r27unkiNNg8ZOiwZXBHwcA==
 
-check-error@^1.0.1:
+check-error@^1.0.2:
   version "1.0.2"
   resolved "https://registry.yarnpkg.com/check-error/-/check-error-1.0.2.tgz#574d312edd88bb5dd8912e9286dd6c0aed4aac82"
+  integrity sha1-V00xLt2Iu13YkS6Sht1sCu1KrII=
 
 chokidar@^2.0.2:
   version "2.0.4"
@@ -6342,9 +6358,10 @@ deep-diff@^0.3.5:
   resolved "https://registry.yarnpkg.com/deep-diff/-/deep-diff-0.3.8.tgz#c01de63efb0eec9798801d40c7e0dae25b582c84"
   integrity sha1-wB3mPvsO7JeYgB1Ax+Da4ltYLIQ=
 
-deep-eql@^3.0.0:
+deep-eql@^3.0.1:
   version "3.0.1"
   resolved "https://registry.yarnpkg.com/deep-eql/-/deep-eql-3.0.1.tgz#dfc9404400ad1c8fe023e7da1df1c147c4b444df"
+  integrity sha512-+QeIQyN5ZuO+3Uk5DYh6/1eKO0m0YmJFGNmFHGACpf1ClL1nmlV/p4gNgbl2pJGxgXb4faqo6UE+M5ACEMyVcw==
   dependencies:
     type-detect "^4.0.0"
 
@@ -13498,9 +13515,10 @@ path@^0.12.7:
     process "^0.11.1"
     util "^0.10.3"
 
-pathval@^1.0.0:
+pathval@^1.1.1:
-  version "1.1.0"
+  version "1.1.1"
-  resolved "https://registry.yarnpkg.com/pathval/-/pathval-1.1.0.tgz#b942e6d4bde653005ef6b71361def8727d0645e0"
+  resolved "https://registry.yarnpkg.com/pathval/-/pathval-1.1.1.tgz#8534e77a77ce7ac5a2512ea21e0fdb8fcf6c3d8d"
+  integrity sha512-Dp6zGqpTdETdR63lehJYPeIOqpiNBNtc7BpWSLrOje7UaIsE5aY92r/AunQA7rsXvet3lrJ3JnZX29UPTKXyKQ==
 
 pbkdf2@^3.0.3:
   version "3.0.14"
@@ -17608,7 +17626,7 @@ type-check@~0.3.2:
   dependencies:
     prelude-ls "~1.1.2"
 
-type-detect@4.0.8, type-detect@^4.0.8:
+type-detect@4.0.8, type-detect@^4.0.5, type-detect@^4.0.8:
   version "4.0.8"
   resolved "https://registry.yarnpkg.com/type-detect/-/type-detect-4.0.8.tgz#7646fb5f18871cfbb7749e69bd39a6388eb7450c"
   integrity sha512-0fr/mIH1dlO+x7TlcMy+bIDqKPsw/70tVyeHW787goQjhmqaZe10uwLujubK9q9Lg6Fiho1KUKDYz0Z7k7g5/g==