Add transaction support to kyber triples

Author: indutny-signal

ts/SignalProtocolStore.ts

@@ -20,7 +20,7 @@ import {
 } from '@signalapp/libsignal-client';
 
 import { DataReader, DataWriter } from './sql/Client.js';
-import type { ItemType } from './sql/Interface.js';
+import type { ItemType, KyberPreKeyTripleType } from './sql/Interface.js';
 import * as Bytes from './Bytes.js';
 import { constantTimeEqual, sha256 } from './Crypto.js';
 import { assertDev, strictAssert } from './util/assert.js';
@@ -213,6 +213,19 @@ export function hydrateSignedPreKey(
   );
 }
 
+// Format: keyId:signedPreKeyId:baseKey
+type KyberTripleCacheKeyPrefixType = `${KyberPreKeyTripleType['id']}:`;
+type KyberTripleCacheKeyType =
+  `${KyberTripleCacheKeyPrefixType}${KyberPreKeyTripleType['signedPreKeyId']}:${string}`;
+
+function getKyberTripleCacheKey({
+  id,
+  signedPreKeyId,
+  baseKey,
+}: KyberPreKeyTripleType): KyberTripleCacheKeyType {
+  return `${id}:${signedPreKeyId}:${Bytes.toHex(baseKey)}`;
+}
+
 type SessionCacheEntry = CacheEntryType<SessionType, SessionRecord>;
 type SenderKeyCacheEntry = CacheEntryType<SenderKeyType, SenderKeyRecord>;
 
