Skip to content

Commit

Permalink
Element-R: handle events which arrive before their keys (#3230)
Browse files Browse the repository at this point in the history 
* minor cleanups to the crypto tests

mostly, this is about using `testUtils.awaitDecryption` rather than custom
code. Some other cleanups too.

* Keep a record of events which are missing their keys

* Retry event decryption when we receive megolm keys
  • Loading branch information
@richvdh
richvdh committed Apr 5, 2023
1 parent e89467c commit 9a840d4
Show file tree
Hide file tree
Showing 6 changed files with 163 additions and 45 deletions.
2 changes: 1 addition & 1 deletion package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@
],
"dependencies": {
"@babel/runtime": "^7.12.5",
"@matrix-org/matrix-sdk-crypto-js": "^0.1.0-alpha.5",
"@matrix-org/matrix-sdk-crypto-js": "^0.1.0-alpha.6",
"another-json": "^0.2.0",
"bs58": "^5.0.0",
"content-type": "^1.0.4",
Expand Down
26 changes: 8 additions & 18 deletions spec/integ/crypto.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -624,10 +624,8 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
expect(decryptedEvent.getContent().body).toEqual("42");
});

oldBackendOnly("Alice receives a megolm message before the session keys", async () => {
it("Alice receives a megolm message before the session keys", async () => {
expectAliceKeyQuery({ device_keys: { "@alice:localhost": {} }, failures: {} });

// https://github.com/vector-im/element-web/issues/2273
await startClientAndAwaitFirstSync();

// if we're using the old crypto impl, stub out some methods in the device manager.
Expand Down Expand Up @@ -667,7 +665,11 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
await syncPromise(aliceClient);

const room = aliceClient.getRoom(ROOM_ID)!;
expect(room.getLiveTimeline().getEvents()[0].getContent().msgtype).toEqual("m.bad.encrypted");
const event = room.getLiveTimeline().getEvents()[0];

// wait for a first attempt at decryption: should fail
await testUtils.awaitDecryption(event);
expect(event.getContent().msgtype).toEqual("m.bad.encrypted");

// now she gets the room_key event
syncResponder.sendOrQueueSyncResponse({
Expand All @@ -678,20 +680,8 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("crypto (%s)", (backend: string,
});
await syncPromise(aliceClient);

const event = room.getLiveTimeline().getEvents()[0];

let decryptedEvent: MatrixEvent;
if (event.getContent().msgtype != "m.bad.encrypted") {
decryptedEvent = event;
} else {
decryptedEvent = await new Promise<MatrixEvent>((resolve) => {
event.once(MatrixEventEvent.Decrypted, (ev) => {
logger.log(`${Date.now()} event ${event.getId()} now decrypted`);
resolve(ev);
});
});
}
expect(decryptedEvent.getContent().body).toEqual("42");
await testUtils.awaitDecryption(event, { waitOnDecryptionFailure: true });
expect(event.getContent().body).toEqual("42");
});

it("Alice gets a second room_key message", async () => {
Expand Down
8 changes: 4 additions & 4 deletions spec/test-utils/test-utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -375,17 +375,17 @@ export async function awaitDecryption(
// already
if (event.getClearContent() !== null) {
if (waitOnDecryptionFailure && event.isDecryptionFailure()) {
logger.log(`${Date.now()} event ${event.getId()} got decryption error; waiting`);
logger.log(`${Date.now()}: event ${event.getId()} got decryption error; waiting`);
} else {
return event;
}
} else {
logger.log(`${Date.now()} event ${event.getId()} is not yet decrypted; waiting`);
logger.log(`${Date.now()}: event ${event.getId()} is not yet decrypted; waiting`);
}

return new Promise((resolve) => {
event.once(MatrixEventEvent.Decrypted, (ev) => {
logger.log(`${Date.now()} event ${event.getId()} now decrypted`);
event.once(MatrixEventEvent.Decrypted, (ev, err) => {
logger.log(`${Date.now()}: MatrixEventEvent.Decrypted for event ${event.getId()}: ${err ?? "success"}`);
resolve(ev);
});
});
Expand Down
3 changes: 3 additions & 0 deletions src/rust-crypto/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,9 @@ export async function initRustCrypto(
// TODO: use the pickle key for the passphrase
const olmMachine = await RustSdkCryptoJs.OlmMachine.initialize(u, d, RUST_SDK_STORE_PREFIX, "test pass");
const rustCrypto = new RustCrypto(olmMachine, http, userId, deviceId);
await olmMachine.registerRoomKeyUpdatedCallback((sessions: RustSdkCryptoJs.RoomKeyInfo[]) =>
rustCrypto.onRoomKeysUpdated(sessions),
);

logger.info("Completed rust crypto-sdk setup");
return rustCrypto;
Expand Down
161 changes: 143 additions & 18 deletions src/rust-crypto/rust-crypto.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
Copyright 2022 The Matrix.org Foundation C.I.C.
Copyright 2022-2023 The Matrix.org Foundation C.I.C.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -29,6 +29,7 @@ import { DeviceTrustLevel, UserTrustLevel } from "../crypto/CrossSigning";
import { RoomEncryptor } from "./RoomEncryptor";
import { OutgoingRequest, OutgoingRequestProcessor } from "./OutgoingRequestProcessor";
import { KeyClaimManager } from "./KeyClaimManager";
import { MapWithDefault } from "../utils";

/**
* An implementation of {@link CryptoBackend} using the Rust matrix-sdk-crypto.
Expand All @@ -45,6 +46,7 @@ export class RustCrypto implements CryptoBackend {
/** mapping of roomId → encryptor class */
private roomEncryptors: Record<string, RoomEncryptor> = {};

private eventDecryptor: EventDecryptor;
private keyClaimManager: KeyClaimManager;
private outgoingRequestProcessor: OutgoingRequestProcessor;

Expand All @@ -56,6 +58,7 @@ export class RustCrypto implements CryptoBackend {
) {
this.outgoingRequestProcessor = new OutgoingRequestProcessor(olmMachine, http);
this.keyClaimManager = new KeyClaimManager(olmMachine, this.outgoingRequestProcessor);
this.eventDecryptor = new EventDecryptor(olmMachine);
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -101,23 +104,7 @@ export class RustCrypto implements CryptoBackend {
// through decryptEvent and hence get rid of this case.
throw new Error("to-device event was not decrypted in preprocessToDeviceMessages");
}
const res = (await this.olmMachine.decryptRoomEvent(
JSON.stringify({
event_id: event.getId(),
type: event.getWireType(),
sender: event.getSender(),
state_key: event.getStateKey(),
content: event.getWireContent(),
origin_server_ts: event.getTs(),
}),
new RustSdkCryptoJs.RoomId(event.getRoomId()!),
)) as RustSdkCryptoJs.DecryptedRoomEvent;
return {
clearEvent: JSON.parse(res.event),
claimedEd25519Key: res.senderClaimedEd25519Key,
senderCurve25519Key: res.senderCurve25519Key,
forwardingCurve25519KeyChain: res.forwardingCurve25519KeyChain,
};
return await this.eventDecryptor.attemptEventDecryption(event);
}

public getEventEncryptionInfo(event: MatrixEvent): IEncryptedEventInfo {
Expand Down Expand Up @@ -303,6 +290,43 @@ export class RustCrypto implements CryptoBackend {
enc.onRoomMembership(member);
}

/** Callback for OlmMachine.registerRoomKeyUpdatedCallback
*
* Called by the rust-sdk whenever there is an update to (megolm) room keys. We
* check if we have any events waiting for the given keys, and schedule them for
* a decryption retry if so.
*
* @param keys - details of the updated keys
*/
public async onRoomKeysUpdated(keys: RustSdkCryptoJs.RoomKeyInfo[]): Promise<void> {
for (const key of keys) {
this.onRoomKeyUpdated(key);
}
}

private onRoomKeyUpdated(key: RustSdkCryptoJs.RoomKeyInfo): void {
logger.debug(`Got update for session ${key.senderKey.toBase64()}|${key.sessionId} in ${key.roomId.toString()}`);
const pendingList = this.eventDecryptor.getEventsPendingRoomKey(key);
if (pendingList.length === 0) return;

logger.debug(
"Retrying decryption on events:",
pendingList.map((e) => `${e.getId()}`),
);

// Have another go at decrypting events with this key.
//
// We don't want to end up blocking the callback from Rust, which could otherwise end up dropping updates,
// so we don't wait for the decryption to complete. In any case, there is no need to wait:
// MatrixEvent.attemptDecryption ensures that there is only one decryption attempt happening at once,
// and deduplicates repeated attempts for the same event.
for (const ev of pendingList) {
ev.attemptDecryption(this, { isRetry: true }).catch((_e) => {
logger.info(`Still unable to decrypt event ${ev.getId()} after receiving key`);
});
}
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Outgoing requests
Expand Down Expand Up @@ -332,3 +356,104 @@ export class RustCrypto implements CryptoBackend {
}
}
}

class EventDecryptor {
/**
* Events which we couldn't decrypt due to unknown sessions / indexes.
*
* Map from senderKey to sessionId to Set of MatrixEvents
*/
private eventsPendingKey = new MapWithDefault<string, MapWithDefault<string, Set<MatrixEvent>>>(
() => new MapWithDefault<string, Set<MatrixEvent>>(() => new Set()),
);

public constructor(private readonly olmMachine: RustSdkCryptoJs.OlmMachine) {}

public async attemptEventDecryption(event: MatrixEvent): Promise<IEventDecryptionResult> {
logger.info("Attempting decryption of event", event);
// add the event to the pending list *before* attempting to decrypt.
// then, if the key turns up while decryption is in progress (and
// decryption fails), we will schedule a retry.
// (fixes https://github.com/vector-im/element-web/issues/5001)
this.addEventToPendingList(event);

const res = (await this.olmMachine.decryptRoomEvent(
JSON.stringify({
event_id: event.getId(),
type: event.getWireType(),
sender: event.getSender(),
state_key: event.getStateKey(),
content: event.getWireContent(),
origin_server_ts: event.getTs(),
}),
new RustSdkCryptoJs.RoomId(event.getRoomId()!),
)) as RustSdkCryptoJs.DecryptedRoomEvent;

// Success. We can remove the event from the pending list, if
// that hasn't already happened.
this.removeEventFromPendingList(event);

return {
clearEvent: JSON.parse(res.event),
claimedEd25519Key: res.senderClaimedEd25519Key,
senderCurve25519Key: res.senderCurve25519Key,
forwardingCurve25519KeyChain: res.forwardingCurve25519KeyChain,
};
}

/**
* Look for events which are waiting for a given megolm session
*
* Returns a list of events which were encrypted by `session` and could not be decrypted
*
* @param session -
*/
public getEventsPendingRoomKey(session: RustSdkCryptoJs.RoomKeyInfo): MatrixEvent[] {
const senderPendingEvents = this.eventsPendingKey.get(session.senderKey.toBase64());
if (!senderPendingEvents) return [];

const sessionPendingEvents = senderPendingEvents.get(session.sessionId);
if (!sessionPendingEvents) return [];

const roomId = session.roomId.toString();
return [...sessionPendingEvents].filter((ev) => ev.getRoomId() === roomId);
}

/**
* Add an event to the list of those awaiting their session keys.
*/
private addEventToPendingList(event: MatrixEvent): void {
const content = event.getWireContent();
const senderKey = content.sender_key;
const sessionId = content.session_id;

const senderPendingEvents = this.eventsPendingKey.getOrCreate(senderKey);
const sessionPendingEvents = senderPendingEvents.getOrCreate(sessionId);
sessionPendingEvents.add(event);
}

/**
* Remove an event from the list of those awaiting their session keys.
*/
private removeEventFromPendingList(event: MatrixEvent): void {
const content = event.getWireContent();
const senderKey = content.sender_key;
const sessionId = content.session_id;

const senderPendingEvents = this.eventsPendingKey.get(senderKey);
if (!senderPendingEvents) return;

const sessionPendingEvents = senderPendingEvents.get(sessionId);
if (!sessionPendingEvents) return;

sessionPendingEvents.delete(event);

// also clean up the higher-level maps if they are now empty
if (sessionPendingEvents.size === 0) {
senderPendingEvents.delete(sessionId);
if (senderPendingEvents.size === 0) {
this.eventsPendingKey.delete(senderKey);
}
}
}
}
8 changes: 4 additions & 4 deletions yarn.lock
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1538,10 +1538,10 @@
dependencies:
lodash "^4.17.21"

"@matrix-org/matrix-sdk-crypto-js@^0.1.0-alpha.5":
version "0.1.0-alpha.5"
resolved "https://registry.yarnpkg.com/@matrix-org/matrix-sdk-crypto-js/-/matrix-sdk-crypto-js-0.1.0-alpha.5.tgz#60ede2c43b9d808ba8cf46085a3b347b290d9658"
integrity sha512-2KjAgWNGfuGLNjJwsrs6gGX157vmcTfNrA4u249utgnMPbJl7QwuUqh1bGxQ0PpK06yvZjgPlkna0lTbuwtuQw==
"@matrix-org/matrix-sdk-crypto-js@^0.1.0-alpha.6":
version "0.1.0-alpha.6"
resolved "https://registry.yarnpkg.com/@matrix-org/matrix-sdk-crypto-js/-/matrix-sdk-crypto-js-0.1.0-alpha.6.tgz#c0bdb9ab0d30179b8ef744d1b4010b0ad0ab9c3a"
integrity sha512-7hMffzw7KijxDyyH/eUyTfrLeCQHuyU3kaPOKGhcl3DZ3vx7bCncqjGMGTnxNPoP23I6gosvKSbO+3wYOT24Xg==

"@matrix-org/olm@https://gitlab.matrix.org/api/v4/projects/27/packages/npm/@matrix-org/olm/-/@matrix-org/olm-3.2.14.tgz":
version "3.2.14"
Expand Down

0 comments on commit 9a840d4

Please sign in to comment.