From 30e186ca9c2c238a4ee11bd9efe535679b445146 Mon Sep 17 00:00:00 2001
From: Travis Ralston <travpc@gmail.com>
Date: Wed, 6 Nov 2019 18:54:59 -0700
Subject: [PATCH] Add a config option to reduce homeserver load during ban
 checks

---
 config/default.yaml     |  6 ++++++
 src/actions/ApplyBan.ts | 28 ++++++++++++++++++----------
 src/config.ts           |  1 +
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/config/default.yaml b/config/default.yaml
index 81fefef..6cfb2a8 100644
--- a/config/default.yaml
+++ b/config/default.yaml
@@ -51,6 +51,12 @@ verifyPermissionsOnStartup: true
 # bot might not be trusted fully, yet. Default false (do bans/ACLs).
 noop: false
 
+# Set to true to use /joined_members instead of /state to figure out who is
+# in the room. Using /state is preferred because it means that users are
+# banned when they are invited instead of just when they join, though if your
+# server struggles with /state requests then set this to true.
+fasterMembershipChecks: false
+
 # A list of rooms to protect (matrix.to URLs)
 protectedRooms:
   - "https://matrix.to/#/#yourroom:example.org"
diff --git a/src/actions/ApplyBan.ts b/src/actions/ApplyBan.ts
index 6d4b991..3e17722 100644
--- a/src/actions/ApplyBan.ts
+++ b/src/actions/ApplyBan.ts
@@ -37,30 +37,38 @@ export async function applyUserBans(lists: BanList[], roomIds: string[], mjolnir
             // We specifically use sendNotice to avoid having to escape HTML
             await logMessage(LogLevel.DEBUG, "ApplyBan", `Updating member bans in ${roomId}`);
 
-            const state = await mjolnir.client.getRoomState(roomId);
-            const members = state.filter(s => s['type'] === 'm.room.member' && !!s['state_key']);
+            let members: { userId: string, membership: string }[];
+
+            if (config.fasterMembershipChecks) {
+                const memberIds = await mjolnir.client.getJoinedRoomMembers(roomId);
+                members = memberIds.map(u => {
+                    return {userId: u, membership: "join"};
+                });
+            } else {
+                const state = await mjolnir.client.getRoomState(roomId);
+                members = state.filter(s => s['type'] === 'm.room.member' && !!s['state_key']).map(s => {
+                    return {userId: s['state_key'], membership: s['content'] ? s['content']['membership'] : 'leave'};
+                });
+            }
 
             for (const member of members) {
-                const content = member['content'];
-                if (!content) continue; // Invalid, but whatever.
-
-                if (content['membership'] === 'ban') {
+                if (member.membership === 'ban') {
                     continue; // user already banned
                 }
 
                 let banned = false;
                 for (const list of lists) {
                     for (const userRule of list.userRules) {
-                        if (userRule.isMatch(member['state_key'])) {
+                        if (userRule.isMatch(member.userId)) {
                             // User needs to be banned
 
                             // We specifically use sendNotice to avoid having to escape HTML
-                            await logMessage(LogLevel.DEBUG, "ApplyBan", `Banning ${member['state_key']} in ${roomId} for: ${userRule.reason}`);
+                            await logMessage(LogLevel.DEBUG, "ApplyBan", `Banning ${member.userId} in ${roomId} for: ${userRule.reason}`);
 
                             if (!config.noop) {
-                                await mjolnir.client.banUser(member['state_key'], roomId, userRule.reason);
+                                await mjolnir.client.banUser(member.userId, roomId, userRule.reason);
                             } else {
-                                await logMessage(LogLevel.WARN, "ApplyBan", `Tried to ban ${member['state_key']} in ${roomId} but Mjolnir is running in no-op mode`);
+                                await logMessage(LogLevel.WARN, "ApplyBan", `Tried to ban ${member.userId} in ${roomId} but Mjolnir is running in no-op mode`);
                             }
 
                             bansApplied++;
diff --git a/src/config.ts b/src/config.ts
index 50a6b2e..2fffbea 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -34,6 +34,7 @@ interface IConfig {
     verifyPermissionsOnStartup: boolean;
     noop: boolean;
     protectedRooms: string[]; // matrix.to urls
+    fasterMembershipChecks: boolean;
 
     /**
      * Config options only set at runtime. Try to avoid using the objects