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/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,
+        },
+      },
+    ]);
+  });
+});