Introduce kyber pre key triple table

Author: indutny-signal

ts/LibSignalStores.ts

@@ -250,10 +250,14 @@ export class KyberPreKeys extends KyberPreKeyStore {
     return kyberPreKey;
   }
 
-  async markKyberPreKeyUsed(id: number): Promise<void> {
+  async markKyberPreKeyUsed(
+    keyId: number,
+    signedPreKeyId: number,
+    baseKey: PublicKey
+  ): Promise<void> {
     await window.textsecure.storage.protocol.maybeRemoveKyberPreKey(
       this.#ourServiceId,
-      id,
+      { keyId, signedPreKeyId, baseKey },
       { zone: this.#zone }
     );
   }

ts/SignalProtocolStore.ts

@@ -503,7 +503,11 @@ export class SignalProtocolStore extends EventEmitter {
 
   async maybeRemoveKyberPreKey(
     ourServiceId: ServiceIdString,
-    keyId: number,
+    {
+      keyId,
+      signedPreKeyId,
+      baseKey,
+    }: { keyId: number; signedPreKeyId: number; baseKey: PublicKey },
     { zone = GLOBAL_ZONE }: SessionTransactionOptions = {}
   ): Promise<void> {
     const id: PreKeyIdType = this.#_getKeyId(ourServiceId, keyId);
@@ -512,14 +516,23 @@ export class SignalProtocolStore extends EventEmitter {
     if (!entry) {
       return;
     }
-    if (entry.fromDB.isLastResort) {
-      log.info(
-        `maybeRemoveKyberPreKey: Not removing kyber prekey ${id}; it's a last resort key`
-      );
+    if (!entry.fromDB.isLastResort) {
+      await this.removeKyberPreKeys(ourServiceId, [keyId], { zone });
       return;
     }
 
-    await this.removeKyberPreKeys(ourServiceId, [keyId], { zone });
+    log.info(
+      `maybeRemoveKyberPreKey: Not removing kyber prekey ${id}; it's a last resort key`
+    );
+
+    const result = await DataWriter.markKyberTripleSeenOrFail({
+      id: `${ourServiceId}:${keyId}`,
+      signedPreKeyId,
+      baseKey: baseKey.serialize(),
+    });
+    if (result === 'fail') {
+      throw new Error(`Duplicate kyber triple ${keyId}:${signedPreKeyId}`);
+    }
   }
 
   async removeKyberPreKeys(

ts/sql/Interface.ts

@@ -590,6 +590,12 @@ export type MediaItemDBType = Readonly<{
   message: MediaItemMessageType;
 }>;
 
+export type KyberPreKeyTripleType = Readonly<{
+  id: PreKeyIdType;
+  signedPreKeyId: number;
+  baseKey: Uint8Array;
+}>;
+
 export const MESSAGE_ATTACHMENT_COLUMNS = [
   'messageId',
   'conversationId',
@@ -956,6 +962,9 @@ type WritableInterface = {
   removeKyberPreKeyById: (id: PreKeyIdType | Array<PreKeyIdType>) => number;
   removeKyberPreKeysByServiceId: (serviceId: ServiceIdString) => void;
   removeAllKyberPreKeys: () => number;
+  markKyberTripleSeenOrFail: (
+    options: KyberPreKeyTripleType
+  ) => 'seen' | 'fail';
 
   removePreKeyById: (id: PreKeyIdType | Array<PreKeyIdType>) => number;
   removePreKeysByServiceId: (serviceId: ServiceIdString) => void;

ts/sql/Server.ts

@@ -140,6 +140,7 @@ import type {
   GetUnreadByConversationAndMarkReadResultType,
   IdentityKeyIdType,
   ItemKeyType,
+  KyberPreKeyTripleType,
   MediaItemDBType,
   MessageAttachmentsCursorType,
   MessageCursorType,
@@ -531,6 +532,7 @@ export const DataWriter: ServerWritableInterface = {
   bulkAddKyberPreKeys,
   removeKyberPreKeyById,
   removeKyberPreKeysByServiceId,
+  markKyberTripleSeenOrFail,
   removeAllKyberPreKeys,
 
   createOrUpdatePreKey,
@@ -1075,6 +1077,34 @@ function removeKyberPreKeysByServiceId(
     serviceId,
   });
 }
+function markKyberTripleSeenOrFail(
+  db: WritableDB,
+  { id, signedPreKeyId, baseKey }: KyberPreKeyTripleType
+): 'seen' | 'fail' {
+  // Notes that `kyberPreKey_triples` has
+  // - Unique constraint on id, signedPreKeyId, baseKey so that we can't insert
+  //   two identical rows
+  // - `ON DELETE CASCADE` trigger linked to `kyberPreKeys` table so that we
+  //   cleanup the triples whenever we remove the key
+  const [query, parameters] = sql`
+    INSERT OR FAIL INTO kyberPreKey_triples
+      (id, signedPreKeyId, baseKey)
+    VALUES
+      (${id}, ${signedPreKeyId}, ${baseKey});
+  `;
+
+  try {
+    db.prepare(query).run(parameters);
+    return 'seen';
+  } catch (error) {
+    if (error.code === 'SQLITE_CONSTRAINT_UNIQUE') {
+      return 'fail';
+    }
+
+    // Unexpected error
+    throw error;
+  }
+}
 function removeAllKyberPreKeys(db: WritableDB): number {
   return removeAllFromTable(db, KYBER_PRE_KEYS_TABLE);
 }

ts/sql/migrations/1470-kyber-triple.ts

@@ -0,0 +1,15 @@
+// Copyright 2025 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import type { WritableDB } from '../Interface.js';
+
+export default function updateToSchemaVersion1460(db: WritableDB): void {
+  db.exec(`
+    CREATE TABLE kyberPreKey_triples (
+      id TEXT NOT NULL REFERENCES kyberPreKeys(id) ON DELETE CASCADE,
+      signedPreKeyId INTEGER NOT NULL,
+      baseKey BLOB NOT NULL,
+      UNIQUE(id, signedPreKeyId, baseKey) ON CONFLICT FAIL
+    ) STRICT;
+  `);
+}

ts/sql/migrations/index.ts

@@ -122,6 +122,7 @@ import updateToSchemaVersion1430 from './1430-call-links-epoch-id.js';
 import updateToSchemaVersion1440 from './1440-chat-folders.js';
 import updateToSchemaVersion1450 from './1450-all-media.js';
 import updateToSchemaVersion1460 from './1460-attachment-duration.js';
+import updateToSchemaVersion1470 from './1470-kyber-triple.js';
 
 import { DataWriter } from '../Server.js';
 
@@ -1601,6 +1602,7 @@ export const SCHEMA_VERSIONS: ReadonlyArray<SchemaUpdateType> = [
   { version: 1440, update: updateToSchemaVersion1440 },
   { version: 1450, update: updateToSchemaVersion1450 },
   { version: 1460, update: updateToSchemaVersion1460 },
+  { version: 1470, update: updateToSchemaVersion1470 },
 ];
 
 export class DBVersionFromFutureError extends Error {

ts/test-electron/SignalProtocolStore_test.ts

@@ -27,6 +27,7 @@ import {
   clampPrivateKey,
   setPublicKeyTypeByte,
   generateSignedPreKey,
+  generateKyberPreKey,
 } from '../Curve.js';
 import type { SignalProtocolStore } from '../SignalProtocolStore.js';
 import { GLOBAL_ZONE } from '../SignalProtocolStore.js';
@@ -597,12 +598,9 @@ describe('SignalProtocolStore', () => {
       });
 
       async function testInvalidAttributes() {
-        try {
-          await store.saveIdentityWithAttributes(theirAci, attributes);
-          throw new Error('saveIdentityWithAttributes should have failed');
-        } catch (error) {
-          // good. we expect to fail with invalid attributes.
-        }
+        await assert.isRejected(
+          store.saveIdentityWithAttributes(theirAci, attributes)
+        );
       }
 
       it('rejects an invalid publicKey', async () => {
@@ -1637,4 +1635,96 @@ describe('SignalProtocolStore', () => {
       // Note: signature is ignored.
     });
   });
+
+  describe('maybeRemoveKyberPreKey', () => {
+    beforeEach(async () => {
+      await store.clearKyberPreKeyStore();
+    });
+
+    afterEach(async () => {
+      await store.clearKyberPreKeyStore();
+    });
+
+    it('should detect duplicate triples', async () => {
+      await store.storeKyberPreKeys(ourAci, [
+        {
+          createdAt: Date.now(),
+          data: generateKyberPreKey(identityKey, 1).serialize(),
+          isConfirmed: true,
+          isLastResort: true,
+          keyId: 1,
+          ourServiceId: ourAci,
+        },
+      ]);
+
+      await store.maybeRemoveKyberPreKey(ourAci, {
+        keyId: 1,
+        signedPreKeyId: 1,
+        baseKey: testKey.publicKey,
+      });
+
+      await assert.isRejected(
+        store.maybeRemoveKyberPreKey(ourAci, {
+          keyId: 1,
+          signedPreKeyId: 1,
+          baseKey: testKey.publicKey,
+        }),
+        'Duplicate kyber triple 1:1'
+      );
+    });
+
+    it('should ignore triples for non last resort keys', async () => {
+      await store.storeKyberPreKeys(ourAci, [
+        {
+          createdAt: Date.now(),
+          data: generateKyberPreKey(identityKey, 1).serialize(),
+          isConfirmed: true,
+          isLastResort: false,
+          keyId: 1,
+          ourServiceId: ourAci,
+        },
+      ]);
+
+      await store.maybeRemoveKyberPreKey(ourAci, {
+        keyId: 1,
+        signedPreKeyId: 1,
+        baseKey: testKey.publicKey,
+      });
+
+      // this should not throw since the key was not last resort
+      await store.maybeRemoveKyberPreKey(ourAci, {
+        keyId: 1,
+        signedPreKeyId: 1,
+        baseKey: testKey.publicKey,
+      });
+    });
+
+    it('should remove triples when removing the key', async () => {
+      await store.storeKyberPreKeys(ourAci, [
+        {
+          createdAt: Date.now(),
+          data: generateKyberPreKey(identityKey, 1).serialize(),
+          isConfirmed: true,
+          isLastResort: true,
+          keyId: 1,
+          ourServiceId: ourAci,
+        },
+      ]);
+
+      await store.maybeRemoveKyberPreKey(ourAci, {
+        keyId: 1,
+        signedPreKeyId: 1,
+        baseKey: testKey.publicKey,
+      });
+
+      await store.removeKyberPreKeys(ourAci, [1]);
+
+      // this should not throw since we removed the key
+      await store.maybeRemoveKyberPreKey(ourAci, {
+        keyId: 1,
+        signedPreKeyId: 1,
+        baseKey: testKey.publicKey,
+      });
+    });
+  });
 });

ts/textsecure/WebsocketResources.ts

@@ -363,7 +363,7 @@ export function connectAuthenticatedLibsignal({
   const listener: LibsignalWebSocketResourceHolder & ChatServiceListener = {
     resource: undefined,
     onIncomingMessage(
-      envelope: Buffer,
+      envelope: Uint8Array,
       timestamp: number,
       ack: ChatServerMessageAck
     ): void {