From 02a0bee7569a419847a1fb29a82607f26003f4d8 Mon Sep 17 00:00:00 2001 From: Erwan Leboucher Date: Tue, 8 Sep 2026 20:06:50 +0200 Subject: [PATCH] Fix first encrypted message delay and key query ordering --- src/app/crypto/engineCrypto/EngineCrypto.ts | 55 ++--- .../engineCrypto/keyQueryOrdering.test.ts | 99 +++++++++ .../engineCrypto/sendConcurrency.test.ts | 188 ++++++++++++++++++ .../crypto/engineCrypto/sendLatency.test.ts | 160 +++++++++++++++ 4 files changed, 477 insertions(+), 25 deletions(-) create mode 100644 src/app/crypto/engineCrypto/keyQueryOrdering.test.ts create mode 100644 src/app/crypto/engineCrypto/sendConcurrency.test.ts create mode 100644 src/app/crypto/engineCrypto/sendLatency.test.ts diff --git a/src/app/crypto/engineCrypto/EngineCrypto.ts b/src/app/crypto/engineCrypto/EngineCrypto.ts index a6f772824..e2fdd4e16 100644 --- a/src/app/crypto/engineCrypto/EngineCrypto.ts +++ b/src/app/crypto/engineCrypto/EngineCrypto.ts @@ -45,7 +45,7 @@ import { 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, @@ -358,6 +358,8 @@ export class EngineCrypto #claimChain: Promise = Promise.resolve(); + #keyQueryChain: Promise = Promise.resolve(); + readonly #encryptionChains = new Map>(); readonly #backupUpload = createCoalescedRunner( @@ -754,13 +756,20 @@ export class EngineCrypto async #sendTracked(request: unknown): Promise { 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 { @@ -832,13 +841,7 @@ export class EngineCrypto // 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 @@ export class EngineCrypto return this.#serializeForRoom(room.roomId, () => this.#encryptEventInner(event, room)); } - async #encryptEventInner(event: MatrixEvent, room: Room): Promise { - // The megolm session has to reach every device in the room before the event does. + async #prepareRoomForEncryption(room: Room): Promise { 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 { + // 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 @@ export class EngineCrypto 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)); } diff --git a/src/app/crypto/engineCrypto/keyQueryOrdering.test.ts b/src/app/crypto/engineCrypto/keyQueryOrdering.test.ts new file mode 100644 index 000000000..cba4ceecd --- /dev/null +++ b/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>(), +})); + +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(); + const acknowledgement = Promise.withResolvers(); + let backgroundAcknowledged = false; + let queries = 0; + const authedRequest = vi.fn<(...args: never[]) => Promise>(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; + } + }); +}); diff --git a/src/app/crypto/engineCrypto/sendConcurrency.test.ts b/src/app/crypto/engineCrypto/sendConcurrency.test.ts new file mode 100644 index 000000000..7d165d671 --- /dev/null +++ b/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>(), +})); + +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) => + 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(); + const queryAcknowledged = Promise.withResolvers(); + const authedRequest = vi.fn<(...args: never[]) => Promise>(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(); + const firstClaimAcknowledged = Promise.withResolvers(); + let claimsSent = 0; + let heldClaimId: string | undefined; + let queries = 0; + const authedRequest = vi.fn<(...args: never[]) => Promise>(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(); + let queries = 0; + const authedRequest = vi.fn<(...args: never[]) => Promise>(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 | 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); + } + }); +}); diff --git a/src/app/crypto/engineCrypto/sendLatency.test.ts b/src/app/crypto/engineCrypto/sendLatency.test.ts new file mode 100644 index 000000000..1a7fba424 --- /dev/null +++ b/src/app/crypto/engineCrypto/sendLatency.test.ts @@ -0,0 +1,160 @@ +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>(), +})); + +const mockInvoke = vi.mocked(engineInvoke); + +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('first encrypted send latency', () => { + afterEach(() => mockInvoke.mockReset()); + + it.each([false, true])( + 'does not wait for an unrelated outgoing drain (%s preparation)', + async (prepareFirst) => { + const signatureUpload = Promise.withResolvers(); + let signatureAcknowledged = false; + const authedRequest = vi.fn<(...args: never[]) => Promise>(async (_method, url) => { + if (url === '/_matrix/client/v3/keys/signatures/upload') return signatureUpload.promise; + return '{}'; + }); + + mockInvoke.mockImplementation(async (_identity, method, args) => { + if (method === 'outgoingRequests') + return signatureAcknowledged ? [] : [{ id: 'signature', type: 4, body: '{}' }]; + if (method === 'markRequestAsSent') { + if ((args as { requestId?: string } | undefined)?.requestId === 'signature') { + signatureAcknowledged = true; + } + return null; + } + if (method === 'queryKeysForUsers') return { id: 'query', type: 1, body: '{}' }; + if (method === 'getMissingSessions') return { id: 'claim', type: 2, body: '{}' }; + if (method === 'shareRoomKey') { + return [ + { + id: 'share', + type: 3, + event_type: 'm.room.encrypted', + txn_id: 'share', + body: '{}', + }, + ]; + } + if (method === 'encryptRoomEvent') return '{}'; + if (method === 'identityKeys') return { ed25519: 'ed', curve25519: 'curve' }; + return null; + }); + + const crypto = new EngineCrypto({ http: { authedRequest } } as unknown as MatrixClient, { + userId: '@me:e.org', + deviceId: 'D', + }); + crypto.onSyncCompleted({}); + await vi.waitFor(() => + expect(authedRequest).toHaveBeenCalledWith( + expect.anything(), + '/_matrix/client/v3/keys/signatures/upload', + expect.anything(), + expect.anything(), + expect.anything() + ) + ); + + let send: Promise | undefined; + try { + if (prepareFirst) crypto.prepareToEncrypt(room); + const encrypted = event(); + send = crypto.encryptEvent(encrypted, room); + await vi.waitFor(() => expect(encrypted.makeEncrypted).toHaveBeenCalledOnce(), { + timeout: 500, + }); + await send; + + expect( + mockInvoke.mock.calls.filter(([, method]) => method === 'queryKeysForUsers') + ).toHaveLength(1); + + const methods = mockInvoke.mock.calls.map(([, method]) => method); + const marked = (requestId: string) => + mockInvoke.mock.calls.findIndex( + ([, method, args]) => + method === 'markRequestAsSent' && + (args as { requestId: string }).requestId === requestId + ); + expect(marked('query')).toBeGreaterThanOrEqual(0); + expect(marked('query')).toBeLessThan(methods.indexOf('getMissingSessions')); + expect(marked('claim')).toBeGreaterThanOrEqual(0); + expect(marked('claim')).toBeLessThan(methods.indexOf('shareRoomKey')); + expect(marked('share')).toBeGreaterThanOrEqual(0); + expect(marked('share')).toBeLessThan(methods.indexOf('encryptRoomEvent')); + } finally { + signatureUpload.resolve('{}'); + await send; + } + } + ); + + it('retries the initial query when preparation could not send it', async () => { + let failFirstQuery = true; + const authedRequest = vi.fn<(...args: never[]) => Promise>(async (_method, url) => { + if (url === '/_matrix/client/v3/keys/query' && failFirstQuery) { + failFirstQuery = false; + throw new Error('query failed'); + } + 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 = new EngineCrypto({ http: { authedRequest } } as unknown as MatrixClient, { + userId: '@me:e.org', + deviceId: 'D', + }); + + crypto.prepareToEncrypt(room); + await vi.waitFor(() => + expect(authedRequest).toHaveBeenCalledWith( + expect.anything(), + '/_matrix/client/v3/keys/query', + expect.anything(), + expect.anything(), + expect.anything() + ) + ); + await vi.waitFor(() => + expect( + mockInvoke.mock.calls.filter(([, method]) => method === 'encryptRoomEvent') + ).toHaveLength(0) + ); + + const encrypted = event(); + await crypto.encryptEvent(encrypted, room); + + expect(encrypted.makeEncrypted).toHaveBeenCalledOnce(); + expect( + authedRequest.mock.calls.filter(([, url]) => url === '/_matrix/client/v3/keys/query') + ).toHaveLength(2); + }); +});