Fix deadlock in saveIdentity

ts/SignalProtocolStore.ts

@@ -1044,15 +1044,27 @@ export class SignalProtocolStore extends EventEmitter {
     });
   }
 
-  private _getIdentityQueue(serviceId: ServiceIdString): PQueue {
+  private _runOnIdentityQueue<T>(
+    serviceId: ServiceIdString,
+    zone: Zone,
+    name: string,
+    body: () => Promise<T>
+  ): Promise<T> {
+    let queue: PQueue;
+
     const cachedQueue = this.identityQueues.get(serviceId);
     if (cachedQueue) {
-      return cachedQueue;
+      queue = cachedQueue;
+    } else {
+      queue = this._createIdentityQueue();
+      this.identityQueues.set(serviceId, queue);
     }
 
-    const freshQueue = this._createIdentityQueue();
-    this.identityQueues.set(serviceId, freshQueue);
-    return freshQueue;
+    // We run the identity queue task in zone because `saveIdentity` needs to
+    // be able to archive sibling sessions on keychange. Not entering the zone
+    // now would mean that we can take locks in different order here and in
+    // MessageReceiver which will lead to a deadlock.
+    return this.withZone(zone, name, () => queue.add(body));
   }
 
   // Sessions
@@ -1991,7 +2003,7 @@ export class SignalProtocolStore extends EventEmitter {
     encodedAddress: Address,
     publicKey: Uint8Array,
     nonblockingApproval = false,
-    { zone }: SessionTransactionOptions = {}
+    { zone = GLOBAL_ZONE }: SessionTransactionOptions = {}
   ): Promise<boolean> {
     if (!this.identityKeys) {
       throw new Error('saveIdentity: this.identityKeys not yet cached!');
@@ -2009,7 +2021,11 @@ export class SignalProtocolStore extends EventEmitter {
       nonblockingApproval = false;
     }
 
-    return this._getIdentityQueue(encodedAddress.serviceId).add(async () => {
+    return this._runOnIdentityQueue(
+      encodedAddress.serviceId,
+      zone,
+      'saveIdentity',
+      async () => {
         const identityRecord = await this.getOrMigrateIdentityRecord(
           encodedAddress.serviceId
         );
@@ -2107,7 +2123,8 @@ export class SignalProtocolStore extends EventEmitter {
         }
 
         return false;
-    });
+      }
+    );
   }
 
   // https://github.com/signalapp/Signal-Android/blob/fc3db538bcaa38dc149712a483d3032c9c1f3998/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalBaseIdentityKeyStore.java#L257
@@ -2125,9 +2142,14 @@ export class SignalProtocolStore extends EventEmitter {
     serviceId: ServiceIdString,
     attributes: Partial<IdentityKeyType>
   ): Promise<void> {
-    return this._getIdentityQueue(serviceId).add(async () => {
+    return this._runOnIdentityQueue(
+      serviceId,
+      GLOBAL_ZONE,
+      'saveIdentityWithAttributes',
+      async () => {
         return this.saveIdentityWithAttributesOnQueue(serviceId, attributes);
-    });
+      }
+    );
   }
 
   private async saveIdentityWithAttributesOnQueue(
@@ -2172,7 +2194,11 @@ export class SignalProtocolStore extends EventEmitter {
       throw new Error('setApproval: Invalid approval status');
     }
 
-    return this._getIdentityQueue(serviceId).add(async () => {
+    return this._runOnIdentityQueue(
+      serviceId,
+      GLOBAL_ZONE,
+      'setApproval',
+      async () => {
         const identityRecord = await this.getOrMigrateIdentityRecord(serviceId);
 
         if (!identityRecord) {
@@ -2181,7 +2207,8 @@ export class SignalProtocolStore extends EventEmitter {
 
         identityRecord.nonblockingApproval = nonblockingApproval;
         await this._saveIdentityKey(identityRecord);
-    });
+      }
+    );
   }
 
   // https://github.com/signalapp/Signal-Android/blob/fc3db538bcaa38dc149712a483d3032c9c1f3998/app/src/main/java/org/thoughtcrime/securesms/crypto/storage/SignalBaseIdentityKeyStore.java#L215
@@ -2198,7 +2225,11 @@ export class SignalProtocolStore extends EventEmitter {
       throw new Error('setVerified: Invalid verified status');
     }
 
-    return this._getIdentityQueue(serviceId).add(async () => {
+    return this._runOnIdentityQueue(
+      serviceId,
+      GLOBAL_ZONE,
+      'setVerified',
+      async () => {
         const identityRecord = await this.getOrMigrateIdentityRecord(serviceId);
 
         if (!identityRecord) {
@@ -2212,7 +2243,8 @@ export class SignalProtocolStore extends EventEmitter {
             verified: verifiedStatus,
           });
         }
-    });
+      }
+    );
   }
 
   async getVerified(serviceId: ServiceIdString): Promise<number> {
@@ -2279,7 +2311,11 @@ export class SignalProtocolStore extends EventEmitter {
       `Invalid verified status: ${verifiedStatus}`
     );
 
-    return this._getIdentityQueue(serviceId).add(async () => {
+    return this._runOnIdentityQueue(
+      serviceId,
+      GLOBAL_ZONE,
+      'updateIdentityAfterSync',
+      async () => {
         const identityRecord = await this.getOrMigrateIdentityRecord(serviceId);
         const hadEntry = identityRecord !== undefined;
         const keyMatches = Boolean(
@@ -2299,10 +2335,18 @@ export class SignalProtocolStore extends EventEmitter {
           });
         }
         if (!hadEntry) {
-        this.checkPreviousKey(serviceId, publicKey, 'updateIdentityAfterSync');
+          this.checkPreviousKey(
+            serviceId,
+            publicKey,
+            'updateIdentityAfterSync'
+          );
         } else if (hadEntry && !keyMatches) {
           try {
-          this.emit('keychange', serviceId, 'updateIdentityAfterSync - change');
+            this.emit(
+              'keychange',
+              serviceId,
+              'updateIdentityAfterSync - change'
+            );
           } catch (error) {
             log.error(
               'updateIdentityAfterSync: error triggering keychange:',
@@ -2328,7 +2372,8 @@ export class SignalProtocolStore extends EventEmitter {
           return true;
         }
         return false;
-    });
+      }
+    );
   }
 
   isUntrusted(

ts/test-electron/SignalProtocolStore_test.ts

@@ -20,6 +20,7 @@ import { v4 as generateUuid } from 'uuid';
 import { signal } from '../protobuf/compiled';
 import { sessionStructureToBytes } from '../util/sessionTranslation';
 import * as durations from '../util/durations';
+import { explodePromise } from '../util/explodePromise';
 import { Zone } from '../util/Zone';
 
 import * as Bytes from '../Bytes';
@@ -262,6 +263,29 @@ describe('SignalProtocolStore', () => {
       await store.saveIdentity(identifier, testKey.pubKey);
       await store.saveIdentity(identifier, newIdentity);
     });
+    it('should not deadlock', async () => {
+      const newIdentity = getPublicKey();
+      const zone = new Zone('zone', {
+        pendingSenderKeys: true,
+        pendingSessions: true,
+        pendingUnprocessed: true,
+      });
+
+      await store.saveIdentity(identifier, testKey.pubKey);
+
+      const { promise, resolve } = explodePromise<void>();
+
+      await Promise.all([
+        store.withZone(zone, 'test', async () => {
+          await promise;
+          return store.saveIdentity(identifier, newIdentity, false, { zone });
+        }),
+        store.saveIdentity(identifier, newIdentity, false, {
+          zone: GLOBAL_ZONE,
+        }),
+        resolve(),
+      ]);
+    });
 
     describe('When there is no existing key (first use)', () => {
       before(async () => {