Migrate sessions table to BLOB column

Co-authored-by: Scott Nonnenberg <scott@signal.org>

ts/sql/Interface.ts

@@ -206,8 +206,7 @@ export type SessionType = {
   serviceId: ServiceIdString;
   conversationId: string;
   deviceId: number;
-  record: string;
-  version?: number;
+  record: Uint8Array;
 };
 export type SessionIdType = SessionType['id'];
 export type SignedPreKeyType = {
@@ -722,7 +721,6 @@ type WritableInterface = {
     sessions: Array<SessionType>;
     unprocessed: Array<UnprocessedType>;
   }): void;
-  bulkAddSessions: (array: Array<SessionType>) => void;
   removeSessionById: (id: SessionIdType) => number;
   removeSessionsByConversation: (conversationId: string) => void;
   removeSessionsByServiceId: (serviceId: ServiceIdString) => void;

ts/sql/Server.ts

@@ -407,7 +407,6 @@ export const DataWriter: ServerWritableInterface = {
   createOrUpdateSession,
   createOrUpdateSessions,
   commitDecryptResult,
-  bulkAddSessions,
   removeSessionById,
   removeSessionsByConversation,
   removeSessionsByServiceId,
@@ -1443,7 +1442,8 @@ function _getAllSentProtoMessageIds(db: ReadableDB): Array<SentMessageDBType> {
 
 const SESSIONS_TABLE = 'sessions';
 function createOrUpdateSession(db: WritableDB, data: SessionType): void {
-  const { id, conversationId, ourServiceId, serviceId } = data;
+  const { id, conversationId, ourServiceId, serviceId, deviceId, record } =
+    data;
   if (!id) {
     throw new Error(
       'createOrUpdateSession: Provided data did not have a truthy id'
@@ -1463,13 +1463,15 @@ function createOrUpdateSession(db: WritableDB, data: SessionType): void {
       conversationId,
       ourServiceId,
       serviceId,
-      json
+      deviceId,
+      record
     ) values (
       $id,
       $conversationId,
       $ourServiceId,
       $serviceId,
-      $json
+      $deviceId,
+      $record
     )
     `
   ).run({
@@ -1477,7 +1479,8 @@ function createOrUpdateSession(db: WritableDB, data: SessionType): void {
     conversationId,
     ourServiceId,
     serviceId,
-    json: objectToJSON(data),
+    deviceId,
+    record,
   });
 }
 
@@ -1519,9 +1522,6 @@ function commitDecryptResult(
   })();
 }
 
-function bulkAddSessions(db: WritableDB, array: Array<SessionType>): void {
-  return bulkAdd(db, SESSIONS_TABLE, array);
-}
 function removeSessionById(db: WritableDB, id: SessionIdType): number {
   return removeById(db, SESSIONS_TABLE, id);
 }
@@ -1555,7 +1555,7 @@ function removeAllSessions(db: WritableDB): number {
   return removeAllFromTable(db, SESSIONS_TABLE);
 }
 function getAllSessions(db: ReadableDB): Array<SessionType> {
-  return getAllFromTable(db, SESSIONS_TABLE);
+  return db.prepare('SELECT * FROM sessions').all();
 }
 // Conversations
 

ts/sql/migrations/1220-blob-sessions.ts

@@ -0,0 +1,218 @@
+// Copyright 2024 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+import assert from 'assert';
+import z from 'zod';
+import type { Database } from '@signalapp/better-sqlite3';
+import type { LoggerType } from '../../types/Logging';
+import * as Errors from '../../types/errors';
+import {
+  sessionRecordToProtobuf,
+  sessionStructureToBytes,
+} from '../../util/sessionTranslation';
+import { getOwn } from '../../util/getOwn';
+import { missingCaseError } from '../../util/missingCaseError';
+
+export const version = 1220;
+
+const identityKeyMapSchema = z.record(
+  z.string(),
+  z.object({
+    privKey: z.string().transform(x => Buffer.from(x, 'base64')),
+    pubKey: z.string().transform(x => Buffer.from(x, 'base64')),
+  })
+);
+
+const registrationIdMapSchema = z.record(z.string(), z.number());
+
+type PreviousSessionRowType = Readonly<{
+  id: string;
+  conversationId: string;
+  ourServiceId: string;
+  serviceId: string;
+  json: string;
+}>;
+
+const previousSessionJsonSchema = z.object({
+  id: z.string(),
+  ourServiceId: z.string(),
+  serviceId: z.string(),
+  conversationId: z.string(),
+  deviceId: z.number(),
+  record: z.string(),
+  version: z.literal(1).or(z.literal(2)),
+});
+
+type NextSessionRowType = Readonly<{
+  id: string;
+  conversationId: string;
+  ourServiceId: string;
+  serviceId: string;
+  deviceId: number;
+  record: Buffer;
+}>;
+
+function migrateSession(
+  row: PreviousSessionRowType,
+  identityKeyMap: z.infer<typeof identityKeyMapSchema>,
+  registrationIdMap: z.infer<typeof registrationIdMapSchema>,
+  logger: LoggerType
+): NextSessionRowType {
+  const { id, conversationId, ourServiceId, serviceId, json } = row;
+  const session = previousSessionJsonSchema.parse(JSON.parse(json));
+
+  assert.strictEqual(session.id, id, 'Invalid id');
+  assert.strictEqual(
+    session.conversationId,
+    conversationId,
+    'Invalid conversationId'
+  );
+  assert.strictEqual(
+    session.ourServiceId,
+    ourServiceId,
+    'Invalid ourServiceId,'
+  );
+  assert.strictEqual(session.serviceId, serviceId, 'Invalid serviceId');
+
+  // Previously migrated session
+  if (session.version === 2) {
+    return {
+      id,
+      conversationId,
+      ourServiceId,
+      serviceId,
+      deviceId: session.deviceId,
+      record: Buffer.from(session.record, 'base64'),
+    };
+  }
+
+  if (session.version === 1) {
+    const keyPair = getOwn(identityKeyMap, ourServiceId);
+    if (!keyPair) {
+      throw new Error('migrateSession: No identity key for ourself!');
+    }
+
+    const localRegistrationId = getOwn(registrationIdMap, ourServiceId);
+    if (localRegistrationId == null) {
+      throw new Error('_maybeMigrateSession: No registration id for ourself!');
+    }
+
+    const localUserData = {
+      identityKeyPublic: keyPair.pubKey,
+      registrationId: localRegistrationId,
+    };
+
+    logger.info(`migrateSession: Migrating session with id ${id}`);
+    const sessionProto = sessionRecordToProtobuf(
+      JSON.parse(session.record),
+      localUserData
+    );
+
+    return {
+      id,
+      conversationId,
+      ourServiceId,
+      serviceId,
+      deviceId: session.deviceId,
+      record: Buffer.from(sessionStructureToBytes(sessionProto)),
+    };
+  }
+
+  throw missingCaseError(session.version);
+}
+
+export function updateToSchemaVersion1220(
+  currentVersion: number,
+  db: Database,
+  logger: LoggerType
+): void {
+  if (currentVersion >= 1220) {
+    return;
+  }
+
+  db.transaction(() => {
+    db.exec(`
+      ALTER TABLE sessions
+        RENAME TO old_sessions;
+
+      CREATE TABLE sessions (
+        id TEXT NOT NULL PRIMARY KEY,
+        ourServiceId TEXT NOT NULL,
+        serviceId TEXT NOT NULL,
+        conversationId TEXT NOT NULL,
+        deviceId INTEGER NOT NULL,
+        record BLOB NOT NULL
+      ) STRICT;
+    `);
+
+    const getItem = db
+      .prepare(
+        `
+          SELECT json -> '$.value' FROM items WHERE id IS ?
+        `
+      )
+      .pluck();
+
+    const identityKeyMapJson = getItem.get('identityKeyMap');
+    const registrationIdMapJson = getItem.get('registrationIdMap');
+
+    // If we don't have private keys - the sessions cannot be used anyway
+    if (!identityKeyMapJson || !registrationIdMapJson) {
+      logger.info('updateToSchemaVersion1220: no identity/registration id');
+      db.exec('DROP TABLE old_sessions');
+      db.pragma('user_version = 1220');
+      return;
+    }
+
+    const identityKeyMap = identityKeyMapSchema.parse(
+      JSON.parse(identityKeyMapJson)
+    );
+    const registrationIdMap = registrationIdMapSchema.parse(
+      JSON.parse(registrationIdMapJson)
+    );
+
+    const getSessionsPage = db.prepare(
+      'DELETE FROM old_sessions RETURNING * LIMIT 1000'
+    );
+    const insertSession = db.prepare(`
+      INSERT INTO sessions
+      (id, ourServiceId, serviceId, conversationId, deviceId, record)
+      VALUES
+      ($id, $ourServiceId, $serviceId, $conversationId, $deviceId, $record)
+    `);
+
+    let migrated = 0;
+    let failed = 0;
+
+    // eslint-disable-next-line no-constant-condition
+    while (true) {
+      const rows: Array<PreviousSessionRowType> = getSessionsPage.all();
+      if (rows.length === 0) {
+        break;
+      }
+
+      for (const row of rows) {
+        try {
+          insertSession.run(
+            migrateSession(row, identityKeyMap, registrationIdMap, logger)
+          );
+          migrated += 1;
+        } catch (error) {
+          failed += 1;
+          logger.error(
+            'updateToSchemaVersion1220: failed to migrate session',
+            Errors.toLogFormat(error)
+          );
+        }
+      }
+    }
+
+    logger.info(
+      `updateToSchemaVersion1220: migrated ${migrated} sessions, ` +
+        `${failed} failed`
+    );
+
+    db.exec('DROP TABLE old_sessions');
+    db.pragma('user_version = 1220');
+  })();
+  logger.info('updateToSchemaVersion1220: success!');
+}

ts/sql/migrations/index.ts

@@ -97,10 +97,11 @@ import { updateToSchemaVersion1170 } from './1170-update-call-history-unread-ind
 import { updateToSchemaVersion1180 } from './1180-add-attachment-download-source';
 import { updateToSchemaVersion1190 } from './1190-call-links-storage';
 import { updateToSchemaVersion1200 } from './1200-attachment-download-source-index';
+import { updateToSchemaVersion1210 } from './1210-call-history-started-id';
 import {
-  updateToSchemaVersion1210,
+  updateToSchemaVersion1220,
   version as MAX_VERSION,
-} from './1210-call-history-started-id';
+} from './1220-blob-sessions';
 
 function updateToSchemaVersion1(
   currentVersion: number,
@@ -2067,6 +2068,7 @@ export const SCHEMA_VERSIONS = [
 
   updateToSchemaVersion1200,
   updateToSchemaVersion1210,
+  updateToSchemaVersion1220,
 ];
 
 export class DBVersionFromFutureError extends Error {