Live location sharing: handle encrypted messages in processBeaconEvents (#2327)

* handle encrypted locations

Signed-off-by: Kerry Archibald <kerrya@element.io>

* fix processBeaconEvents to handle encrypted events

Signed-off-by: Kerry Archibald <kerrya@element.io>

spec/unit/matrix-client.spec.ts

@@ -1047,7 +1047,7 @@ describe("MatrixClient", function() {
                 expect(roomStateProcessSpy).not.toHaveBeenCalled();
             });
 
-            it('calls room states processBeaconEvents with m.beacon events', () => {
+            it('calls room states processBeaconEvents with events', () => {
                 const room = new Room(roomId, client, userId);
                 const roomStateProcessSpy = jest.spyOn(room.currentState, 'processBeaconEvents');
 
@@ -1055,7 +1055,7 @@ describe("MatrixClient", function() {
                 const beaconEvent = makeBeaconEvent(userId);
 
                 client.processBeaconEvents(room, [messageEvent, beaconEvent]);
-                expect(roomStateProcessSpy).toHaveBeenCalledWith([beaconEvent]);
+                expect(roomStateProcessSpy).toHaveBeenCalledWith([messageEvent, beaconEvent], client);
             });
         });
     });

spec/unit/room-state.spec.js

@@ -3,6 +3,12 @@ import { makeBeaconEvent, makeBeaconInfoEvent } from "../test-utils/beacon";
 import { filterEmitCallsByEventType } from "../test-utils/emitter";
 import { RoomState, RoomStateEvent } from "../../src/models/room-state";
 import { BeaconEvent, getBeaconInfoIdentifier } from "../../src/models/beacon";
+import { EventType, RelationType } from "../../src/@types/event";
+import {
+    MatrixEvent,
+    MatrixEventEvent,
+} from "../../src/models/event";
+import { M_BEACON } from "../../src/@types/beacon";
 
 describe("RoomState", function() {
     const roomId = "!foo:bar";
@@ -717,26 +723,62 @@ describe("RoomState", function() {
         const beacon1 = makeBeaconInfoEvent(userA, roomId, {}, '$beacon1', '$beacon1');
         const beacon2 = makeBeaconInfoEvent(userB, roomId, {}, '$beacon2', '$beacon2');
 
+        const mockClient = { decryptEventIfNeeded: jest.fn() };
+
+        beforeEach(() => {
+            mockClient.decryptEventIfNeeded.mockClear();
+        });
+
         it('does nothing when state has no beacons', () => {
             const emitSpy = jest.spyOn(state, 'emit');
-            state.processBeaconEvents([makeBeaconEvent(userA, { beaconInfoId: '$beacon1' })]);
+            state.processBeaconEvents([makeBeaconEvent(userA, { beaconInfoId: '$beacon1' })], mockClient);
             expect(emitSpy).not.toHaveBeenCalled();
+            expect(mockClient.decryptEventIfNeeded).not.toHaveBeenCalled();
         });
 
         it('does nothing when there are no events', () => {
             state.setStateEvents([beacon1, beacon2]);
             const emitSpy = jest.spyOn(state, 'emit').mockClear();
-            state.processBeaconEvents([]);
+            state.processBeaconEvents([], mockClient);
             expect(emitSpy).not.toHaveBeenCalled();
+            expect(mockClient.decryptEventIfNeeded).not.toHaveBeenCalled();
         });
 
+        describe('without encryption', () => {
             it('discards events for beacons that are not in state', () => {
                 const location = makeBeaconEvent(userA, {
                     beaconInfoId: 'some-other-beacon',
                 });
+                const otherRelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessage,
+                    content: {
+                        ['m.relates_to']: {
+                            event_id: 'whatever',
+                        },
+                    },
+                });
                 state.setStateEvents([beacon1, beacon2]);
                 const emitSpy = jest.spyOn(state, 'emit').mockClear();
-            state.processBeaconEvents([location]);
+                state.processBeaconEvents([location, otherRelatedEvent], mockClient);
+                expect(emitSpy).not.toHaveBeenCalled();
+            });
+
+            it('discards events that are not beacon type', () => {
+                // related to beacon1
+                const otherRelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessage,
+                    content: {
+                        ['m.relates_to']: {
+                            rel_type: RelationType.Reference,
+                            event_id: beacon1.getId(),
+                        },
+                    },
+                });
+                state.setStateEvents([beacon1, beacon2]);
+                const emitSpy = jest.spyOn(state, 'emit').mockClear();
+                state.processBeaconEvents([otherRelatedEvent], mockClient);
                 expect(emitSpy).not.toHaveBeenCalled();
             });
 
@@ -751,18 +793,168 @@ describe("RoomState", function() {
                     beaconInfoId: 'some-other-beacon',
                 });
 
-            state.setStateEvents([beacon1, beacon2]);
+                state.setStateEvents([beacon1, beacon2], mockClient);
 
                 expect(state.beacons.size).toEqual(2);
 
                 const beaconInstance = state.beacons.get(getBeaconInfoIdentifier(beacon1));
                 const addLocationsSpy = jest.spyOn(beaconInstance, 'addLocations');
 
-            state.processBeaconEvents([location1, location2, location3]);
+                state.processBeaconEvents([location1, location2, location3], mockClient);
 
-            expect(addLocationsSpy).toHaveBeenCalledTimes(1);
+                expect(addLocationsSpy).toHaveBeenCalledTimes(2);
                 // only called with locations for beacon1
-            expect(addLocationsSpy).toHaveBeenCalledWith([location1, location2]);
+                expect(addLocationsSpy).toHaveBeenCalledWith([location1]);
+                expect(addLocationsSpy).toHaveBeenCalledWith([location2]);
+            });
+        });
+
+        describe('with encryption', () => {
+            const beacon1RelationContent = { ['m.relates_to']: {
+                rel_type: RelationType.Reference,
+                event_id: beacon1.getId(),
+            } };
+            const relatedEncryptedEvent = new MatrixEvent({
+                sender: userA,
+                type: EventType.RoomMessageEncrypted,
+                content: beacon1RelationContent,
+            });
+            const decryptingRelatedEvent = new MatrixEvent({
+                sender: userA,
+                type: EventType.RoomMessageEncrypted,
+                content: beacon1RelationContent,
+            });
+            jest.spyOn(decryptingRelatedEvent, 'isBeingDecrypted').mockReturnValue(true);
+
+            const failedDecryptionRelatedEvent = new MatrixEvent({
+                sender: userA,
+                type: EventType.RoomMessageEncrypted,
+                content: beacon1RelationContent,
+            });
+            jest.spyOn(failedDecryptionRelatedEvent, 'isDecryptionFailure').mockReturnValue(true);
+
+            it('discards events without relations', () => {
+                const unrelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessageEncrypted,
+                });
+                state.setStateEvents([beacon1, beacon2]);
+                const emitSpy = jest.spyOn(state, 'emit').mockClear();
+                state.processBeaconEvents([unrelatedEvent], mockClient);
+                expect(emitSpy).not.toHaveBeenCalled();
+                // discard unrelated events early
+                expect(mockClient.decryptEventIfNeeded).not.toHaveBeenCalled();
+            });
+
+            it('discards events for beacons that are not in state', () => {
+                const location = makeBeaconEvent(userA, {
+                    beaconInfoId: 'some-other-beacon',
+                });
+                const otherRelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessageEncrypted,
+                    content: {
+                        ['m.relates_to']: {
+                            rel_type: RelationType.Reference,
+                            event_id: 'whatever',
+                        },
+                    },
+                });
+                state.setStateEvents([beacon1, beacon2]);
+
+                const beacon = state.beacons.get(getBeaconInfoIdentifier(beacon1));
+                const addLocationsSpy = jest.spyOn(beacon, 'addLocations').mockClear();
+                state.processBeaconEvents([location, otherRelatedEvent], mockClient);
+                expect(addLocationsSpy).not.toHaveBeenCalled();
+                // discard unrelated events early
+                expect(mockClient.decryptEventIfNeeded).not.toHaveBeenCalled();
+            });
+
+            it('decrypts related events if needed', () => {
+                const location = makeBeaconEvent(userA, {
+                    beaconInfoId: beacon1.getId(),
+                });
+                state.setStateEvents([beacon1, beacon2]);
+                state.processBeaconEvents([location, relatedEncryptedEvent], mockClient);
+                // discard unrelated events early
+                expect(mockClient.decryptEventIfNeeded).toHaveBeenCalledTimes(2);
+            });
+
+            it('listens for decryption on events that are being decrypted', () => {
+                const decryptingRelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessageEncrypted,
+                    content: beacon1RelationContent,
+                });
+                jest.spyOn(decryptingRelatedEvent, 'isBeingDecrypted').mockReturnValue(true);
+                // spy on event.once
+                const eventOnceSpy = jest.spyOn(decryptingRelatedEvent, 'once');
+
+                state.setStateEvents([beacon1, beacon2]);
+                state.processBeaconEvents([decryptingRelatedEvent], mockClient);
+
+                // listener was added
+                expect(eventOnceSpy).toHaveBeenCalled();
+            });
+
+            it('listens for decryption on events that have decryption failure', () => {
+                const failedDecryptionRelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessageEncrypted,
+                    content: beacon1RelationContent,
+                });
+                jest.spyOn(failedDecryptionRelatedEvent, 'isDecryptionFailure').mockReturnValue(true);
+                // spy on event.once
+                const eventOnceSpy = jest.spyOn(decryptingRelatedEvent, 'once');
+
+                state.setStateEvents([beacon1, beacon2]);
+                state.processBeaconEvents([decryptingRelatedEvent], mockClient);
+
+                // listener was added
+                expect(eventOnceSpy).toHaveBeenCalled();
+            });
+
+            it('discard events that are not m.beacon type after decryption', () => {
+                const decryptingRelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessageEncrypted,
+                    content: beacon1RelationContent,
+                });
+                jest.spyOn(decryptingRelatedEvent, 'isBeingDecrypted').mockReturnValue(true);
+                state.setStateEvents([beacon1, beacon2]);
+                const beacon = state.beacons.get(getBeaconInfoIdentifier(beacon1));
+                const addLocationsSpy = jest.spyOn(beacon, 'addLocations').mockClear();
+                state.processBeaconEvents([decryptingRelatedEvent], mockClient);
+
+                // this event is a message after decryption
+                decryptingRelatedEvent.type = EventType.RoomMessage;
+                decryptingRelatedEvent.emit(MatrixEventEvent.Decrypted);
+
+                expect(addLocationsSpy).not.toHaveBeenCalled();
+            });
+
+            it('adds locations to beacons after decryption', () => {
+                const decryptingRelatedEvent = new MatrixEvent({
+                    sender: userA,
+                    type: EventType.RoomMessageEncrypted,
+                    content: beacon1RelationContent,
+                });
+                const locationEvent = makeBeaconEvent(userA, {
+                    beaconInfoId: '$beacon1', timestamp: Date.now() + 1,
+                });
+                jest.spyOn(decryptingRelatedEvent, 'isBeingDecrypted').mockReturnValue(true);
+                state.setStateEvents([beacon1, beacon2]);
+                const beacon = state.beacons.get(getBeaconInfoIdentifier(beacon1));
+                const addLocationsSpy = jest.spyOn(beacon, 'addLocations').mockClear();
+                state.processBeaconEvents([decryptingRelatedEvent], mockClient);
+
+                // update type after '''decryption'''
+                decryptingRelatedEvent.event.type = M_BEACON.name;
+                decryptingRelatedEvent.event.content = locationEvent.content;
+                decryptingRelatedEvent.emit(MatrixEventEvent.Decrypted);
+
+                expect(addLocationsSpy).toHaveBeenCalledWith([decryptingRelatedEvent]);
+            });
         });
     });
 });

