Use untagged pnis in storage service

ts/sql/migrations/960-untag-pni.ts

@@ -0,0 +1,265 @@
+// 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 type {
+  ServiceIdString,
+  AciString,
+  PniString,
+} from '../../types/ServiceId';
+import { normalizePni } from '../../types/ServiceId';
+import { normalizeAci } from '../../util/normalizeAci';
+import type { JSONWithUnknownFields } from '../../types/Util';
+
+export const version = 960;
+
+export function updateToSchemaVersion960(
+  currentVersion: number,
+  db: Database,
+  logger: LoggerType
+): void {
+  if (currentVersion >= 960) {
+    return;
+  }
+
+  db.transaction(() => {
+    const ourServiceIds = migratePni(db, logger);
+    if (!ourServiceIds) {
+      logger.info('updateToSchemaVersion960: not running, pni is normalized');
+      return;
+    }
+
+    // Migrate JSON fields
+    db.prepare(
+      `
+        UPDATE conversations
+        SET json = json_set(json, '$.pni', $pni)
+        WHERE serviceId IS $aci
+      `
+    ).run({
+      aci: ourServiceIds.aci,
+      pni: ourServiceIds.pni,
+    });
+
+    migratePreKeys(db, 'preKeys', ourServiceIds, logger);
+    migratePreKeys(db, 'signedPreKeys', ourServiceIds, logger);
+    migratePreKeys(db, 'kyberPreKeys', ourServiceIds, logger);
+
+    db.pragma('user_version = 960');
+  })();
+
+  logger.info('updateToSchemaVersion960: success!');
+}
+
+//
+// migratePni checks if `pni` needs normalization:
+//
+// * If yes - return legacy and updated pni
+// * It no - return undefined
+//
+
+type OurServiceIds = Readonly<{
+  aci: AciString;
+  legacyPni: string;
+  pni: PniString;
+}>;
+
+function migratePni(
+  db: Database,
+  logger: LoggerType
+): OurServiceIds | undefined {
+  // Get our ACI and PNI
+  const uuidIdJson = db
+    .prepare(
+      `
+        SELECT json
+        FROM items
+        WHERE id IS 'uuid_id'
+      `
+    )
+    .pluck()
+    .get();
+  const pniJson = db
+    .prepare(
+      `
+        SELECT json
+        FROM items
+        WHERE id IS 'pni'
+      `
+    )
+    .pluck()
+    .get();
+
+  let aci: string | undefined;
+  try {
+    [aci] = JSON.parse(uuidIdJson).value.split('.', 2);
+  } catch (error) {
+    if (uuidIdJson) {
+      logger.warn(
+        'updateToSchemaVersion960: failed to parse uuid_id item',
+        error
+      );
+    } else {
+      logger.info('updateToSchemaVersion960: Our ACI not found');
+    }
+  }
+  if (!aci) {
+    return undefined;
+  }
+
+  let legacyPni: string | undefined;
+  try {
+    legacyPni = JSON.parse(pniJson).value;
+  } catch (error) {
+    if (pniJson) {
+      logger.warn('updateToSchemaVersion960: failed to parse pni item', error);
+    } else {
+      logger.info('updateToSchemaVersion960: Our PNI not found');
+    }
+  }
+  if (!legacyPni) {
+    return undefined;
+  }
+
+  const pni = prefixPni(legacyPni, 'pni', logger);
+  if (!pni || pni === legacyPni) {
+    return undefined;
+  }
+
+  const maps: Array<{ id: string; json: string }> = db
+    .prepare(
+      `
+      SELECT id, json
+      FROM items
+      WHERE id IN ('identityKeyMap', 'registrationIdMap');
+    `
+    )
+    .all();
+
+  const updateStmt = db.prepare(
+    'UPDATE items SET json = $json WHERE id IS $id'
+  );
+
+  updateStmt.run({
+    id: 'pni',
+    json: JSON.stringify({ id: 'pni', value: pni }),
+  });
+
+  for (const { id, json } of maps) {
+    try {
+      const data: { id: string; value: Record<string, unknown> } =
+        JSON.parse(json);
+
+      const pniValue = data.value[legacyPni];
+      if (pniValue) {
+        delete data.value[legacyPni];
+        data.value[pni] = pniValue;
+      }
+
+      updateStmt.run({ id, json: JSON.stringify(data) });
+    } catch (error) {
+      logger.warn(
+        `updateToSchemaVersion960: failed to parse ${id} item`,
+        error
+      );
+    }
+  }
+  return {
+    aci: normalizeAci(aci, 'uuid_id', logger),
+    pni,
+    legacyPni,
+  };
+}
+
+// migratePreKeys does the following:
+//
+// 1. Update `ourServiceId` to prefixed PNI
+// 2. Update `id` to use new `ourServiceId` value
+//    (the schema is `${ourServiceId}:${keyId}`)
+//
+
+function migratePreKeys(
+  db: Database,
+  table: string,
+  { legacyPni, pni }: OurServiceIds,
+  logger: LoggerType
+): void {
+  const preKeys = db
+    .prepare(`SELECT id, json FROM ${table} WHERE ourServiceId IS $legacyPni`)
+    .all({ legacyPni });
+
+  const updateStmt = db.prepare(`
+    UPDATE ${table}
+    SET id = $newId, json = $newJson
+    WHERE id = $id
+  `);
+
+  logger.info(`updateToSchemaVersion960: updating ${preKeys.length} ${table}`);
+  for (const { id, json } of preKeys) {
+    const match = id.match(/^(.*):(.*)$/);
+    if (!match) {
+      logger.warn(`updateToSchemaVersion960: invalid ${table} id ${id}`);
+      continue;
+    }
+
+    let legacyData: JSONWithUnknownFields<Record<string, unknown>>;
+    try {
+      legacyData = JSON.parse(json);
+    } catch (error) {
+      logger.warn(
+        `updateToSchemaVersion960: failed to parse ${table} ${id}`,
+        error
+      );
+      continue;
+    }
+
+    const [, ourServiceId, keyId] = match;
+    if (ourServiceId !== legacyPni) {
+      logger.warn(
+        'updateToSchemaVersion960: unexpected ourServiceId',
+        ourServiceId,
+        legacyPni
+      );
+      continue;
+    }
+
+    const newId = `${pni}:${keyId}`;
+
+    const newData: JSONWithUnknownFields<{
+      id: string;
+      ourServiceId: ServiceIdString;
+    }> = {
+      ...legacyData,
+      id: newId,
+      ourServiceId: pni,
+    };
+
+    updateStmt.run({
+      id,
+      newId,
+      newJson: JSON.stringify(newData),
+    });
+  }
+}
+
+//
+// Various utility methods below.
+//
+
+function prefixPni(
+  legacyPni: string | null | undefined,
+  context: string,
+  logger: LoggerType
+): PniString | undefined {
+  if (legacyPni == null) {
+    return undefined;
+  }
+
+  if (legacyPni.toLowerCase().startsWith('pni:')) {
+    return normalizePni(legacyPni, context, logger);
+  }
+
+  return normalizePni(`PNI:${legacyPni}`, context, logger);
+}

ts/sql/migrations/index.ts

@@ -70,10 +70,11 @@ import updateToSchemaVersion91 from './91-clean-keys';
 import { updateToSchemaVersion920 } from './920-clean-more-keys';
 import { updateToSchemaVersion930 } from './930-fts5-secure-delete';
 import { updateToSchemaVersion940 } from './940-fts5-revert';
+import { updateToSchemaVersion950 } from './950-fts5-secure-delete';
 import {
   version as MAX_VERSION,
-  updateToSchemaVersion950,
-} from './950-fts5-secure-delete';
+  updateToSchemaVersion960,
+} from './960-untag-pni';
 
 function updateToSchemaVersion1(
   currentVersion: number,
@@ -2011,6 +2012,7 @@ export const SCHEMA_VERSIONS = [
   updateToSchemaVersion930,
   updateToSchemaVersion940,
   updateToSchemaVersion950,
+  updateToSchemaVersion960,
 ];
 
 export class DBVersionFromFutureError extends Error {