Extra tests for SignalProtocolStore migration

ts/sql/Server.ts

@@ -2125,6 +2125,36 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
     )
     .pluck();
 
+  const getConversationStats = db.prepare<Query>(
+    `
+      SELECT uuid, e164, active_at
+      FROM
+        conversations
+      WHERE
+        id = $conversationId
+      `
+  );
+
+  const compareConvoRecency = (a: string, b: string): number => {
+    const aStats = getConversationStats.get({ conversationId: a });
+    const bStats = getConversationStats.get({ conversationId: b });
+
+    const isAComplete = Boolean(aStats?.uuid && aStats?.e164);
+    const isBComplete = Boolean(bStats?.uuid && bStats?.e164);
+
+    if (!isAComplete && !isBComplete) {
+      return 0;
+    }
+    if (!isAComplete) {
+      return -1;
+    }
+    if (!isBComplete) {
+      return 1;
+    }
+
+    return aStats.active_at - bStats.active_at;
+  };
+
   const clearSessionsAndKeys = () => {
     // ts/background.ts will ask user to relink so all that matters here is
     // to maintain an invariant:
@@ -2196,20 +2226,7 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
 
   const prefixKeys = (ourUuid: string) => {
     for (const table of ['signedPreKeys', 'preKeys']) {
-      // Add numeric `keyId` field to keys
+      // Update id to include suffix, add `ourUuid` and `keyId` fields.
-      db.prepare<EmptyQuery>(
-        `
-        UPDATE ${table}
-        SET
-          json = json_insert(
-            json,
-            '$.keyId',
-            json_extract(json, '$.id')
-          )
-        `
-      ).run();
-
-      // Update id to include suffix and add `ourUuid` field
       db.prepare<Query>(
         `
         UPDATE ${table}
@@ -2219,17 +2236,26 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
             json,
             '$.id',
             $ourUuid || ':' || json_extract(json, '$.id'),
+            '$.keyId',
+            json_extract(json, '$.id'),
             '$.ourUuid',
             $ourUuid
           )
         `
       ).run({ ourUuid });
     }
+  };
 
+  const updateSenderKeys = (ourUuid: string) => {
     const senderKeys: ReadonlyArray<{
       id: string;
       senderId: string;
-    }> = db.prepare<EmptyQuery>('SELECT id, senderId FROM senderKeys').all();
+      lastUpdatedDate: number;
+    }> = db
+      .prepare<EmptyQuery>(
+        'SELECT id, senderId, lastUpdatedDate FROM senderKeys'
+      )
+      .all();
 
     console.log(`Updating ${senderKeys.length} sender keys`);
 
@@ -2248,9 +2274,18 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
       'DELETE FROM senderKeys WHERE id = $id'
     );
 
+    const pastKeys = new Map<
+      string,
+      {
+        conversationId: string;
+        lastUpdatedDate: number;
+      }
+    >();
+
     let updated = 0;
     let deleted = 0;
-    for (const { id, senderId } of senderKeys) {
+    let skipped = 0;
+    for (const { id, senderId, lastUpdatedDate } of senderKeys) {
       const [conversationId] = Helpers.unencodeNumber(senderId);
       const uuid = getConversationUuid.get({ conversationId });
 
@@ -2260,17 +2295,40 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
         continue;
       }
 
-      updated += 1;
+      const newId = `${ourUuid}:${id.replace(conversationId, uuid)}`;
+
+      const existing = pastKeys.get(newId);
+
+      // We are going to delete on of the keys anyway
+      if (existing) {
+        skipped += 1;
+      } else {
+        updated += 1;
+      }
+
+      const isOlder =
+        existing &&
+        (lastUpdatedDate < existing.lastUpdatedDate ||
+          compareConvoRecency(conversationId, existing.conversationId) < 0);
+      if (isOlder) {
+        deleteSenderKey.run({ id });
+        continue;
+      } else if (existing) {
+        deleteSenderKey.run({ id: newId });
+      }
+
+      pastKeys.set(newId, { conversationId, lastUpdatedDate });
+
       updateSenderKey.run({
         id,
-        newId: `${ourUuid}:${id.replace(conversationId, uuid)}`,
+        newId,
         newSenderId: `${senderId.replace(conversationId, uuid)}`,
       });
     }
 
     console.log(
       `Updated ${senderKeys.length} sender keys: ` +
-        `updated: ${updated}, deleted: ${deleted}`
+        `updated: ${updated}, deleted: ${deleted}, skipped: ${skipped}`
     );
   };
 