src/client.ts

@@ -180,7 +180,7 @@ import { MediaHandler } from "./webrtc/mediaHandler";
 import { IRefreshTokenResponse } from "./@types/auth";
 import { TypedEventEmitter } from "./models/typed-event-emitter";
 import { Thread, THREAD_RELATION_TYPE } from "./models/thread";
-import { MBeaconInfoEventContent, M_BEACON, M_BEACON_INFO } from "./@types/beacon";
+import { MBeaconInfoEventContent, M_BEACON_INFO } from "./@types/beacon";
 
 export type Store = IStore;
 export type SessionStore = WebStorageSessionStore;
@@ -5192,7 +5192,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
 
                 const [timelineEvents, threadedEvents] = room.partitionThreadedEvents(matrixEvents);
 
-                this.processBeaconEvents(room, matrixEvents);
+                this.processBeaconEvents(room, timelineEvents);
                 room.addEventsToTimeline(timelineEvents, true, room.getLiveTimeline());
                 await this.processThreadEvents(room, threadedEvents, true);
 
@@ -5335,7 +5335,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         timelineSet.addEventsToTimeline(timelineEvents, true, timeline, res.start);
         // The target event is not in a thread but process the contextual events, so we can show any threads around it.
         await this.processThreadEvents(timelineSet.room, threadedEvents, true);
