Include sender keys in SignalProtocolStore zones

ts/LibSignalStores.ts

@@ -214,15 +214,19 @@ export class PreKeys extends PreKeyStore {
 }
 
 export type SenderKeysOptions = Readonly<{
-  ourUuid: UUID;
+  readonly ourUuid: UUID;
+  readonly zone: Zone | undefined;
 }>;
 
 export class SenderKeys extends SenderKeyStore {
   private readonly ourUuid: UUID;
 
-  constructor({ ourUuid }: SenderKeysOptions) {
+  readonly zone: Zone | undefined;
+
+  constructor({ ourUuid, zone }: SenderKeysOptions) {
     super();
     this.ourUuid = ourUuid;
+    this.zone = zone;
   }
 
   async saveSenderKey(
@@ -235,7 +239,8 @@ export class SenderKeys extends SenderKeyStore {
     await window.textsecure.storage.protocol.saveSenderKey(
       encodedAddress,
       distributionId,
-      record
+      record,
+      { zone: this.zone }
     );
   }
 
@@ -247,7 +252,8 @@ export class SenderKeys extends SenderKeyStore {
 
     const senderKey = await window.textsecure.storage.protocol.getSenderKey(
       encodedAddress,
-      distributionId
+      distributionId,
+      { zone: this.zone }
     );
 
     return senderKey || null;

ts/SignalProtocolStore.ts

@@ -191,6 +191,7 @@ const EventsMixin = function EventsMixin(this: unknown) {
 } as any as typeof window.Backbone.EventsMixin;
 
 type SessionCacheEntry = CacheEntryType<SessionType, SessionRecord>;
+type SenderKeyCacheEntry = CacheEntryType<SenderKeyType, SenderKeyRecord>;
 
 type ZoneQueueEntryType = Readonly<{
   zone: Zone;
@@ -213,10 +214,7 @@ export class SignalProtocolStore extends EventsMixin {
     CacheEntryType<IdentityKeyType, PublicKey>
   >;
 
-  senderKeys?: Map<
+  senderKeys?: Map<SenderKeyIdType, SenderKeyCacheEntry>;
-    SenderKeyIdType,
-    CacheEntryType<SenderKeyType, SenderKeyRecord>
-  >;
 
   sessions?: Map<SessionIdType, SessionCacheEntry>;
 
@@ -239,6 +237,8 @@ export class SignalProtocolStore extends EventsMixin {
 
   private pendingSessions = new Map<SessionIdType, SessionCacheEntry>();
 
+  private pendingSenderKeys = new Map<SenderKeyIdType, SenderKeyCacheEntry>();
+
   private pendingUnprocessed = new Map<string, UnprocessedType>();
 
   async hydrateCaches(): Promise<void> {
@@ -501,7 +501,21 @@ export class SignalProtocolStore extends EventsMixin {
     await window.Signal.Data.removeAllSignedPreKeys();
   }
 
-  // Sender Key Queue
+  // Sender Key
+
+  // Re-entrant sender key transaction routine. Only one sender key transaction could
+  // be running at the same time.
+  //
+  // While in transaction:
+  //
+  // - `saveSenderKey()` adds the updated session to the `pendingSenderKeys`
+  // - `getSenderKey()` looks up the session first in `pendingSenderKeys` and only
+  //   then in the main `senderKeys` store
+  //
+  // When transaction ends:
+  //
+  // - successfully: pending sender key stores are batched into the database
+  // - with an error: pending sender key stores are reverted
 
   async enqueueSenderKeyJob<T>(
     qualifiedAddress: QualifiedAddress,
@@ -534,8 +548,6 @@ export class SignalProtocolStore extends EventsMixin {
     return freshQueue;
   }
 
-  // Sender Keys
-
   private getSenderKeyId(
     senderKeyId: QualifiedAddress,
     distributionId: string
@@ -546,8 +558,10 @@ export class SignalProtocolStore extends EventsMixin {
   async saveSenderKey(
     qualifiedAddress: QualifiedAddress,
     distributionId: string,
-    record: SenderKeyRecord
+    record: SenderKeyRecord,
+    { zone = GLOBAL_ZONE }: SessionTransactionOptions = {}
   ): Promise<void> {
+    await this.withZone(zone, 'saveSenderKey', async () => {
       if (!this.senderKeys) {
         throw new Error('saveSenderKey: this.senderKeys not yet cached!');
       }
@@ -565,25 +579,31 @@ export class SignalProtocolStore extends EventsMixin {
           lastUpdatedDate: Date.now(),
         };
 
-      await window.Signal.Data.createOrUpdateSenderKey(fromDB);
+        this.pendingSenderKeys.set(id, {
-
-      this.senderKeys.set(id, {
           hydrated: true,
           fromDB,
           item: record,
         });
+
+        // Current zone doesn't support pending sessions - commit immediately
+        if (!zone.supportsPendingSenderKeys()) {
+          await this.commitZoneChanges('saveSenderKey');
+        }
       } catch (error) {
         const errorString = error && error.stack ? error.stack : error;
         log.error(
           `saveSenderKey: failed to save senderKey ${senderId}/${distributionId}: ${errorString}`
         );
       }
+    });
   }
 
   async getSenderKey(
     qualifiedAddress: QualifiedAddress,
-    distributionId: string
+    distributionId: string,
+    { zone = GLOBAL_ZONE }: SessionTransactionOptions = {}
   ): Promise<SenderKeyRecord | undefined> {
+    return this.withZone(zone, 'getSenderKey', async () => {
       if (!this.senderKeys) {
         throw new Error('getSenderKey: this.senderKeys not yet cached!');
       }
@@ -593,7 +613,11 @@ export class SignalProtocolStore extends EventsMixin {
       try {
         const id = this.getSenderKeyId(qualifiedAddress, distributionId);
 
-      const entry = this.senderKeys.get(id);
+        const map = this.pendingSenderKeys.has(id)
+          ? this.pendingSenderKeys
+          : this.senderKeys;
+        const entry = map.get(id);
+
         if (!entry) {
           log.error('Failed to fetch sender key:', id);
           return undefined;
@@ -604,7 +628,9 @@ export class SignalProtocolStore extends EventsMixin {
           return entry.item;
         }
 
-      const item = SenderKeyRecord.deserialize(Buffer.from(entry.fromDB.data));
+        const item = SenderKeyRecord.deserialize(
+          Buffer.from(entry.fromDB.data)
+        );
         this.senderKeys.set(id, {
           hydrated: true,
           item,
@@ -619,6 +645,7 @@ export class SignalProtocolStore extends EventsMixin {
         );
         return undefined;
       }
+    });
   }
 
   async removeSenderKey(
@@ -645,11 +672,16 @@ export class SignalProtocolStore extends EventsMixin {
     }
   }
 
-  async clearSenderKeyStore(): Promise<void> {
+  async removeAllSenderKeys(): Promise<void> {
+    return this.withZone(GLOBAL_ZONE, 'removeAllSenderKeys', async () => {
       if (this.senderKeys) {
         this.senderKeys.clear();
       }
+      if (this.pendingSenderKeys) {
+        this.pendingSenderKeys.clear();
+      }
       await window.Signal.Data.removeAllSenderKeys();
+    });
   }
 
   // Session Queue
@@ -700,6 +732,7 @@ export class SignalProtocolStore extends EventsMixin {
   //
   // - successfully: pending session stores are batched into the database
   // - with an error: pending session stores are reverted
+
   public async withZone<T>(
     zone: Zone,
     name: string,
@@ -753,45 +786,66 @@ export class SignalProtocolStore extends EventsMixin {
   }
 
   private async commitZoneChanges(name: string): Promise<void> {
-    const { pendingSessions, pendingUnprocessed } = this;
+    const { pendingSenderKeys, pendingSessions, pendingUnprocessed } = this;
 
-    if (pendingSessions.size === 0 && pendingUnprocessed.size === 0) {
+    if (
+      pendingSenderKeys.size === 0 &&
+      pendingSessions.size === 0 &&
+      pendingUnprocessed.size === 0
+    ) {
       return;
     }
 
     log.info(
-      `commitZoneChanges(${name}): pending sessions ${pendingSessions.size} ` +
+      `commitZoneChanges(${name}): ` +
+        `pending sender keys ${pendingSenderKeys.size}, ` +
+        `pending sessions ${pendingSessions.size}, ` +
         `pending unprocessed ${pendingUnprocessed.size}`
     );
 
+    this.pendingSenderKeys = new Map();
     this.pendingSessions = new Map();
     this.pendingUnprocessed = new Map();
 
-    // Commit both unprocessed and sessions in the same database transaction
+    // Commit both sender keys, sessions and unprocessed in the same database transaction
     //   to unroll both on error.
-    await window.Signal.Data.commitSessionsAndUnprocessed({
+    await window.Signal.Data.commitDecryptResult({
+      senderKeys: Array.from(pendingSenderKeys.values()).map(
+        ({ fromDB }) => fromDB
+      ),
       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.
+
+    const { sessions } = this;
+    assert(sessions !== undefined, "Can't commit unhydrated session storage");
     pendingSessions.forEach((value, key) => {
       sessions.set(key, value);
     });
+
+    const { senderKeys } = this;
+    assert(
+      senderKeys !== undefined,
+      "Can't commit unhydrated sender key storage"
+    );
+    pendingSenderKeys.forEach((value, key) => {
+      senderKeys.set(key, value);
+    });
   }
 
   private async revertZoneChanges(name: string, error: Error): Promise<void> {
     log.info(
       `revertZoneChanges(${name}): ` +
-        `pending sessions size ${this.pendingSessions.size} ` +
+        `pending sender keys size ${this.pendingSenderKeys.size}, ` +
+        `pending sessions size ${this.pendingSessions.size}, ` +
         `pending unprocessed size ${this.pendingUnprocessed.size}`,
       error && error.stack
     );
+    this.pendingSenderKeys.clear();
     this.pendingSessions.clear();
     this.pendingUnprocessed.clear();
   }

ts/sql/Client.ts

@@ -186,7 +186,7 @@ const dataInterface: ClientInterface = {
 
   createOrUpdateSession,
   createOrUpdateSessions,
-  commitSessionsAndUnprocessed,
+  commitDecryptResult,
   bulkAddSessions,
   removeSessionById,
   removeSessionsByConversation,
@@ -921,11 +921,12 @@ async function createOrUpdateSession(data: SessionType) {
 async function createOrUpdateSessions(array: Array<SessionType>) {
   await channels.createOrUpdateSessions(array);
 }
-async function commitSessionsAndUnprocessed(options: {
+async function commitDecryptResult(options: {
+  senderKeys: Array<SenderKeyType>;
   sessions: Array<SessionType>;
   unprocessed: Array<UnprocessedType>;
 }) {
-  await channels.commitSessionsAndUnprocessed(options);
+  await channels.commitDecryptResult(options);
 }
 async function bulkAddSessions(array: Array<SessionType>) {
   await channels.bulkAddSessions(array);

ts/sql/Interface.ts

@@ -329,7 +329,8 @@ export type DataInterface = {
 
   createOrUpdateSession: (data: SessionType) => Promise<void>;
   createOrUpdateSessions: (array: Array<SessionType>) => Promise<void>;
-  commitSessionsAndUnprocessed(options: {
+  commitDecryptResult(options: {
+    senderKeys: Array<SenderKeyType>;
     sessions: Array<SessionType>;
     unprocessed: Array<UnprocessedType>;
   }): Promise<void>;

ts/sql/Server.ts

@@ -182,7 +182,7 @@ const dataInterface: ServerInterface = {
 
   createOrUpdateSession,
   createOrUpdateSessions,
-  commitSessionsAndUnprocessed,
+  commitDecryptResult,
   bulkAddSessions,
   removeSessionById,
   removeSessionsByConversation,
@@ -757,6 +757,10 @@ async function removeAllItems(): Promise<void> {
 }
 
 async function createOrUpdateSenderKey(key: SenderKeyType): Promise<void> {
+  createOrUpdateSenderKeySync(key);
+}
+
+function createOrUpdateSenderKeySync(key: SenderKeyType): void {
   const db = getInstance();
 
   prepare(
@@ -1175,16 +1179,22 @@ async function createOrUpdateSessions(
   })();
 }
 
-async function commitSessionsAndUnprocessed({
+async function commitDecryptResult({
+  senderKeys,
   sessions,
   unprocessed,
 }: {
+  senderKeys: Array<SenderKeyType>;
   sessions: Array<SessionType>;
   unprocessed: Array<UnprocessedType>;
 }): Promise<void> {
   const db = getInstance();
 
   db.transaction(() => {
+    for (const item of senderKeys) {
+      assertSync(createOrUpdateSenderKeySync(item));
+    }
+
     for (const item of sessions) {
       assertSync(createOrUpdateSessionSync(item));
     }

ts/test-electron/SignalProtocolStore_test.ts

@@ -1442,7 +1442,9 @@ describe('SignalProtocolStore', () => {
   });
 
   describe('zones', () => {
+    const distributionId = UUID.generate().toString();
     const zone = new Zone('zone', {
+      pendingSenderKeys: true,
       pendingSessions: true,
       pendingUnprocessed: true,
     });
@@ -1450,6 +1452,7 @@ describe('SignalProtocolStore', () => {
     beforeEach(async () => {
       await store.removeAllUnprocessed();
       await store.removeAllSessions(theirUuid.toString());
+      await store.removeAllSenderKeys();
     });
 
     it('should not store pending sessions in global zone', async () => {
@@ -1467,12 +1470,29 @@ describe('SignalProtocolStore', () => {
       assert.equal(await store.loadSession(id), testRecord);
     });
 
-    it('commits session stores and unprocessed on success', async () => {
+    it('should not store pending sender keys in global zone', async () => {
       const id = new QualifiedAddress(ourUuid, new Address(theirUuid, 1));
-      const testRecord = getSessionRecord();
+      const testRecord = getSenderKeyRecord();
+
+      await assert.isRejected(
+        store.withZone(GLOBAL_ZONE, 'test', async () => {
+          await store.saveSenderKey(id, distributionId, testRecord);
+          throw new Error('Failure');
+        }),
+        'Failure'
+      );
+
+      assert.equal(await store.getSenderKey(id, distributionId), testRecord);
+    });
+
+    it('commits sender keys, sessions and unprocessed on success', async () => {
+      const id = new QualifiedAddress(ourUuid, new Address(theirUuid, 1));
+      const testSession = getSessionRecord();
+      const testSenderKey = getSenderKeyRecord();
 
       await store.withZone(zone, 'test', async () => {
-        await store.storeSession(id, testRecord, { zone });
+        await store.storeSession(id, testSession, { zone });
+        await store.saveSenderKey(id, distributionId, testSenderKey, { zone });
 
         await store.addUnprocessed(
           {
@@ -1484,10 +1504,16 @@ describe('SignalProtocolStore', () => {
           },
           { zone }
         );
-        assert.equal(await store.loadSession(id, { zone }), testRecord);
+
+        assert.equal(await store.loadSession(id, { zone }), testSession);
+        assert.equal(
+          await store.getSenderKey(id, distributionId, { zone }),
+          testSenderKey
+        );
       });
 
-      assert.equal(await store.loadSession(id), testRecord);
+      assert.equal(await store.loadSession(id), testSession);
+      assert.equal(await store.getSenderKey(id, distributionId), testSenderKey);
 
       const allUnprocessed = await store.getAllUnprocessed();
       assert.deepEqual(
@@ -1496,18 +1522,31 @@ describe('SignalProtocolStore', () => {
       );
     });
 
-    it('reverts session stores and unprocessed on error', async () => {
+    it('reverts sender keys, sessions and unprocessed on error', async () => {
       const id = new QualifiedAddress(ourUuid, new Address(theirUuid, 1));
-      const testRecord = getSessionRecord();
+      const testSession = getSessionRecord();
-      const failedRecord = getSessionRecord();
+      const failedSession = getSessionRecord();
+      const testSenderKey = getSenderKeyRecord();
+      const failedSenderKey = getSenderKeyRecord();
 
-      await store.storeSession(id, testRecord);
+      await store.storeSession(id, testSession);
-      assert.equal(await store.loadSession(id), testRecord);
+      assert.equal(await store.loadSession(id), testSession);
+
+      await store.saveSenderKey(id, distributionId, testSenderKey);
+      assert.equal(await store.getSenderKey(id, distributionId), testSenderKey);
 
       await assert.isRejected(
         store.withZone(zone, 'test', async () => {
-          await store.storeSession(id, failedRecord, { zone });
+          await store.storeSession(id, failedSession, { zone });
-          assert.equal(await store.loadSession(id, { zone }), failedRecord);
+          assert.equal(await store.loadSession(id, { zone }), failedSession);
+
+          await store.saveSenderKey(id, distributionId, failedSenderKey, {
+            zone,
+          });
+          assert.equal(
+            await store.getSenderKey(id, distributionId, { zone }),
+            failedSenderKey
+          );
 
           await store.addUnprocessed(
             {
@@ -1525,7 +1564,8 @@ describe('SignalProtocolStore', () => {
         'Failure'
       );
 
-      assert.equal(await store.loadSession(id), testRecord);
+      assert.equal(await store.loadSession(id), testSession);
+      assert.equal(await store.getSenderKey(id, distributionId), testSenderKey);
       assert.deepEqual(await store.getAllUnprocessed(), []);
     });
 

ts/textsecure/MessageReceiver.ts

@@ -138,6 +138,7 @@ type CacheAddItemType = {
 };
 
 type LockedStores = {
+  readonly senderKeyStore: SenderKeys;
   readonly sessionStore: Sessions;
   readonly identityKeyStore: IdentityKeys;
   readonly zone?: Zone;
@@ -779,6 +780,7 @@ export default class MessageReceiver
 
     try {
       const zone = new Zone('decryptAndCacheBatch', {
+        pendingSenderKeys: true,
         pendingSessions: true,
         pendingUnprocessed: true,
       });
@@ -814,6 +816,10 @@ export default class MessageReceiver
               let stores = storesMap.get(destinationUuid.toString());
               if (!stores) {
                 stores = {
+                  senderKeyStore: new SenderKeys({
+                    ourUuid: destinationUuid,
+                    zone,
+                  }),
                   sessionStore: new Sessions({
                     zone,
                     ourUuid: destinationUuid,
@@ -1301,7 +1307,7 @@ export default class MessageReceiver
   }
 
   private async decryptSealedSender(
-    { sessionStore, identityKeyStore, zone }: LockedStores,
+    { senderKeyStore, sessionStore, identityKeyStore, zone }: LockedStores,
     envelope: UnsealedEnvelope,
     ciphertext: Uint8Array
   ): Promise<DecryptSealedSenderResult> {
@@ -1352,7 +1358,6 @@ export default class MessageReceiver
       );
       const sealedSenderIdentifier = certificate.senderUuid();
       const sealedSenderSourceDevice = certificate.senderDeviceId();
-      const senderKeyStore = new SenderKeys({ ourUuid: destinationUuid });
 
       const address = new QualifiedAddress(
         destinationUuid,
@@ -2000,7 +2005,6 @@ export default class MessageReceiver
         Buffer.from(distributionMessage)
       );
     const { destinationUuid } = envelope;
-    const senderKeyStore = new SenderKeys({ ourUuid: destinationUuid });
     const address = new QualifiedAddress(
       destinationUuid,
       Address.create(identifier, sourceDevice)
@@ -2012,7 +2016,7 @@ export default class MessageReceiver
         processSenderKeyDistributionMessage(
           sender,
           senderKeyDistributionMessage,
-          senderKeyStore
+          stores.senderKeyStore
         ),
       stores.zone
     );

ts/textsecure/SendMessage.ts

@@ -15,6 +15,7 @@ import {
   SenderKeyDistributionMessage,
 } from '@signalapp/signal-client';
 
+import { GLOBAL_ZONE } from '../SignalProtocolStore';
 import { assert } from '../util/assert';
 import { parseIntOrThrow } from '../util/parseIntOrThrow';
 import { Address } from '../types/Address';
@@ -1874,7 +1875,7 @@ export default class MessageSender {
       ourUuid,
       new Address(ourUuid, ourDeviceId)
     );
-    const senderKeyStore = new SenderKeys({ ourUuid });
+    const senderKeyStore = new SenderKeys({ ourUuid, zone: GLOBAL_ZONE });
 
     return window.textsecure.storage.protocol.enqueueSenderKeyJob(
       address,

ts/util/Zone.ts

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
 export type ZoneOptions = {
+  readonly pendingSenderKeys?: boolean;
   readonly pendingSessions?: boolean;
   readonly pendingUnprocessed?: boolean;
 };
@@ -12,6 +13,10 @@ export class Zone {
     private readonly options: ZoneOptions = {}
   ) {}
 
+  public supportsPendingSenderKeys(): boolean {
+    return this.options.pendingSenderKeys === true;
+  }
+
   public supportsPendingSessions(): boolean {
     return this.options.pendingSessions === true;
   }

ts/util/sendToGroup.ts

@@ -55,6 +55,7 @@ import * as RemoteConfig from '../RemoteConfig';
 
 import { strictAssert } from './assert';
 import * as log from '../logging/log';
+import { GLOBAL_ZONE } from '../SignalProtocolStore';
 
 const ERROR_EXPIRED_OR_MISSING_DEVICES = 409;
 const ERROR_STALE_DEVICES = 410;
@@ -861,7 +862,7 @@ async function encryptForSenderKey({
     parseIntOrThrow(ourDeviceId, 'encryptForSenderKey, ourDeviceId')
   );
   const ourAddress = getOurAddress();
-  const senderKeyStore = new SenderKeys({ ourUuid });
+  const senderKeyStore = new SenderKeys({ ourUuid, zone: GLOBAL_ZONE });
   const message = Buffer.from(padMessage(contentMessage));
 
   const ciphertextMessage =