Fix first encrypted message delay and key query ordering

src/app/crypto/engineCrypto/EngineCrypto.ts

@@ -45,7 +45,7 @@
   type EngineVerificationState,
 } from '../verification/state';
 import { engineInvoke, type EngineIdentity } from '../olmMachine/engineInvoke';
-import { sendOutgoingRequest, type OutgoingRequest } from './outgoing';
+import { RequestType, sendOutgoingRequest, type OutgoingRequest } from './outgoing';
 import { createCoalescedRunner } from './coalescedRunner';
 import type {
   BackupDecryptor,
@@ -133,7 +133,7 @@
 
   const entries = Object.entries(value as Record<string, unknown>)
     .filter(([, item]) => item !== undefined)
     .sort(([a], [b]) => (a < b ? -1 : 1));
   return `{${entries.map(([key, item]) => `${JSON.stringify(key)}:${canonicalJson(item)}`).join(',')}}`;
 };
 
@@ -358,6 +358,8 @@
 
   #claimChain: Promise<unknown> = Promise.resolve();
 
+  #keyQueryChain: Promise<unknown> = Promise.resolve();
+
   readonly #encryptionChains = new Map<string, Promise<unknown>>();
 
   readonly #backupUpload = createCoalescedRunner(
@@ -754,13 +756,20 @@
 
   async #sendTracked(request: unknown): Promise<void> {
     if (!isOutgoingRequest(request)) return;
-    const response = await sendOutgoingRequest(this.#mx, request);
-    if (typeof request.id !== 'string') return;
-    await this.#call('markRequestAsSent', {
-      requestId: request.id,
-      requestType: request.type,
-      response,
-    });
+    const send = async () => {
+      const response = await sendOutgoingRequest(this.#mx, request);
+      if (typeof request.id !== 'string') return;
+      await this.#call('markRequestAsSent', {
+        requestId: request.id,
+        requestType: request.type,
+        response,
+      });
+    };
+    if (request.type !== RequestType.KeysQuery) return send();
+
+    const next = this.#keyQueryChain.catch(() => undefined).then(send);
+    this.#keyQueryChain = next;
+    await next;
   }
 
   async onIncomingKeyVerificationRequest(sender: string, transactionId: string): Promise<boolean> {
@@ -832,13 +841,7 @@
         // Sequential: the engine's queue is ordered and later requests can depend on
         // earlier ones having landed.
         // eslint-disable-next-line no-await-in-loop
-        const response = await sendOutgoingRequest(this.#mx, request);
-        // eslint-disable-next-line no-await-in-loop
-        await this.#call('markRequestAsSent', {
-          requestId: request.id,
-          requestType: request.type,
-          response,
-        });
+        await this.#sendTracked(request);
         sent += 1;
       } catch (error) {
         // Loud: a request the engine never marks sent is retried on every sync forever.
@@ -1089,19 +1092,24 @@
     return this.#serializeForRoom(room.roomId, () => this.#encryptEventInner(event, room));
   }
 
-  async #encryptEventInner(event: MatrixEvent, room: Room): Promise<void> {
-    // The megolm session has to reach every device in the room before the event does.
+  async #prepareRoomForEncryption(room: Room): Promise<string[]> {
     const members = await room.getEncryptionTargetMembers();
     const users = members.map((member) => member.userId);
 
-    if (this.#roomsWithTrackedMembers.has(room.roomId)) {
-      void this.#flushOutgoingRequests();
-    } else {
+    if (!this.#roomsWithTrackedMembers.has(room.roomId)) {
       await this.#trackUsers(users);
-      await this.#flushOutgoingRequests();
+      await this.#sendTracked(await this.#call('queryKeysForUsers', { users }));
       this.#roomsWithTrackedMembers.add(room.roomId);
     }
 
+    void this.#flushOutgoingRequests();
+    return users;
+  }
+
+  async #encryptEventInner(event: MatrixEvent, room: Room): Promise<void> {
+    // The megolm session has to reach every device in the room before the event does.
+    const users = await this.#prepareRoomForEncryption(room);
+
     await this.#ensureSessionsForUsers(users);
 
     const shared = ((await this.#call('shareRoomKey', {
@@ -1493,11 +1501,8 @@
 
   prepareToEncrypt(room: Room): void {
     void this.#serializeForRoom(room.roomId, async () => {
-      const members = await room.getEncryptionTargetMembers();
-      const users = members.map((member) => member.userId);
-      await this.#trackUsers(users);
+      const users = await this.#prepareRoomForEncryption(room);
       await this.#ensureSessionsForUsers(users);
-      await this.#flushOutgoingRequests();
     }).catch((error: unknown) => engineCryptoLog.warn('general', 'prepareToEncrypt failed', error));
   }
 