@@ -2310,8 +2368,16 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
       'DELETE FROM sessions WHERE id = $id'
     );
 
+    const pastSessions = new Map<
+      string,
+      {
+        conversationId: string;
+      }
+    >();
+
     let updated = 0;
     let deleted = 0;
+    let skipped = 0;
     for (const { id, conversationId } of allSessions) {
       const uuid = getConversationUuid.get({ conversationId });
       if (!uuid) {
@@ -2322,7 +2388,27 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
 
       const newId = `${ourUuid}:${id.replace(conversationId, uuid)}`;
 
-      updated += 1;
+      const existing = pastSessions.get(newId);
+
+      // We are going to delete on of the keys anyway
+      if (existing) {
+        skipped += 1;
+      } else {
+        updated += 1;
+      }
+
+      const isOlder =
+        existing &&
+        compareConvoRecency(conversationId, existing.conversationId) < 0;
+      if (isOlder) {
+        deleteSession.run({ id });
+        continue;
+      } else if (existing) {
+        deleteSession.run({ id: newId });
+      }
+
+      pastSessions.set(newId, { conversationId });
+
       updateSession.run({
         id,
         newId,
@@ -2333,7 +2419,7 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
 
     console.log(
       `Updated ${allSessions.length} sessions: ` +
-        `updated: ${updated}, deleted: ${deleted}`
+        `updated: ${updated}, deleted: ${deleted}, skipped: ${skipped}`
     );
   };
 
@@ -2429,6 +2515,8 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
 
     prefixKeys(ourUuid);
 
+    updateSenderKeys(ourUuid);
+
     updateSessions(ourUuid);
 
     moveIdentityKeyToMap(ourUuid);
@@ -2440,7 +2528,7 @@ function updateToSchemaVersion41(currentVersion: number, db: Database) {
   console.log('updateToSchemaVersion41: success!');
 }
 
-const SCHEMA_VERSIONS = [
+export const SCHEMA_VERSIONS = [
   updateToSchemaVersion1,
   updateToSchemaVersion2,
   updateToSchemaVersion3,
@@ -2484,7 +2572,7 @@ const SCHEMA_VERSIONS = [
   updateToSchemaVersion41,
 ];
 
-function updateSchema(db: Database): void {
+export function updateSchema(db: Database) {
   const sqliteVersion = getSQLiteVersion(db);
   const sqlcipherVersion = getSQLCipherVersion(db);
   const userVersion = getUserVersion(db);
@@ -2502,7 +2590,8 @@ function updateSchema(db: Database): void {
 
   if (userVersion > maxUserVersion) {
     throw new Error(
-      `SQL: User version is ${userVersion} but the expected maximum version is ${maxUserVersion}. Did you try to start an old version of Signal?`
+      `SQL: User version is ${userVersion} but the expected maximum version ` +
+        `is ${maxUserVersion}. Did you try to start an old version of Signal?`
     );
   }
 

ts/test-node/sql_migrations_test.ts

@@ -0,0 +1,522 @@
+// Copyright 2021 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import SQL, { Database } from 'better-sqlite3';
+import { v4 as generateGuid } from 'uuid';
+
+import { SCHEMA_VERSIONS } from '../sql/Server';
+
+const THEIR_UUID = generateGuid();
+const THEIR_CONVO = generateGuid();
+const ANOTHER_CONVO = generateGuid();
+const THIRD_CONVO = generateGuid();
+const OUR_UUID = generateGuid();
+
+describe('SQL migrations test', () => {
+  let db: Database;
+
+  const updateToVersion = (version: number) => {
+    const startVersion = db.pragma('user_version', { simple: true });
+
+    for (const run of SCHEMA_VERSIONS) {
+      run(startVersion, db);
+
+      const currentVersion = db.pragma('user_version', { simple: true });
+
+      if (currentVersion === version) {
+        return;
+      }
+    }
+
+    throw new Error(`Migration to ${version} not found`);
+  };
+
+  const addOurUuid = () => {
+    const value = {
+      id: 'uuid_id',
+      value: `${OUR_UUID}.1`,
+    };
+    db.exec(
+      `
+      INSERT INTO items (id, json) VALUES
+        ('uuid_id', '${JSON.stringify(value)}');
+      `
+    );
+  };
+
+  const parseItems = (
+    items: ReadonlyArray<{ json: string }>
+  ): Array<unknown> => {
+    return items.map(item => {
+      return {
+        ...item,
+        json: JSON.parse(item.json),
+      };
+    });
+  };
+
+  const insertSession = (
+    conversationId: string,
+    deviceId: number,
+    data: Record<string, unknown> = {}
+  ): void => {
+    const id = `${conversationId}.${deviceId}`;
+    db.prepare(
+      `
+        INSERT INTO sessions (id, conversationId, json)
+        VALUES ($id, $conversationId, $json)
+      `
+    ).run({
+      id,
+      conversationId,
+      json: JSON.stringify({
+        ...data,
+        id,
+        conversationId,
+      }),
+    });
+  };
+
+  beforeEach(() => {
+    db = new SQL(':memory:');
+  });
+
+  afterEach(() => {
+    db.close();
+  });
+
+  describe('updateToSchemaVersion41', () => {
+    it('clears sessions and keys if UUID is not available', () => {
+      updateToVersion(40);
+
+      db.exec(
+        `
+        INSERT INTO senderKeys
+          (id, senderId, distributionId, data, lastUpdatedDate)
+          VALUES
+          ('1', '1', '1', '1', 1);
+        INSERT INTO sessions (id, conversationId, json) VALUES
+          ('1', '1', '{}');
+        INSERT INTO signedPreKeys (id, json) VALUES
+          ('1', '{}');
+        INSERT INTO preKeys (id, json) VALUES
+          ('1', '{}');
+        INSERT INTO items (id, json) VALUES
+          ('identityKey', '{}'),
+          ('registrationId', '{}');
+        `
+      );
+
+      const senderKeyCount = db
+        .prepare('SELECT COUNT(*) FROM senderKeys')
+        .pluck();
+      const sessionCount = db.prepare('SELECT COUNT(*) FROM sessions').pluck();
+      const signedPreKeyCount = db
+        .prepare('SELECT COUNT(*) FROM signedPreKeys')
+        .pluck();
+      const preKeyCount = db.prepare('SELECT COUNT(*) FROM preKeys').pluck();
+      const itemCount = db.prepare('SELECT COUNT(*) FROM items').pluck();
+
+      assert.strictEqual(senderKeyCount.get(), 1);
+      assert.strictEqual(sessionCount.get(), 1);
+      assert.strictEqual(signedPreKeyCount.get(), 1);
+      assert.strictEqual(preKeyCount.get(), 1);
+      assert.strictEqual(itemCount.get(), 2);
+
+      updateToVersion(41);
+
+      assert.strictEqual(senderKeyCount.get(), 0);
+      assert.strictEqual(sessionCount.get(), 0);
+      assert.strictEqual(signedPreKeyCount.get(), 0);
+      assert.strictEqual(preKeyCount.get(), 0);
+      assert.strictEqual(itemCount.get(), 0);
+    });
+
+    it('adds prefix to preKeys/signedPreKeys', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      const signedKeyItem = { id: 1 };
+      const preKeyItem = { id: 2 };
+
+      db.exec(
+        `
+        INSERT INTO signedPreKeys (id, json) VALUES
+          (1, '${JSON.stringify(signedKeyItem)}');
+        INSERT INTO preKeys (id, json) VALUES
+          (2, '${JSON.stringify(preKeyItem)}');
+        `
+      );
+
+      updateToVersion(41);
+
+      assert.deepStrictEqual(
+        parseItems(db.prepare('SELECT * FROM signedPreKeys').all()),
+        [
+          {
+            id: `${OUR_UUID}:1`,
+            json: {
+              id: `${OUR_UUID}:1`,
+              keyId: 1,
+              ourUuid: OUR_UUID,
+            },
+          },
+        ]
+      );
+      assert.deepStrictEqual(
+        parseItems(db.prepare('SELECT * FROM preKeys').all()),
+        [
+          {
+            id: `${OUR_UUID}:2`,
+            json: {
+              id: `${OUR_UUID}:2`,
+              keyId: 2,
+              ourUuid: OUR_UUID,
+            },
+          },
+        ]
+      );
+    });
+
+    it('migrates senderKeys', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      db.exec(
+        `
+        INSERT INTO conversations (id, uuid) VALUES
+          ('${THEIR_CONVO}', '${THEIR_UUID}');
+
+        INSERT INTO senderKeys
+          (id, senderId, distributionId, data, lastUpdatedDate)
+          VALUES
+          ('${THEIR_CONVO}.1--234', '${THEIR_CONVO}.1', '234', '1', 1);
+        `
+      );
+
+      updateToVersion(41);
+
+      assert.deepStrictEqual(db.prepare('SELECT * FROM senderKeys').all(), [
+        {
+          id: `${OUR_UUID}:${THEIR_UUID}.1--234`,
+          distributionId: '234',
+          data: '1',
+          lastUpdatedDate: 1,
+          senderId: `${THEIR_UUID}.1`,
+        },
+      ]);
+    });
+
+    it('removes senderKeys that do not have conversation uuid', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      db.exec(
+        `
+        INSERT INTO conversations (id) VALUES
+          ('${THEIR_CONVO}');
+
+        INSERT INTO senderKeys
+          (id, senderId, distributionId, data, lastUpdatedDate)
+          VALUES
+          ('${THEIR_CONVO}.1--234', '${THEIR_CONVO}.1', '234', '1', 1),
+          ('${ANOTHER_CONVO}.1--234', '${ANOTHER_CONVO}.1', '234', '1', 1);
+        `
+      );
+
+      updateToVersion(41);
+
+      assert.strictEqual(
+        db.prepare('SELECT COUNT(*) FROM senderKeys').pluck().get(),
+        0
+      );
+    });
+
+    it('correctly merges senderKeys for conflicting conversations', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      const fullA = generateGuid();
+      const fullB = generateGuid();
+      const fullC = generateGuid();
+      const partial = generateGuid();
+
+      // When merging two keys for different conversations with the same uuid
+      // only the most recent key would be kept in the database. We prefer keys
+      // with either:
+      //
+      // 1. more recent lastUpdatedDate column
+      // 2. conversation with both e164 and uuid
+      // 3. conversation with more recent active_at
+      db.exec(
+        `
+        INSERT INTO conversations (id, uuid, e164, active_at) VALUES
+          ('${fullA}', '${THEIR_UUID}', '+12125555555', 1),
+          ('${fullB}', '${THEIR_UUID}', '+12125555555', 2),
+          ('${fullC}', '${THEIR_UUID}', '+12125555555', 3),
+          ('${partial}', '${THEIR_UUID}', NULL, 3);
+
+        INSERT INTO senderKeys
+          (id, senderId, distributionId, data, lastUpdatedDate)
+        VALUES
+          ('${fullA}.1--234', '${fullA}.1', 'fullA', '1', 1),
+          ('${fullC}.1--234', '${fullC}.1', 'fullC', '2', 2),
+          ('${fullB}.1--234', '${fullB}.1', 'fullB', '3', 2),
+          ('${partial}.1--234', '${partial}.1', 'partial', '4', 2);
+        `
+      );
+
+      updateToVersion(41);
+
+      assert.deepStrictEqual(db.prepare('SELECT * FROM senderKeys').all(), [
+        {
+          id: `${OUR_UUID}:${THEIR_UUID}.1--234`,
+          senderId: `${THEIR_UUID}.1`,
+          distributionId: 'fullC',
+          lastUpdatedDate: 2,
+          data: '2',
+        },
+      ]);
+    });
+
+    it('migrates sessions', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      db.exec(
+        `
+        INSERT INTO conversations (id, uuid) VALUES
+          ('${THEIR_CONVO}', '${THEIR_UUID}');
+        `
+      );
+
+      insertSession(THEIR_CONVO, 1);
+
+      updateToVersion(41);
+
+      assert.deepStrictEqual(
+        parseItems(db.prepare('SELECT * FROM sessions').all()),
+        [
+          {
+            conversationId: THEIR_CONVO,
+            id: `${OUR_UUID}:${THEIR_UUID}.1`,
+            uuid: THEIR_UUID,
+            ourUuid: OUR_UUID,
+            json: {
+              id: `${OUR_UUID}:${THEIR_UUID}.1`,
+              conversationId: THEIR_CONVO,
+              uuid: THEIR_UUID,
+              ourUuid: OUR_UUID,
+            },
+          },
+        ]
+      );
+    });
+
+    it('removes sessions that do not have conversation id', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      insertSession(THEIR_CONVO, 1);
+
+      updateToVersion(41);
+
+      assert.strictEqual(
+        db.prepare('SELECT COUNT(*) FROM sessions').pluck().get(),
+        0
+      );
+    });
+
+    it('removes sessions that do not have conversation uuid', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      db.exec(
+        `
+        INSERT INTO conversations (id) VALUES ('${THEIR_CONVO}');
+        `
+      );
+
+      insertSession(THEIR_CONVO, 1);
+
+      updateToVersion(41);
+
+      assert.strictEqual(
+        db.prepare('SELECT COUNT(*) FROM sessions').pluck().get(),
+        0
+      );
+    });
+
+    it('correctly merges sessions for conflicting conversations', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      const fullA = generateGuid();
+      const fullB = generateGuid();
+      const partial = generateGuid();
+
+      // Similar merging logic to senderkeys above. We prefer sessions with
+      // either:
+      //
+      // 1. conversation with both e164 and uuid
+      // 2. conversation with more recent active_at
+      db.exec(
+        `
+        INSERT INTO conversations (id, uuid, e164, active_at) VALUES
+          ('${fullA}', '${THEIR_UUID}', '+12125555555', 1),
+          ('${fullB}', '${THEIR_UUID}', '+12125555555', 2),
+          ('${partial}', '${THEIR_UUID}', NULL, 3);
+        `
+      );
+
+      insertSession(fullA, 1, { name: 'A' });
+      insertSession(fullB, 1, { name: 'B' });
+      insertSession(partial, 1, { name: 'C' });
+
+      updateToVersion(41);
+
+      assert.deepStrictEqual(
+        parseItems(db.prepare('SELECT * FROM sessions').all()),
+        [
+          {
+            id: `${OUR_UUID}:${THEIR_UUID}.1`,
+            conversationId: fullB,
+            ourUuid: OUR_UUID,
+            uuid: THEIR_UUID,
+            json: {
+              id: `${OUR_UUID}:${THEIR_UUID}.1`,
+              conversationId: fullB,
+              ourUuid: OUR_UUID,
+              uuid: THEIR_UUID,
+              name: 'B',
+            },
+          },
+        ]
+      );
+    });
+
+    it('moves identity key and registration id into a map', () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      const items = [
+        { id: 'identityKey', value: 'secret' },
+        { id: 'registrationId', value: 42 },
+      ];
+
+      for (const item of items) {
+        db.prepare(
+          `
+          INSERT INTO items (id, json) VALUES ($id, $json);
+          `
+        ).run({
+          id: item.id,
+          json: JSON.stringify(item),
+        });
+      }
+
+      updateToVersion(41);
+
+      assert.deepStrictEqual(
+        parseItems(db.prepare('SELECT * FROM items ORDER BY id').all()),
+        [
+          {
+            id: 'identityKeyMap',
+            json: {
+              id: 'identityKeyMap',
+              value: { [OUR_UUID]: 'secret' },
+            },
+          },
+          {
+            id: 'registrationIdMap',
+            json: {
+              id: 'registrationIdMap',
+              value: { [OUR_UUID]: 42 },
+            },
+          },
+          {
+            id: 'uuid_id',
+            json: {
+              id: 'uuid_id',
+              value: `${OUR_UUID}.1`,
+            },
+          },
+        ]
+      );
+    });
+
+    it("migrates other users' identity keys", () => {
+      updateToVersion(40);
+
+      addOurUuid();
+
+      db.exec(
+        `
+        INSERT INTO conversations (id, uuid) VALUES
+          ('${THEIR_CONVO}', '${THEIR_UUID}'),
+          ('${ANOTHER_CONVO}', NULL);
+        `
+      );
+
+      const identityKeys = [
+        { id: THEIR_CONVO },
+        { id: ANOTHER_CONVO },
+        { id: THIRD_CONVO },
+      ];
+      for (const key of identityKeys) {
+        db.prepare(
+          `
+            INSERT INTO identityKeys (id, json) VALUES ($id, $json);
+          `
+        ).run({
+          id: key.id,
+          json: JSON.stringify(key),
+        });
+      }
+
+      updateToVersion(41);
+
+      assert.deepStrictEqual(
+        parseItems(db.prepare('SELECT * FROM identityKeys ORDER BY id').all()),
+        [
+          {
+            id: THEIR_UUID,
+            json: {
+              id: THEIR_UUID,
+            },
+          },
+          {
+            id: `conversation:${ANOTHER_CONVO}`,
+            json: {
+              id: `conversation:${ANOTHER_CONVO}`,
+            },
+          },
+          {
+            id: `conversation:${THIRD_CONVO}`,
+            json: {
+              id: `conversation:${THIRD_CONVO}`,
+            },
+          },
+        ].sort((a, b) => {
+          if (a.id === b.id) {
+            return 0;
+          }
+          if (a.id < b.id) {
+            return -1;
+          }
+          return 1;
+        })
+      );
+    });
+  });
+});