Eliminate extra preKeys, fail early on key creation if no PNI identity key

Co-authored-by: Fedor Indutny <79877362+indutny-signal@users.noreply.github.com>

ts/sql/migrations/47-further-optimize.ts

@@ -122,7 +122,7 @@ export default function updateToSchemaVersion47(
 
     const ourUuid = getOurUuid(db);
     if (!ourUuid) {
-      logger.warn('updateToSchemaVersion47: our UUID not found');
+      logger.info('updateToSchemaVersion47: our UUID not found');
     } else {
       db.prepare<Query>(
         `

ts/sql/migrations/88-service-ids.ts

@@ -509,14 +509,25 @@ function migrateItems(db: Database, logger: LoggerType): OurServiceIds {
   try {
     [legacyAci] = JSON.parse(uuidIdJson).value.split('.', 2);
   } catch (error) {
-    logger.warn('updateToSchemaVersion88: failed to parse uuid_id item', error);
+    if (uuidIdJson) {
+      logger.warn(
+        'updateToSchemaVersion88: failed to parse uuid_id item',
+        error
+      );
+    } else {
+      logger.info('updateToSchemaVersion88: Our UUID not found');
+    }
   }
 
   let legacyPni: string | undefined;
   try {
     legacyPni = JSON.parse(pniJson).value;
   } catch (error) {
-    logger.warn('updateToSchemaVersion88: failed to parse pni item', error);
+    if (pniJson) {
+      logger.warn('updateToSchemaVersion88: failed to parse pni item', error);
+    } else {
+      logger.info('updateToSchemaVersion88: Our PNI not found');
+    }
   }
 
   const aci = normalizeAci(legacyAci, 'uuid_id', logger);

ts/sql/migrations/91-clean-keys.ts

@@ -0,0 +1,177 @@
+// Copyright 2023 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import type { Database } from '@signalapp/better-sqlite3';
+
+import type { LoggerType } from '../../types/Logging';
+import { sql } from '../util';
+import type { PniString } from '../../types/ServiceId';
+import { normalizePni } from '../../types/ServiceId';
+import * as Errors from '../../types/errors';
+
+export default function updateToSchemaVersion91(
+  currentVersion: number,
+  db: Database,
+  logger: LoggerType
+): void {
+  if (currentVersion >= 91) {
+    return;
+  }
+
+  db.transaction(() => {
+    // Fix the ourServiceId column so it's generated from the right JSON field
+
+    db.exec(`
+      --- First, prekeys
+      DROP INDEX preKeys_ourServiceId;
+      
+      ALTER TABLE preKeys
+        DROP COLUMN ourServiceId;
+      ALTER TABLE preKeys
+        ADD COLUMN ourServiceId NUMBER
+        GENERATED ALWAYS AS (json_extract(json, '$.ourServiceId'));
+      
+      CREATE INDEX preKeys_ourServiceId ON preKeys (ourServiceId);    
+
+      -- Second, kyber prekeys
+
+      DROP INDEX kyberPreKeys_ourServiceId;
+      
+      ALTER TABLE kyberPreKeys
+        DROP COLUMN ourServiceId;
+      ALTER TABLE kyberPreKeys
+        ADD COLUMN ourServiceId NUMBER
+        GENERATED ALWAYS AS (json_extract(json, '$.ourServiceId'));
+      
+      CREATE INDEX kyberPreKeys_ourServiceId ON kyberPreKeys (ourServiceId); 
+
+      -- Finally, signed prekeys
+
+      DROP INDEX signedPreKeys_ourServiceId;
+      
+      ALTER TABLE signedPreKeys
+        DROP COLUMN ourServiceId;
+      ALTER TABLE signedPreKeys
+        ADD COLUMN ourServiceId NUMBER
+        GENERATED ALWAYS AS (json_extract(json, '$.ourServiceId'));
+      
+      CREATE INDEX signedPreKeys_ourServiceId ON signedPreKeys (ourServiceId); 
+    `);
+
+    // Do overall count - if it's less than 1000, move on
+
+    const totalKeys = db
+      .prepare('SELECT count(*) FROM preKeys;')
+      .pluck(true)
+      .get();
+    logger.info(`updateToSchemaVersion91: Found ${totalKeys} keys`);
+    if (totalKeys < 1000) {
+      db.pragma('user_version = 91');
+      return;
+    }
+
+    // Grab our PNI
+
+    let pni: PniString;
+    const pniJson = db
+      .prepare("SELECT json FROM items WHERE id IS 'pni'")
+      .pluck()
+      .get();
+    try {
+      const pniData = JSON.parse(pniJson);
+      pni = normalizePni(pniData.value, 'updateToSchemaVersion91');
+    } catch (error) {
+      db.pragma('user_version = 91');
+      if (pniJson) {
+        logger.warn(
+          'updateToSchemaVersion91: PNI found but did not parse',
+          Errors.toLogFormat(error)
+        );
+      } else {
+        logger.info('updateToSchemaVersion91: Our PNI not found');
+      }
+      return;
+    }
+
+    // Create index to help us with all these queries
+
+    db.exec(`
+      ALTER TABLE preKeys
+        ADD COLUMN createdAt NUMBER
+          GENERATED ALWAYS AS (json_extract(json, '$.createdAt'));
+      
+      CREATE INDEX preKeys_date
+        ON preKeys (ourServiceId, createdAt);
+    `);
+
+    // Grab PNI-specific count - if it's less than 1000, move on
+
+    const [
+      beforeQuery,
+      beforeParams,
+    ] = sql`SELECT count(*) from preKeys WHERE ourServiceId = ${pni}`;
+    const beforeKeys = db.prepare(beforeQuery).pluck(true).get(beforeParams);
+    logger.info(`updateToSchemaVersion91: Found ${beforeKeys} preKeys for PNI`);
+
+    // Fetch 500th-oldest timestamp for PNI
+
+    const [oldQuery, oldParams] = sql`
+      SELECT createdAt
+      FROM preKeys
+      WHERE 
+        createdAt IS NOT NULL AND
+        ourServiceId = ${pni}
+      ORDER BY createdAt ASC
+      LIMIT 1
+      OFFSET 499
+    `;
+    const oldBoundary = db.prepare(oldQuery).pluck(true).get(oldParams);
+
+    // Fetch 500th-newest timestamp for PNI
+
+    const [newQuery, newParams] = sql`
+      SELECT createdAt
+      FROM preKeys
+      WHERE 
+        createdAt IS NOT NULL AND
+        ourServiceId = ${pni}
+      ORDER BY createdAt DESC
+      LIMIT 1
+      OFFSET 499
+    `;
+    const newBoundary = db.prepare(newQuery).pluck(true).get(newParams);
+
+    // Delete everything in between for PNI
+
+    const [deleteQuery, deleteParams] = sql`
+      DELETE FROM preKeys
+      WHERE 
+        createdAt IS NOT NULL AND
+        createdAt > ${oldBoundary} AND 
+        createdAt < ${newBoundary} AND
+        ourServiceId = ${pni};
+    `;
+    db.prepare(deleteQuery).run(deleteParams);
+
+    // Get updated count for PNI
+
+    const [afterQuery, afterParams] = sql`
+      SELECT count(*)
+      FROM preKeys
+      WHERE ourServiceId = ${pni};
+    `;
+    const afterCount = db.prepare(afterQuery).pluck(true).get(afterParams);
+    logger.info(
+      `updateToSchemaVersion91: Found ${afterCount} preKeys for PNI after delete`
+    );
+
+    db.exec(`
+      DROP INDEX preKeys_date;
+      ALTER TABLE preKeys DROP COLUMN createdAt;
+    `);
+
+    db.pragma('user_version = 91');
+  })();
+
+  logger.info('updateToSchemaVersion91: success!');
+}

ts/sql/migrations/index.ts

@@ -65,6 +65,7 @@ import updateToSchemaVersion86 from './86-story-replies-index';
 import updateToSchemaVersion88 from './88-service-ids';
 import updateToSchemaVersion89 from './89-call-history';
 import updateToSchemaVersion90 from './90-delete-story-reply-screenshot';
+import updateToSchemaVersion91 from './91-clean-keys';
 
 function updateToSchemaVersion1(
   currentVersion: number,
@@ -2001,6 +2002,8 @@ export const SCHEMA_VERSIONS = [
   updateToSchemaVersion88,
   updateToSchemaVersion89,
   updateToSchemaVersion90,
+
+  updateToSchemaVersion91,
 ];
 
 export function updateSchema(db: Database, logger: LoggerType): void {