src/app/crypto/engineCrypto/keyQueryOrdering.test.ts

@@ -0,0 +1,99 @@
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import type { MatrixClient, MatrixEvent, Room } from '$types/matrix-sdk';
+import { engineInvoke } from '../olmMachine/engineInvoke';
+import { EngineCrypto } from './EngineCrypto';
+
+vi.mock('../olmMachine/engineInvoke', () => ({
+  engineInvoke: vi.fn<(...args: never[]) => Promise<unknown>>(),
+}));
+
+const mockInvoke = vi.mocked(engineInvoke);
+const identity = { userId: '@me:e.org', deviceId: 'D' };
+
+const room = {
+  roomId: '!room:e.org',
+  getEncryptionTargetMembers: async () => [{ userId: '@a:e.org' }],
+  getHistoryVisibility: () => 'shared',
+  getBlacklistUnverifiedDevices: () => false,
+  currentState: { getStateEvents: () => null },
+} as unknown as Room;
+
+const event = () =>
+  ({
+    getType: () => 'm.room.message',
+    getContent: () => ({ body: 'hello' }),
+    makeEncrypted: vi.fn<() => void>(),
+  }) as unknown as MatrixEvent;
+
+describe('keys/query ordering', () => {
+  afterEach(() => mockInvoke.mockReset());
+
+  it('waits for a background query HTTP response and acknowledgement before a first-room query', async () => {
+    const http = Promise.withResolvers<string>();
+    const acknowledgement = Promise.withResolvers<void>();
+    let backgroundAcknowledged = false;
+    let queries = 0;
+    const authedRequest = vi.fn<(...args: never[]) => Promise<string>>(async (_method, url) => {
+      if (url !== '/_matrix/client/v3/keys/query') return '{}';
+      queries += 1;
+      if (queries === 1) return http.promise;
+      return '{}';
+    });
+    mockInvoke.mockImplementation(async (_identity, method, args) => {
+      if (method === 'outgoingRequests')
+        return backgroundAcknowledged ? [] : [{ id: 'background', type: 1, body: '{}' }];
+      if (method === 'markRequestAsSent') {
+        if ((args as { requestId: string }).requestId === 'background') {
+          await acknowledgement.promise;
+          backgroundAcknowledged = true;
+        }
+        return null;
+      }
+      if (method === 'queryKeysForUsers') return { id: 'room', type: 1, body: '{}' };
+      if (method === 'getMissingSessions') return null;
+      if (method === 'shareRoomKey') return [];
+      if (method === 'encryptRoomEvent') return '{}';
+      if (method === 'identityKeys') return { ed25519: 'ed', curve25519: 'curve' };
+      return null;
+    });
+    const crypto = new EngineCrypto(
+      { http: { authedRequest } } as unknown as MatrixClient,
+      identity
+    );
+
+    crypto.onSyncCompleted({});
+    const keyQueries = () =>
+      authedRequest.mock.calls.filter(([, url]) => url === '/_matrix/client/v3/keys/query');
+    await vi.waitFor(() => expect(keyQueries()).toHaveLength(1));
+
+    const send = crypto.encryptEvent(event(), room);
+    try {
+      await vi.waitFor(() =>
+        expect(mockInvoke).toHaveBeenCalledWith(
+          expect.anything(),
+          'queryKeysForUsers',
+          expect.anything()
+        )
+      );
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(keyQueries()).toHaveLength(1);
+
+      http.resolve('{}');
+      await vi.waitFor(() =>
+        expect(mockInvoke).toHaveBeenCalledWith(
+          expect.anything(),
+          'markRequestAsSent',
+          expect.anything()
+        )
+      );
+      expect(keyQueries()).toHaveLength(1);
+
+      acknowledgement.resolve();
+      await vi.waitFor(() => expect(keyQueries()).toHaveLength(2));
+    } finally {
+      http.resolve('{}');
+      acknowledgement.resolve();
+      await send;
+    }
+  });
+});