-        this.processBeaconEvents(timelineSet.room, events);
+        this.processBeaconEvents(timelineSet.room, timelineEvents);
 
         // There is no guarantee that the event ended up in "timeline" (we might have switched to a neighbouring
         // timeline) - so check the room's index again. On the other hand, there's no guarantee the event ended up
@@ -5503,7 +5503,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
                 const timelineSet = eventTimeline.getTimelineSet();
                 const [timelineEvents, threadedEvents] = timelineSet.room.partitionThreadedEvents(matrixEvents);
                 timelineSet.addEventsToTimeline(timelineEvents, backwards, eventTimeline, token);
-                this.processBeaconEvents(timelineSet.room, matrixEvents);
+                this.processBeaconEvents(timelineSet.room, timelineEvents);
                 await this.processThreadEvents(room, threadedEvents, backwards);
 
                 // if we've hit the end of the timeline, we need to stop trying to
@@ -8929,8 +8929,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         if (!events?.length) {
             return;
         }
-        const beaconEvents = events.filter(event => M_BEACON.matches(event.getType()));
-        room.currentState.processBeaconEvents(beaconEvents);
+        room.currentState.processBeaconEvents(events, this);
     }
 
     /**

src/models/room-state.ts

@@ -22,15 +22,14 @@ import { RoomMember } from "./room-member";
 import { logger } from '../logger';
 import * as utils from "../utils";
 import { EventType } from "../@types/event";
-import { MatrixEvent } from "./event";
+import { MatrixEvent, MatrixEventEvent } from "./event";
 import { MatrixClient } from "../client";
 import { GuestAccess, HistoryVisibility, IJoinRuleEventContent, JoinRule } from "../@types/partials";
 import { TypedEventEmitter } from "./typed-event-emitter";
 import { Beacon, BeaconEvent, BeaconEventHandlerMap } from "./beacon";
 import { TypedReEmitter } from "../ReEmitter";
-import { M_BEACON_INFO } from "../@types/beacon";
-import { getBeaconInfoIdentifier } from "./beacon";
-import { BeaconIdentifier } from "..";
+import { M_BEACON, M_BEACON_INFO } from "../@types/beacon";
+import { getBeaconInfoIdentifier, BeaconIdentifier } from "./beacon";
 
 // possible statuses for out-of-band member loading
 enum OobStatus {
@@ -411,7 +410,7 @@ export class RoomState extends TypedEventEmitter<EmittedEvents, EventHandlerMap>
         this.emit(RoomStateEvent.Update, this);
     }
 
-    public processBeaconEvents(events: MatrixEvent[]): void {
+    public processBeaconEvents(events: MatrixEvent[], matrixClient: MatrixClient): void {
         if (
             !events.length ||
             // discard locations if we have no beacons
@@ -420,24 +419,38 @@ export class RoomState extends TypedEventEmitter<EmittedEvents, EventHandlerMap>
             return;
         }
 
-        // names are confusing here
-        // a Beacon is the parent event, but event type is 'm.beacon_info'
-        // a location is the 'child' related to the Beacon, but the event type is 'm.beacon'
-        // group locations by beaconInfo event id
-        const locationEventsByBeaconEventId = events.reduce<Record<string, MatrixEvent[]>>((acc, event) => {
-            const beaconInfoEventId = event.getRelation()?.event_id;
-            if (!acc[beaconInfoEventId]) {
-                acc[beaconInfoEventId] = [];
-            }
-            acc[beaconInfoEventId].push(event);
-            return acc;
-        }, {});
+        const beaconByEventIdDict: Record<string, Beacon> =
+            [...this.beacons.values()].reduce((dict, beacon) => ({ ...dict, [beacon.beaconInfoId]: beacon }), {});
 
-        Object.entries(locationEventsByBeaconEventId).forEach(([beaconInfoEventId, events]) => {
-            const beacon = [...this.beacons.values()].find(beacon => beacon.beaconInfoId === beaconInfoEventId);
+        const processBeaconRelation = (beaconInfoEventId: string, event: MatrixEvent): void => {
+            if (!M_BEACON.matches(event.getType())) {
+                return;
+            }
+
+            const beacon = beaconByEventIdDict[beaconInfoEventId];
 
             if (beacon) {
-                beacon.addLocations(events);
+                beacon.addLocations([event]);
+            }
+        };
+
+        events.forEach((event: MatrixEvent) => {
+            const relatedToEventId = event.getRelation()?.event_id;
+            // not related to a beacon we know about
+            // discard
+            if (!beaconByEventIdDict[relatedToEventId]) {
+                return;
+            }
+
+            matrixClient.decryptEventIfNeeded(event);
+
+            if (event.isBeingDecrypted() || event.isDecryptionFailure()) {
+                // add an event listener for once the event is decrypted.
+                event.once(MatrixEventEvent.Decrypted, async () => {
+                    processBeaconRelation(relatedToEventId, event);
+                });
+            } else {
+                processBeaconRelation(relatedToEventId, event);
             }
         });
     }