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 {

ts/test-node/sql/migration_91_test.ts

@@ -0,0 +1,232 @@
+// Copyright 2023 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import type { Database } from '@signalapp/better-sqlite3';
+import SQL from '@signalapp/better-sqlite3';
+import { v4 as generateGuid } from 'uuid';
+import { range } from 'lodash';
+
+import { getTableData, insertData, updateToVersion } from './helpers';
+import type { ServiceIdString } from '../../types/ServiceId';
+import { normalizeAci, normalizePni } from '../../types/ServiceId';
+import type { PreKeyType } from '../../sql/Interface';
+
+type TestingPreKey = Omit<
+  PreKeyType,
+  'privateKey' | 'publicKey' | 'createdAt'
+> & {
+  createdAt: number | undefined;
+};
+
+describe('SQL/updateToSchemaVersion91', () => {
+  let db: Database;
+
+  const OUR_ACI = normalizeAci(generateGuid(), 'updateToSchemaVersion91 test');
+  const OUR_PNI = normalizePni(
+    `PNI:${generateGuid()}`,
+    'updateToSchemaVersion91 test'
+  );
+  let idCount = 0;
+
+  beforeEach(() => {
+    db = new SQL(':memory:');
+    updateToVersion(db, 90);
+  });
+
+  afterEach(() => {
+    db.close();
+  });
+
+  function addPni() {
+    insertData(db, 'items', [
+      {
+        id: 'pni',
+        json: {
+          id: 'pni',
+          value: OUR_PNI,
+        },
+      },
+    ]);
+  }
+
+  function getCountOfKeys(): number {
+    return db.prepare('SELECT count(*) FROM preKeys;').pluck(true).get();
+  }
+
+  function getPragma(): number {
+    return db.prepare('PRAGMA user_version;').pluck(true).get();
+  }
+
+  function generateKey(
+    createdAt: number | undefined,
+    ourServiceId: ServiceIdString
+  ): TestingPreKey {
+    idCount += 1;
+
+    return {
+      createdAt,
+      id: `${ourServiceId}:${idCount}`,
+      keyId: idCount,
+      ourServiceId,
+    };
+  }
+
+  function getRangeOfKeysForInsert(
+    start: number,
+    end: number,
+    ourServiceId: ServiceIdString,
+    options?: {
+      clearCreatedAt?: boolean;
+    }
+  ): Array<{ id: string; json: TestingPreKey }> {
+    return range(start, end).map(createdAt => {
+      const key = generateKey(
+        options?.clearCreatedAt ? undefined : createdAt,
+        ourServiceId
+      );
+
+      return {
+        id: key.id,
+        json: key,
+      };
+    });
+  }
+
+  it('handles missing PNI', () => {
+    assert.strictEqual(0, getCountOfKeys());
+    insertData(db, 'preKeys', getRangeOfKeysForInsert(0, 1500, OUR_ACI));
+    assert.strictEqual(1500, getCountOfKeys());
+    assert.strictEqual(90, getPragma());
+
+    updateToVersion(db, 91);
+
+    assert.strictEqual(91, getPragma());
+    assert.strictEqual(1500, getCountOfKeys());
+  });
+
+  it('deletes 500 extra keys', () => {
+    assert.strictEqual(0, getCountOfKeys());
+    addPni();
+    insertData(db, 'preKeys', getRangeOfKeysForInsert(0, 1500, OUR_PNI));
+    assert.strictEqual(1500, getCountOfKeys());
+    assert.strictEqual(90, getPragma());
+
+    updateToVersion(db, 91);
+
+    assert.strictEqual(91, getPragma());
+    assert.strictEqual(1000, getCountOfKeys());
+  });
+
+  it('leaves 1000 existing keys alone', () => {
+    assert.strictEqual(0, getCountOfKeys());
+    addPni();
+    insertData(db, 'preKeys', getRangeOfKeysForInsert(0, 1000, OUR_PNI));
+    assert.strictEqual(1000, getCountOfKeys());
+    assert.strictEqual(90, getPragma());
+
+    updateToVersion(db, 91);
+
+    assert.strictEqual(91, getPragma());
+    assert.strictEqual(1000, getCountOfKeys());
+  });
+
+  it('leaves keys with missing createdAt alone', () => {
+    assert.strictEqual(0, getCountOfKeys());
+    addPni();
+    insertData(
+      db,
+      'preKeys',
+      getRangeOfKeysForInsert(0, 1500, OUR_PNI, { clearCreatedAt: true })
+    );
+    assert.strictEqual(1500, getCountOfKeys());
+    assert.strictEqual(90, getPragma());
+
+    updateToVersion(db, 91);
+
+    assert.strictEqual(91, getPragma());
+    assert.strictEqual(1500, getCountOfKeys());
+  });
+
+  it('leaves extra ACI keys alone, even if above 1000', () => {
+    assert.strictEqual(0, getCountOfKeys());
+    addPni();
+    insertData(db, 'preKeys', getRangeOfKeysForInsert(0, 1500, OUR_ACI));
+    assert.strictEqual(1500, getCountOfKeys());
+    assert.strictEqual(90, getPragma());
+
+    updateToVersion(db, 91);
+
+    assert.strictEqual(91, getPragma());
+    assert.strictEqual(1500, getCountOfKeys());
+  });
+
+  it('fixes ourServiceId generated column in preKeys table', () => {
+    updateToVersion(db, 91);
+    const id = 1;
+
+    insertData(db, 'preKeys', [
+      {
+        id,
+        json: {
+          ourServiceId: OUR_ACI,
+        },
+      },
+    ]);
+    assert.deepEqual(getTableData(db, 'preKeys'), [
+      {
+        id,
+        ourServiceId: OUR_ACI,
+        json: {
+          ourServiceId: OUR_ACI,
+        },
+      },
+    ]);
+  });
+
+  it('fixes ourServiceId generated column in kyberPreKeys table', () => {
+    updateToVersion(db, 91);
+    const id = 1;
+
+    insertData(db, 'kyberPreKeys', [
+      {
+        id,
+        json: {
+          ourServiceId: OUR_ACI,
+        },
+      },
+    ]);
+    assert.deepEqual(getTableData(db, 'kyberPreKeys'), [
+      {
+        id,
+        ourServiceId: OUR_ACI,
+        json: {
+          ourServiceId: OUR_ACI,
+        },
+      },
+    ]);
+  });
+
+  it('fixes ourServiceId generated column in signedPreKeys table', () => {
+    updateToVersion(db, 91);
+    const id = 1;
+
+    insertData(db, 'signedPreKeys', [
+      {
+        id,
+        json: {
+          ourServiceId: OUR_ACI,
+        },
+      },
+    ]);
+    assert.deepEqual(getTableData(db, 'signedPreKeys'), [
+      {
+        id,
+        ourServiceId: OUR_ACI,
+        json: {
+          ourServiceId: OUR_ACI,
+        },
+      },
+    ]);
+  });
+});

ts/textsecure/AccountManager.ts

@@ -489,6 +489,24 @@ export default class AccountManager extends EventTarget {
   async maybeUpdateKeys(serviceIdKind: ServiceIdKind): Promise<void> {
     const logId = `maybeUpdateKeys(${serviceIdKind})`;
     await this.queueTask(async () => {
+      const { storage } = window.textsecure;
+      let identityKey: KeyPairType;
+
+      try {
+        const ourServiceId = storage.user.getCheckedServiceId(serviceIdKind);
+        identityKey = this.getIdentityKeyOrThrow(ourServiceId);
+      } catch (error) {
+        if (serviceIdKind === ServiceIdKind.PNI) {
+          log.info(
+            `${logId}: Not enough information to update PNI keys`,
+            Errors.toLogFormat(error)
+          );
+          return;
+        }
+
+        throw error;
+      }
+
       const { count: preKeyCount, pqCount: kyberPreKeyCount } =
         await this.server.getMyKeyCounts(serviceIdKind);
 
@@ -544,9 +562,6 @@ export default class AccountManager extends EventTarget {
       }
       log.info(`${logId}: Uploading with ${keySummary.join(', ')}`);
 
-      const { storage } = window.textsecure;
-      const ourServiceId = storage.user.getCheckedServiceId(serviceIdKind);
-      const identityKey = this.getIdentityKeyOrThrow(ourServiceId);
       const toUpload = {
         identityKey: identityKey.pubKey,
         preKeys,