src/app/crypto/engineCrypto/sendConcurrency.test.ts

@@ -0,0 +1,188 @@
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import type { MatrixClient, MatrixEvent, Room } from '$types/matrix-sdk';
+import { engineInvoke } from '../olmMachine/engineInvoke';
+import { EngineCrypto } from './EngineCrypto';
+
+vi.mock('../olmMachine/engineInvoke', () => ({
+  engineInvoke: vi.fn<(...args: never[]) => Promise<unknown>>(),
+}));
+
+const mockInvoke = vi.mocked(engineInvoke);
+
+const room = (roomId: string) =>
+  ({
+    roomId,
+    getEncryptionTargetMembers: async () => [{ userId: `@${roomId.slice(1, 2)}:e.org` }],
+    getHistoryVisibility: () => 'shared',
+    getBlacklistUnverifiedDevices: () => false,
+    currentState: { getStateEvents: () => null },
+  }) as unknown as Room;
+
+const event = (body: string) =>
+  ({
+    getType: () => 'm.room.message',
+    getContent: () => ({ body }),
+    makeEncrypted: vi.fn<() => void>(),
+  }) as unknown as MatrixEvent;
+
+const cryptoWith = (authedRequest: ReturnType<typeof vi.fn>) =>
+  new EngineCrypto({ http: { authedRequest } } as unknown as MatrixClient, {
+    userId: '@me:e.org',
+    deviceId: 'D',
+  });
+
+describe('encrypted send concurrency', () => {
+  afterEach(() => mockInvoke.mockReset());
+
+  it('holds same-room sends behind preparation response and acknowledgement', async () => {
+    const queryResponse = Promise.withResolvers<string>();
+    const queryAcknowledged = Promise.withResolvers<void>();
+    const authedRequest = vi.fn<(...args: never[]) => Promise<string>>(async (_method, url) => {
+      if (url === '/_matrix/client/v3/keys/query') return queryResponse.promise;
+      return '{}';
+    });
+    mockInvoke.mockImplementation(async (_identity, method, args) => {
+      if (method === 'queryKeysForUsers') return { id: 'query', type: 1, body: '{}' };
+      if (method === 'markRequestAsSent' && (args as { requestId: string }).requestId === 'query') {
+        await queryAcknowledged.promise;
+      }
+      if (method === 'encryptRoomEvent') return '{}';
+      if (method === 'identityKeys') return { ed25519: 'ed', curve25519: 'curve' };
+      return null;
+    });
+
+    const crypto = cryptoWith(authedRequest);
+    const target = room('!room:e.org');
+    crypto.prepareToEncrypt(target);
+    const sends = [
+      crypto.encryptEvent(event('first'), target),
+      crypto.encryptEvent(event('second'), target),
+    ];
+
+    try {
+      await vi.waitFor(() => expect(authedRequest).toHaveBeenCalledOnce());
+      expect(
+        mockInvoke.mock.calls.filter(([, method]) => method === 'getMissingSessions')
+      ).toHaveLength(0);
+
+      queryResponse.resolve('{}');
+      await vi.waitFor(() =>
+        expect(mockInvoke.mock.calls).toContainEqual([
+          expect.anything(),
+          'markRequestAsSent',
+          expect.objectContaining({ requestId: 'query' }),
+        ])
+      );
+      expect(
+        mockInvoke.mock.calls.filter(([, method]) => method === 'getMissingSessions')
+      ).toHaveLength(0);
+
+      queryAcknowledged.resolve();
+      await Promise.all(sends);
+
+      expect(
+        mockInvoke.mock.calls.filter(([, method]) => method === 'queryKeysForUsers')
+      ).toHaveLength(1);
+      expect(
+        mockInvoke.mock.calls
+          .filter(([, method]) => method === 'encryptRoomEvent')
+          .map(([, , args]) => JSON.parse((args as { content: string }).content).body)
+      ).toEqual(['first', 'second']);
+    } finally {
+      queryResponse.resolve('{}');
+      queryAcknowledged.resolve();
+      await Promise.allSettled(sends);
+    }
+  });
+
+  it('keeps claims for separate rooms serialized until their response is acknowledged', async () => {
+    const firstClaimResponse = Promise.withResolvers<string>();
+    const firstClaimAcknowledged = Promise.withResolvers<void>();
+    let claimsSent = 0;
+    let heldClaimId: string | undefined;
+    let queries = 0;
+    const authedRequest = vi.fn<(...args: never[]) => Promise<string>>(async (_method, url) => {
+      if (url === '/_matrix/client/v3/keys/claim' && claimsSent++ === 0) {
+        return firstClaimResponse.promise;
+      }
+      return '{}';
+    });
+    mockInvoke.mockImplementation(async (_identity, method, args) => {
+      if (method === 'queryKeysForUsers') return { id: `query-${queries++}`, type: 1, body: '{}' };
+      if (method === 'getMissingSessions') {
+        const user = (args as { users: string[] }).users[0];
+        return { id: `claim-${user}`, type: 2, body: '{}' };
+      }
+      if (method === 'markRequestAsSent') {
+        const requestId = (args as { requestId: string }).requestId;
+        if (requestId.startsWith('claim-') && !heldClaimId) {
+          heldClaimId = requestId;
+          await firstClaimAcknowledged.promise;
+        }
+      }
+      if (method === 'encryptRoomEvent') return '{}';
+      if (method === 'identityKeys') return { ed25519: 'ed', curve25519: 'curve' };
+      return null;
+    });
+
+    const crypto = cryptoWith(authedRequest);
+    const sends = [
+      crypto.encryptEvent(event('a'), room('!a:e.org')),
+      crypto.encryptEvent(event('b'), room('!b:e.org')),
+    ];
+
+    try {
+      await vi.waitFor(() => expect(claimsSent).toBe(1));
+      expect(
+        mockInvoke.mock.calls.filter(([, method]) => method === 'getMissingSessions')
+      ).toHaveLength(1);
+
+      firstClaimResponse.resolve('{}');
+      await vi.waitFor(() => expect(heldClaimId).toBeDefined());
+      expect(
+        mockInvoke.mock.calls.filter(([, method]) => method === 'getMissingSessions')
+      ).toHaveLength(1);
+
+      firstClaimAcknowledged.resolve();
+      await Promise.all(sends);
+      expect(
+        mockInvoke.mock.calls.filter(([, method]) => method === 'getMissingSessions')
+      ).toHaveLength(2);
+    } finally {
+      firstClaimResponse.resolve('{}');
+      firstClaimAcknowledged.resolve();
+      await Promise.allSettled(sends);
+    }
+  });
+
+  it('lets a queued send retry after preparation query failure', async () => {
+    const failedQuery = Promise.withResolvers<string>();
+    let queries = 0;
+    const authedRequest = vi.fn<(...args: never[]) => Promise<string>>(async (_method, url) => {
+      if (url === '/_matrix/client/v3/keys/query' && queries++ === 0) return failedQuery.promise;
+      return '{}';
+    });
+    mockInvoke.mockImplementation(async (_identity, method) => {
+      if (method === 'queryKeysForUsers') return { id: 'query', type: 1, body: '{}' };
+      if (method === 'encryptRoomEvent') return '{}';
+      if (method === 'identityKeys') return { ed25519: 'ed', curve25519: 'curve' };
+      return null;
+    });
+
+    const crypto = cryptoWith(authedRequest);
+    const target = room('!room:e.org');
+    crypto.prepareToEncrypt(target);
+    let send: Promise<void> | undefined;
+    try {
+      await vi.waitFor(() => expect(queries).toBe(1));
+      send = crypto.encryptEvent(event('after failure'), target);
+      failedQuery.reject(new Error('query failed'));
+
+      await expect(send).resolves.toBeUndefined();
+      expect(queries).toBe(2);
+    } finally {
+      failedQuery.reject(new Error('query failed'));
+      await send?.catch(() => undefined);
+    }
+  });
+});