From 1b47999e80c8b7ab62288441244985702546e019 Mon Sep 17 00:00:00 2001
From: Travis Ralston <travpc@gmail.com>
Date: Tue, 2 Jul 2019 13:12:29 -0600
Subject: [PATCH] Handle self read receipts for fixing e2e notification counts

Fixes https://github.com/vector-im/riot-web/issues/9421

This also adds a context to the ReEmitter so we have access to the Room at the time of read receipt. Without this, we have to bind handlers to every encrypted room (which is tedious to maintain) or figure out which room `$something` belong to (CPU intensive).
---
 src/ReEmitter.js |  8 +++++++-
 src/client.js    | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/ReEmitter.js b/src/ReEmitter.js
index 51a4ab3b21d..4ddf3229a14 100644
--- a/src/ReEmitter.js
+++ b/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);
         }
     }
diff --git a/src/client.js b/src/client.js
index b2cdcd47764..9b3ade274ab 100644
--- a/src/client.js
+++ b/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);