@@ -254,6 +267,8 @@ export class SignalProtocolStore extends EventEmitter {
     CacheEntryType<SignedPreKeyType, SignedPreKeyRecord>
   >;
 
+  readonly #kyberTriples = new Set<KyberTripleCacheKeyType>();
+
   senderKeyQueues = new Map<QualifiedAddressStringType, PQueue>();
 
   sessionQueues = new Map<SessionIdType, PQueue>();
@@ -267,6 +282,10 @@ export class SignalProtocolStore extends EventEmitter {
   #pendingSessions = new Map<SessionIdType, SessionCacheEntry>();
   #pendingSenderKeys = new Map<SenderKeyIdType, SenderKeyCacheEntry>();
   #pendingUnprocessed = new Map<string, UnprocessedType>();
+  #pendingKyberTriples = new Map<
+    KyberTripleCacheKeyType,
+    KyberPreKeyTripleType
+  >();
 
   async hydrateCaches(): Promise<void> {
     await Promise.all([
@@ -310,6 +329,15 @@ export class SignalProtocolStore extends EventEmitter {
           this.#ourRegistrationIds.set(serviceId, map.value[serviceId]);
         }
       })(),
+      (async () => {
+        this.#kyberTriples.clear();
+
+        const triples = await DataReader.getAllKyberTriples();
+
+        for (const t of triples) {
+          this.#kyberTriples.add(getKyberTripleCacheKey(t));
+        }
+      })(),
       _fillCaches<string, IdentityKeyType, PublicKey>(
         this,
         'identityKeys',
@@ -525,14 +553,30 @@ export class SignalProtocolStore extends EventEmitter {
       `maybeRemoveKyberPreKey: Not removing kyber prekey ${id}; it's a last resort key`
     );
 
-    const result = await DataWriter.markKyberTripleSeenOrFail({
-      id: `${ourServiceId}:${keyId}`,
-      signedPreKeyId,
-      baseKey: baseKey.serialize(),
+    await this.withZone(zone, 'maybeRemoveKyberPreKey', async () => {
+      const triple: KyberPreKeyTripleType = {
+        id: `${ourServiceId}:${keyId}`,
+        signedPreKeyId,
+        baseKey: baseKey.serialize(),
+      };
+
+      const cacheKey = getKyberTripleCacheKey(triple);
+
+      // Note: we don't have to check for `#pendingKyberPreKeysToRemove` since
+      // it makes the key in question inaccessible to begin with.
+      if (
+        this.#kyberTriples.has(cacheKey) ||
+        this.#pendingKyberTriples.has(cacheKey)
+      ) {
+        throw new Error(`Duplicate kyber triple ${keyId}:${signedPreKeyId}`);
+      }
+
+      this.#pendingKyberTriples.set(cacheKey, triple);
+
+      if (!zone.supportsPendingKyberPreKeysToRemove()) {
+        await this.#commitZoneChanges('removeKyberPreKeys');
+      }
     });
-    if (result === 'fail') {
-      throw new Error(`Duplicate kyber triple ${keyId}:${signedPreKeyId}`);
-    }
   }
 
   async removeKyberPreKeys(
@@ -1169,13 +1213,15 @@ export class SignalProtocolStore extends EventEmitter {
     const pendingSenderKeys = this.#pendingSenderKeys;
     const pendingSessions = this.#pendingSessions;
     const pendingUnprocessed = this.#pendingUnprocessed;
+    const pendingKyberTriples = this.#pendingKyberTriples;
 
     if (
       pendingKyberPreKeysToRemove.size === 0 &&
       pendingPreKeysToRemove.size === 0 &&
       pendingSenderKeys.size === 0 &&
       pendingSessions.size === 0 &&
-      pendingUnprocessed.size === 0
+      pendingUnprocessed.size === 0 &&
+      pendingKyberTriples.size === 0
     ) {
       return;
     }
@@ -1186,14 +1232,16 @@ export class SignalProtocolStore extends EventEmitter {
         `pending preKeysToRemove ${pendingKyberPreKeysToRemove.size}, ` +
         `pending senderKeys ${pendingSenderKeys.size}, ` +
         `pending sessions ${pendingSessions.size}, ` +
-        `pending unprocessed ${pendingUnprocessed.size}`
+        `pending unprocessed ${pendingUnprocessed.size}, ` +
+        `pending kyberTriples ${pendingKyberTriples.size}`
     );
 
     this.#pendingKyberPreKeysToRemove = new Set();
     this.#pendingPreKeysToRemove = new Set();
     this.#pendingSenderKeys = new Map();
     this.#pendingSessions = new Map();
     this.#pendingUnprocessed = new Map();
+    this.#pendingKyberTriples = new Map();
 
     // Commit both sender keys, sessions and unprocessed in the same database transaction
     //   to unroll both on error.
@@ -1207,17 +1255,30 @@ export class SignalProtocolStore extends EventEmitter {
         ({ fromDB }) => fromDB
       ),
       unprocessed: Array.from(pendingUnprocessed.values()),
+      kyberTriples: Array.from(pendingKyberTriples.values()),
     });
 
     // Apply changes to in-memory storage after successful DB write.
 
+    for (const cacheKey of pendingKyberTriples.keys()) {
+      this.#kyberTriples.add(cacheKey);
+    }
+
     const { kyberPreKeys } = this;
     assertDev(
       kyberPreKeys !== undefined,
       "Can't commit unhydrated kyberPreKeys storage"
     );
-    pendingKyberPreKeysToRemove.forEach(value => {
+    pendingKyberPreKeysToRemove.forEach((value: PreKeyIdType) => {
       kyberPreKeys.delete(value);
+
+      // Remove all cached kyber triples for this key.
+      const prefix: KyberTripleCacheKeyPrefixType = `${value}:`;
+      for (const key of this.#kyberTriples.keys()) {
+        if (key.startsWith(prefix)) {
+          this.#kyberTriples.delete(key);
+        }
+      }
     });
     if (kyberPreKeys.size < LOW_KEYS_THRESHOLD) {
       this.#emitLowKeys(`removeKyberPreKeys@${kyberPreKeys.size}`);

ts/sql/Interface.ts

@@ -738,6 +738,8 @@ type ReadableInterface = {
 
   getAllSessions: () => Array<SessionType>;
 
+  getAllKyberTriples: () => Array<KyberPreKeyTripleType>;
+
   getConversationCount: () => number;
   getConversationById: (id: string) => ConversationType | undefined;
 
@@ -962,9 +964,6 @@ 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;
@@ -1017,6 +1016,7 @@ type WritableInterface = {
     senderKeys: Array<SenderKeyType>;
     sessions: Array<SessionType>;
     unprocessed: Array<UnprocessedType>;
+    kyberTriples: Array<KyberPreKeyTripleType>;
   }): void;
   removeSessionById: (id: SessionIdType) => number;
   removeSessionsByConversation: (conversationId: string) => void;

ts/sql/Server.ts

@@ -397,6 +397,8 @@ export const DataReader: ServerReadableInterface = {
 
   getAllSessions,
 
+  getAllKyberTriples,
+
   getConversationCount,
   getConversationById,
 
@@ -532,7 +534,6 @@ export const DataWriter: ServerWritableInterface = {
   bulkAddKyberPreKeys,
   removeKyberPreKeyById,
   removeKyberPreKeysByServiceId,
-  markKyberTripleSeenOrFail,
   removeAllKyberPreKeys,
 
   createOrUpdatePreKey,
@@ -1080,7 +1081,7 @@ function removeKyberPreKeysByServiceId(
 function markKyberTripleSeenOrFail(
   db: WritableDB,
   { id, signedPreKeyId, baseKey }: KyberPreKeyTripleType
-): 'seen' | 'fail' {
+): void {
   // Notes that `kyberPreKey_triples` has
   // - Unique constraint on id, signedPreKeyId, baseKey so that we can't insert
   //   two identical rows
@@ -1095,10 +1096,9 @@ function markKyberTripleSeenOrFail(
 
   try {
     db.prepare(query).run(parameters);
-    return 'seen';
   } catch (error) {
     if (error.code === 'SQLITE_CONSTRAINT_UNIQUE') {
-      return 'fail';
+      throw new Error(`Duplicate kyber triple ${id}:${signedPreKeyId}`);
     }
 
     // Unexpected error
@@ -1699,12 +1699,14 @@ function commitDecryptResult(
     senderKeys,
     sessions,
     unprocessed,
+    kyberTriples,
   }: {
     kyberPreKeysToRemove: Array<PreKeyIdType>;
     preKeysToRemove: Array<PreKeyIdType>;
     senderKeys: Array<SenderKeyType>;
     sessions: Array<SessionType>;
     unprocessed: Array<UnprocessedType>;
+    kyberTriples: Array<KyberPreKeyTripleType>;
   }
 ): void {
   db.transaction(() => {
@@ -1745,6 +1747,10 @@ function commitDecryptResult(
     for (const item of unprocessed) {
       saveUnprocessed(db, item);
     }
+
+    for (const item of kyberTriples) {
+      markKyberTripleSeenOrFail(db, item);
+    }
   })();
 }
 
@@ -1783,6 +1789,9 @@ function removeAllSessions(db: WritableDB): number {
 function getAllSessions(db: ReadableDB): Array<SessionType> {
   return db.prepare('SELECT * FROM sessions').all();
 }
+function getAllKyberTriples(db: ReadableDB): Array<KyberPreKeyTripleType> {
+  return db.prepare('SELECT * FROM kyberPreKey_triples').all();
+}
 // Conversations
 
 function getConversationCount(db: ReadableDB): number {