From ff019a24902bac15a62b39d5daf939998144129c Mon Sep 17 00:00:00 2001
From: Fedor Indutny <79877362+indutny-signal@users.noreply.github.com>
Date: Thu, 5 Mar 2026 11:11:57 -0800
Subject: [PATCH] Simplify prekey id generation

---
 ts/SignalProtocolStore.preload.ts       |  5 ++--
 ts/textsecure/AccountManager.preload.ts | 33 +++++++++++++++++--------
 ts/util/wrappingAdd.std.ts              | 12 +++++++++
 3 files changed, 38 insertions(+), 12 deletions(-)
 create mode 100644 ts/util/wrappingAdd.std.ts

diff --git a/ts/SignalProtocolStore.preload.ts b/ts/SignalProtocolStore.preload.ts
index d1af40d82f..64caaf3b38 100644
--- a/ts/SignalProtocolStore.preload.ts
+++ b/ts/SignalProtocolStore.preload.ts
@@ -72,6 +72,7 @@ import {
 } from './textsecure/AccountManager.preload.js';
 import { formatGroups, groupWhile } from './util/groupWhile.std.js';
 import { parseUnknown } from './util/schemas.std.js';
+import { wrappingAdd24 } from './util/wrappingAdd.std.js';
 import { itemStorage } from './textsecure/Storage.preload.js';
 
 const { omit } = lodash;
@@ -2652,7 +2653,7 @@ export class SignalProtocolStore extends EventEmitter {
         [pni]: registrationId,
       }),
       (async () => {
-        const newId = signedPreKey.id() + 1;
+        const newId = wrappingAdd24(signedPreKey.id(), 1);
         log.warn(`${logId}: Updating next signed pre key id to ${newId}`);
         await itemStorage.put(SIGNED_PRE_KEY_ID_KEY[ServiceIdKind.PNI], newId);
       })(),
@@ -2670,7 +2671,7 @@ export class SignalProtocolStore extends EventEmitter {
         if (!lastResortKyberPreKey) {
           return;
         }
-        const newId = lastResortKyberPreKey.id() + 1;
+        const newId = wrappingAdd24(lastResortKyberPreKey.id(), 1);
         log.warn(`${logId}: Updating next kyber pre key id to ${newId}`);
         await itemStorage.put(KYBER_KEY_ID_KEY[ServiceIdKind.PNI], newId);
       })(),
diff --git a/ts/textsecure/AccountManager.preload.ts b/ts/textsecure/AccountManager.preload.ts
index a14bfa0cf5..1457d8fa39 100644
--- a/ts/textsecure/AccountManager.preload.ts
+++ b/ts/textsecure/AccountManager.preload.ts
@@ -34,7 +34,11 @@ import type {
 import createTaskWithTimeout from './TaskWithTimeout.std.js';
 import * as Bytes from '../Bytes.std.js';
 import * as Errors from '../types/errors.std.js';
-import { isMockEnvironment } from '../environment.std.js';
+import {
+  isTestEnvironment,
+  isMockEnvironment,
+  getEnvironment,
+} from '../environment.std.js';
 import { senderCertificateService } from '../services/senderCertificate.preload.js';
 import {
   decryptDeviceName,
@@ -81,6 +85,7 @@ import { canAttemptRemoteBackupDownload } from '../util/isBackupEnabled.preload.
 import { signalProtocolStore } from '../SignalProtocolStore.preload.js';
 import { itemStorage } from './Storage.preload.js';
 import { deriveAccessKeyFromProfileKey } from '../util/zkgroup.node.js';
+import { wrappingAdd24 } from '../util/wrappingAdd.std.js';
 
 const { isNumber, omit, orderBy } = lodash;
 
@@ -92,7 +97,6 @@ type StorageKeyByServiceIdKind = {
 
 const DAY = 24 * 60 * 60 * 1000;
 
-const STARTING_KEY_ID = 1;
 const PROFILE_KEY_LENGTH = 32;
 const MASTER_KEY_LENGTH = 32;
 const KEY_TOO_OLD_THRESHOLD = 14 * DAY;
@@ -200,12 +204,12 @@ function getNextKeyId(
     return id;
   }
 
-  // For PNI ids, start with existing ACI id
-  if (kind === ServiceIdKind.PNI) {
-    return itemStorage.get(keys[ServiceIdKind.ACI], STARTING_KEY_ID);
+  if (isTestEnvironment(getEnvironment())) {
+    return 1;
   }
 
-  return STARTING_KEY_ID;
+  // eslint-disable-next-line no-bitwise
+  return Buffer.from(getRandomBytes(4)).readUint32LE(0) & 0xffffff;
 }
 
 function kyberPreKeyToUploadSignedPreKey(
@@ -471,7 +475,10 @@ export default class AccountManager extends EventTarget {
 
     await Promise.all([
       signalProtocolStore.storePreKeys(ourServiceId, toSave),
-      itemStorage.put(PRE_KEY_ID_KEY[serviceIdKind], startId + count),
+      itemStorage.put(
+        PRE_KEY_ID_KEY[serviceIdKind],
+        wrappingAdd24(startId, count)
+      ),
     ]);
 
     return toSave.map(key => ({
@@ -520,7 +527,10 @@ export default class AccountManager extends EventTarget {
 
     await Promise.all([
       signalProtocolStore.storeKyberPreKeys(ourServiceId, toSave),
-      itemStorage.put(KYBER_KEY_ID_KEY[serviceIdKind], startId + count),
+      itemStorage.put(
+        KYBER_KEY_ID_KEY[serviceIdKind],
+        wrappingAdd24(startId, count)
+      ),
     ]);
 
     return toUpload;
@@ -679,7 +689,7 @@ export default class AccountManager extends EventTarget {
 
     await itemStorage.put(
       SIGNED_PRE_KEY_ID_KEY[serviceIdKind],
-      signedKeyId + 1
+      wrappingAdd24(signedKeyId, 1)
     );
 
     return key;
@@ -749,7 +759,10 @@ export default class AccountManager extends EventTarget {
     const record = await generateKyberPreKey(identityKey, keyId);
     log.info(`${logId}: Saving new last resort prekey`, keyId);
 
-    await itemStorage.put(KYBER_KEY_ID_KEY[serviceIdKind], kyberKeyId + 1);
+    await itemStorage.put(
+      KYBER_KEY_ID_KEY[serviceIdKind],
+      wrappingAdd24(kyberKeyId, 1)
+    );
 
     return record;
   }
diff --git a/ts/util/wrappingAdd.std.ts b/ts/util/wrappingAdd.std.ts
new file mode 100644
index 0000000000..e6f3029f35
--- /dev/null
+++ b/ts/util/wrappingAdd.std.ts
@@ -0,0 +1,12 @@
+// Copyright 2026 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+// Add two unsigned 24-bit integers and return truncated 24-bit result
+export function wrappingAdd24(a: number, b: number): number {
+  if (!Number.isFinite(a) || !Number.isFinite(b)) {
+    throw new Error('Invalid arguments');
+  }
+
+  // eslint-disable-next-line no-bitwise
+  return (a + b) & 0xffffff;
+}