From 373f1218876e4a6f6ac4276c338781f5213c6712 Mon Sep 17 00:00:00 2001
From: Fedor Indutny <79877362+indutny-signal@users.noreply.github.com>
Date: Tue, 27 Sep 2022 09:33:56 -0700
Subject: [PATCH] Extend cleanSignedPreKeys to support PNI

---
 .../textsecure/AccountManager_test.ts         | 11 ++++----
 ts/textsecure/AccountManager.ts               | 26 +++++++++++--------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/ts/test-electron/textsecure/AccountManager_test.ts b/ts/test-electron/textsecure/AccountManager_test.ts
index 6fb9ba11e..d585a7b36 100644
--- a/ts/test-electron/textsecure/AccountManager_test.ts
+++ b/ts/test-electron/textsecure/AccountManager_test.ts
@@ -2,12 +2,11 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
 import { assert } from 'chai';
-import { v4 as getGuid } from 'uuid';
 
 import { getRandomBytes } from '../../Crypto';
 import AccountManager from '../../textsecure/AccountManager';
 import type { OuterSignedPrekeyType } from '../../textsecure/Types.d';
-import { UUID } from '../../types/UUID';
+import { UUID, UUIDKind } from '../../types/UUID';
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
@@ -32,7 +31,7 @@ describe('AccountManager', () => {
     const identityKey = window.Signal.Curve.generateKeyPair();
 
     beforeEach(async () => {
-      const ourUuid = new UUID(getGuid());
+      const ourUuid = UUID.generate();
 
       originalGetUuid = window.textsecure.storage.user.getUuid;
       originalGetIdentityKeyPair =
@@ -107,7 +106,7 @@ describe('AccountManager', () => {
       ];
 
       // should be no calls to store.removeSignedPreKey, would cause crash
-      return accountManager.cleanSignedPreKeys();
+      return accountManager.cleanSignedPreKeys(UUIDKind.ACI);
     });
 
     it('eliminates oldest keys, even if recent key is unconfirmed', async () => {
@@ -170,7 +169,7 @@ describe('AccountManager', () => {
         count += 1;
       };
 
-      await accountManager.cleanSignedPreKeys();
+      await accountManager.cleanSignedPreKeys(UUIDKind.ACI);
       assert.strictEqual(count, 1);
     });
 
@@ -211,7 +210,7 @@ describe('AccountManager', () => {
         throw new Error('None should be removed!');
       };
 
-      await accountManager.cleanSignedPreKeys();
+      await accountManager.cleanSignedPreKeys(UUIDKind.ACI);
     });
   });
 });
diff --git a/ts/textsecure/AccountManager.ts b/ts/textsecure/AccountManager.ts
index 961a01f3a..0dfb2f0d5 100644
--- a/ts/textsecure/AccountManager.ts
+++ b/ts/textsecure/AccountManager.ts
@@ -441,7 +441,10 @@ export default class AccountManager extends EventTarget {
       ]);
 
       try {
-        await this.cleanSignedPreKeys();
+        await Promise.all([
+          this.cleanSignedPreKeys(UUIDKind.ACI),
+          this.cleanSignedPreKeys(UUIDKind.PNI),
+        ]);
       } catch (_error) {
         // Ignoring the error
       }
@@ -455,9 +458,10 @@ export default class AccountManager extends EventTarget {
     return this.pendingQueue.add(taskWithTimeout);
   }
 
-  async cleanSignedPreKeys(): Promise<void> {
-    const ourUuid = window.textsecure.storage.user.getCheckedUuid();
+  async cleanSignedPreKeys(uuidKind: UUIDKind): Promise<void> {
+    const ourUuid = window.textsecure.storage.user.getCheckedUuid(uuidKind);
     const store = window.textsecure.storage.protocol;
+    const logId = `AccountManager.cleanSignedPreKeys(${uuidKind})`;
 
     const allKeys = await store.loadSignedPreKeys(ourUuid);
     allKeys.sort((a, b) => (b.created_at || 0) - (a.created_at || 0));
@@ -467,15 +471,13 @@ export default class AccountManager extends EventTarget {
     const recent = allKeys[0] ? allKeys[0].keyId : 'none';
     const recentConfirmed = confirmed[0] ? confirmed[0].keyId : 'none';
     const recentUnconfirmed = unconfirmed[0] ? unconfirmed[0].keyId : 'none';
-    log.info(`cleanSignedPreKeys: Most recent signed key: ${recent}`);
+    log.info(`${logId}: Most recent signed key: ${recent}`);
+    log.info(`${logId}: Most recent confirmed signed key: ${recentConfirmed}`);
     log.info(
-      `cleanSignedPreKeys: Most recent confirmed signed key: ${recentConfirmed}`
+      `${logId}: Most recent unconfirmed signed key: ${recentUnconfirmed}`
     );
     log.info(
-      `cleanSignedPreKeys: Most recent unconfirmed signed key: ${recentUnconfirmed}`
-    );
-    log.info(
-      'cleanSignedPreKeys: Total signed key count:',
+      `${logId}: Total signed key count:`,
       allKeys.length,
       '-',
       confirmed.length,
@@ -494,7 +496,8 @@ export default class AccountManager extends EventTarget {
           const timestamp = new Date(createdAt).toJSON();
           const confirmedText = key.confirmed ? ' (confirmed)' : '';
           log.info(
-            `Removing signed prekey: ${key.keyId} with timestamp ${timestamp}${confirmedText}`
+            `${logId}: Removing signed prekey: ${key.keyId} with ` +
+              `timestamp ${timestamp}${confirmedText}`
           );
           await store.removeSignedPreKey(ourUuid, key.keyId);
         }
@@ -830,7 +833,8 @@ export default class AccountManager extends EventTarget {
     await Promise.all(promises);
 
     // This is primarily for the signed prekey summary it logs out
-    this.cleanSignedPreKeys();
+    this.cleanSignedPreKeys(UUIDKind.ACI);
+    this.cleanSignedPreKeys(UUIDKind.PNI);
 
     return {
       ...result,