ChatApp/packages/shared/src/chat/userKeyMigration.test.ts

import { describe, expect, it, beforeAll } from 'vitest';

import { setCryptoBackend, getCryptoBackend } from '../crypto/backend';
import { makeWasmTestBackend } from '../crypto/testBackend';
import { encryptFor } from '../crypto/box';
import { migrateOwnLegacyBundles } from './userKeyMigration';

beforeAll(async () => { setCryptoBackend(await makeWasmTestBackend()); });

describe('migrateOwnLegacyBundles', () => {
  it('re-wraps legacy device bundles to user recipients and skips non-own rows', async () => {
    const backend = getCryptoBackend();
    const senderKp = backend.generateKeyPair();
    const oldDeviceKp = backend.generateKeyPair();
    const newUserKp = backend.generateKeyPair();
    const otherDeviceKp = backend.generateKeyPair();

    const convKey = backend.randomBytes(backend.secretboxKeyLength);
    const wrappedForOldDevice = await encryptFor(convKey, oldDeviceKp.publicKey, senderKp.privateKey);
    const wrappedForOther     = await encryptFor(convKey, otherDeviceKp.publicKey, senderKp.privateKey);

    const calls: { name: string; params: unknown }[] = [];
    const stubClient = {
      from: (table: string) => {
        if (table === 'conversation_keys') {
          // Mirrors the new chain: .select(...).is('recipient_user_id', null).not('recipient_device_id', 'is', null)
          return {
            select: () => ({
              is: () => ({
                not: () => Promise.resolve({
                  data: [
                    {
                      conversation_id: 'conv-1',
                      key_version: 1,
                      recipient_device_id: 'dev-old',
                      sender_device_id: 'dev-sender',
                      sender_user_id: 'sender-user',
                      encrypted_key: '\\x' + Buffer.from(wrappedForOldDevice.ciphertext).toString('hex'),
                      nonce:         '\\x' + Buffer.from(wrappedForOldDevice.nonce).toString('hex'),
                    },
                    {
                      conversation_id: 'conv-2',
                      key_version: 1,
                      recipient_device_id: 'dev-other',
                      sender_device_id: 'dev-sender',
                      sender_user_id: 'sender-user',
                      encrypted_key: '\\x' + Buffer.from(wrappedForOther.ciphertext).toString('hex'),
                      nonce:         '\\x' + Buffer.from(wrappedForOther.nonce).toString('hex'),
                    },
                  ],
                  error: null,
                }),
              }),
            }),
          };
        }
        if (table === 'devices') {
          return {
            select: () => ({
              in: () => Promise.resolve({
                data: [{ id: 'dev-sender', user_id: 'sender-user', public_key: '\\x' + Buffer.from(senderKp.publicKey).toString('hex') }],
                error: null,
              }),
            }),
          };
        }
        throw new Error('unexpected table ' + table);
      },
      rpc: (name: string, params: unknown) => {
        calls.push({ name, params });
        return Promise.resolve({ data: 1, error: null });
      },
    } as unknown as Parameters<typeof migrateOwnLegacyBundles>[0]['client'];

    const result = await migrateOwnLegacyBundles({
      client: stubClient,
      ownUserId: 'me-user',
      ownNewPublicKey: newUserKp.publicKey,
      ownNewPrivateKey: newUserKp.privateKey,
      ownLegacyDeviceIds: ['dev-old'],
      ownLegacyDevicePrivateKeys: { 'dev-old': oldDeviceKp.privateKey },
    });

    expect(result.migratedConversations).toBe(1);
    expect(calls).toHaveLength(1);
    expect(calls[0]!.name).toBe('migrate_user_key_recipients');
    const params = calls[0]!.params as { p_conv_id: string; p_user_id: string };
    expect(params.p_conv_id).toBe('conv-1');
    expect(params.p_user_id).toBe('me-user');
  });
});