Handle self read receipts for fixing e2e notification counts

src/ReEmitter.js

@@ -34,12 +34,18 @@ export default class Reemitter {
     }
 
     reEmit(source, eventNames) {
+        // We include the source as the last argument for event handlers which may need it,
+        // such as read receipt listeners on the client class which won't have the context
+        // of the room.
+        const forSource = (handler, ...args) => {
+            handler(...args, source);
+        };
         for (const eventName of eventNames) {
             if (this.boundHandlers[eventName] === undefined) {
                 this.boundHandlers[eventName] = this._handleEvent.bind(this, eventName);
             }
-            const boundHandler = this.boundHandlers[eventName];
 
+            const boundHandler = forSource.bind(this, this.boundHandlers[eventName]);
             source.on(eventName, boundHandler);
         }
     }

src/client.js

@@ -278,6 +278,46 @@ function MatrixClient(opts) {
             }
         }
     });
+
+    // Like above, we have to listen for read receipts from ourselves in order to
+    // correctly handle notification counts on encrypted rooms.
+    // This fixes https://github.com/vector-im/riot-web/issues/9421
+    this.on("Room.receipt", (event, room) => {
+        if (room && this.isRoomEncrypted(room.roomId)) {
+            // Figure out if we've read something or if it's just informational
+            const content = event.getContent();
+            const isSelf = Object.keys(content).filter(eid => {
+                return Object.keys(content[eid]['m.read']).includes(this.getUserId());
+            }).length > 0;
+
+            if (!isSelf) return;
+
+            // Work backwards to determine how many events are unread. We also set
+            // a limit for how back we'll look to avoid spinning CPU for too long.
+            // If we hit the limit, we assume the count is unchanged.
+            const maxHistory = 20;
+            const events = room.getLiveTimeline().getEvents();
+
+            let highlightCount = 0;
+
+            for (let i = events.length - 1; i >= 0; i--) {
+                if (i === events.length - maxHistory) return; // limit reached
+
+                const event = events[i];
+
+                if (room.hasUserReadEvent(this.getUserId(), event.getId())) {
+                    // If the user has read the event, then the counting is done.
+                    break;
+                }
+
+                highlightCount += event.getPushActions().tweaks.highlight ? 1 : 0;
+            }
+
+            // Note: we don't need to handle 'total' notifications because the counts
+            // will come from the server.
+            room.setUnreadNotificationCount("highlight", highlightCount);
+        }
+    });
 }
 utils.inherits(MatrixClient, EventEmitter);
 utils.extend(MatrixClient.prototype, MatrixBaseApis.prototype);