Implement MAB's UserActivityTracker and blocking the bridge upon exce…

…eding the user limit

src/bridge/IrcBridge.ts

@@ -32,6 +32,7 @@ import {
     EphemeralEvent,
     MembershipQueue,
     BridgeInfoStateSyncer,
+    UserActivityState,
 } from "matrix-appservice-bridge";
 import { IrcAction } from "../models/IrcAction";
 import { DataStore } from "../datastore/DataStore";
@@ -91,6 +92,7 @@ export class IrcBridge {
         networkId: string;
     }>;
     private privacyProtection: PrivacyProtection;
+    private isBlocked = false;
 
     constructor(public readonly config: BridgeConfig, private registration: AppServiceRegistration) {
         // TODO: Don't log this to stdout
@@ -127,8 +129,9 @@ export class IrcBridge {
         if (this.config.database.engine === "nedb") {
             const dirPath = this.config.database.connectionString.substring("nedb://".length);
             bridgeStoreConfig = {
-                roomStore: `${dirPath}/rooms.db`,
-                userStore: `${dirPath}/users.db`,
+                roomStore:         `${dirPath}/rooms.db`,
+                userStore:         `${dirPath}/users.db`,
+                userActivityStore: `${dirPath}/user-activity.db`,
             };
         }
         else {
@@ -159,6 +162,8 @@ export class IrcBridge {
                     getLocation: this.getThirdPartyLocation.bind(this),
                     getUser: this.getThirdPartyUser.bind(this),
                 },
+                getUserActivity: () => this.getStore().getUserActivity(),
+                onUserActivityChanged: this.onUserActivityChanged.bind(this),
             },
             ...bridgeStoreConfig,
             disableContext: true,
@@ -570,6 +575,7 @@ export class IrcBridge {
             }
             this.dataStore = new NeDBDataStore(
                 userStore,
+                this.bridge.getUserActivityStore()!,
                 roomStore,
                 this.config.homeserver.domain,
                 pkeyPath,
@@ -734,6 +740,9 @@ export class IrcBridge {
         });
 
         log.info("Startup complete.");
+
+        this.checkLimits(this.bridge.getUserActivityTracker()!.countActiveUsers().allUsers);
+
         this.startedUp = true;
     }
 
@@ -839,6 +848,10 @@ export class IrcBridge {
     }
 
     public async sendMatrixAction(room: MatrixRoom, from: MatrixUser, action: MatrixAction): Promise<void> {
+        if (this.isBlocked) {
+            log.info("Bridge is blocked, dropping Matrix action");
+            return;
+        }
         const intent = this.bridge.getIntent(from.userId);
         const extraContent: Record<string, unknown> = {};
         if (action.replyEvent) {
@@ -947,6 +960,10 @@ export class IrcBridge {
     }
 
     public onEvent(request: BridgeRequestEvent) {
+        if (this.isBlocked) {
+            log.info("Bridge is blocked, dropping Matrix event");
+            return;
+        }
         request.outcomeFrom(this._onEvent(request));
     }
 
@@ -1289,6 +1306,10 @@ export class IrcBridge {
     }
 
     public async sendIrcAction(ircRoom: IrcRoom, bridgedClient: BridgedClient, action: IrcAction) {
+        if (this.isBlocked) {
+            log.info("Bridge is blocked, dropping IRC action");
+            return;
+        }
         log.info(
             "Sending IRC message in %s as %s (connected=%s)",
             ircRoom.channel, bridgedClient.nick, Boolean(bridgedClient.status === BridgedClientStatus.CONNECTED)
@@ -1470,4 +1491,29 @@ export class IrcBridge {
         const current = await this.dataStore.getRoomCount();
         return current >= limit;
     }
+
+    private checkLimits(rmau: number) {
+        log.debug(`Bridge now serving ${rmau} RMAU`);
+
+        const limit = this.config.ircService.RMAUlimit;
+        if (!limit) return;
+        if (rmau > limit) {
+            if (!this.isBlocked) {
+                this.isBlocked = true;
+                log.info(`Bridge has reached the user limit of ${limit} and is now blocked`);
+            }
+        } else {
+            if (this.isBlocked) {
+                this.isBlocked = false;
+                log.info(`Bridge has has gone below the user limit of ${limit} and is now unblocked`);
+            }
+        }
+    }
+
+    private onUserActivityChanged(userActivity: UserActivityState) {
+        for (const userId of userActivity.changed) {
+            this.getStore().storeUserActivity(userId, userActivity.dataSet.users[userId]);
+        }
+        this.checkLimits(userActivity.activeUsers);
+    }
 }

src/config/BridgeConfig.ts

@@ -58,6 +58,7 @@ export interface BridgeConfig {
             [userIdOrDomain: string]: "admin";
         };
         perRoomConfig?: RoomConfigConfig;
+        RMAUlimit?: number;
     };
     sentry?: {
         enabled: boolean;

src/datastore/DataStore.ts

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-import { MatrixRoom, MatrixUser, RoomBridgeStoreEntry as Entry} from "matrix-appservice-bridge";
+import { MatrixRoom, MatrixUser, RoomBridgeStoreEntry as Entry, UserActivity, UserActivitySet } from "matrix-appservice-bridge";
 import Bluebird from "bluebird";
 import { IrcRoom } from "../models/IrcRoom";
 import { IrcClientConfig } from "../models/IrcClientConfig";
@@ -158,6 +158,10 @@ export interface DataStore {
 
     storeUserFeatures(userId: string, features: UserFeatures): Promise<void>;
 
+    getUserActivity(): Promise<UserActivitySet>;
+
+    storeUserActivity(userId: string, activity: UserActivity): Promise<void>;
+
     storePass(userId: string, domain: string, pass: string): Promise<void>;
 
     removePass(userId: string, domain: string): Promise<void>;

src/datastore/NedbDataStore.ts

@@ -20,7 +20,7 @@ import { IrcClientConfig, IrcClientConfigSeralized } from "../models/IrcClientCo
 import { getLogger } from "../logging";
 
 import { MatrixRoom, MatrixUser, RemoteUser, RemoteRoom,
-    UserBridgeStore, RoomBridgeStore, RoomBridgeStoreEntry as Entry } from "matrix-appservice-bridge";
+    UserBridgeStore, UserActivityStore, RoomBridgeStore, RoomBridgeStoreEntry as Entry, UserActivity, UserActivitySet } from "matrix-appservice-bridge";
 import { DataStore, RoomOrigin, ChannelMappings, UserFeatures } from "./DataStore";
 import { IrcServer, IrcServerConfig } from "../irc/IrcServer";
 import { StringCrypto } from "./StringCrypto";
@@ -36,6 +36,7 @@ export class NeDBDataStore implements DataStore {
     private cryptoStore?: StringCrypto;
     constructor(
         private userStore: UserBridgeStore,
+        private userActivityStore: UserActivityStore,
         private roomStore: RoomBridgeStore,
         private bridgeDomain: string,
         pkeyPath?: string) {
@@ -590,6 +591,14 @@ export class NeDBDataStore implements DataStore {
         await this.userStore.setMatrixUser(matrixUser);
     }
 
+    public async getUserActivity(): Promise<UserActivitySet> {
+        return this.userActivityStore.getActivitySet();
+    }
+
+    public async storeUserActivity(userId: string, activity: UserActivity) {
+        this.userActivityStore.storeUserActivity(userId, activity);
+    }
+
     public async storePass(userId: string, domain: string, pass: string) {
         const config = await this.getIrcClientConfig(userId, domain);
         if (!config) {

src/datastore/postgres/PgDataStore.ts

@@ -21,7 +21,9 @@ import {
     MatrixRoom,
     RemoteRoom,
     RoomBridgeStoreEntry as Entry,
-    MatrixRoomData
+    MatrixRoomData,
+    UserActivitySet,
+    UserActivity,
 } from "matrix-appservice-bridge";
 import { DataStore, RoomOrigin, ChannelMappings, UserFeatures } from "../DataStore";
 import { IrcRoom } from "../../models/IrcRoom";
@@ -52,7 +54,7 @@ interface RoomRecord {
 export class PgDataStore implements DataStore {
     private serverMappings: {[domain: string]: IrcServer} = {};
 
-    public static readonly LATEST_SCHEMA = 6;
+    public static readonly LATEST_SCHEMA = 7;
     private pgPool: Pool;
     private hasEnded = false;
     private cryptoStore?: StringCrypto;
@@ -555,6 +557,24 @@ export class PgDataStore implements DataStore {
         await this.pgPool.query(statement, [userId, JSON.stringify(features)]);
     }
 
+    public async getUserActivity(): Promise<UserActivitySet> {
+        const res = await this.pgPool.query('SELECT * FROM user_activity');
+        const users: {[mxid: string]: any} = {};
+        for (const row of res.rows) {
+            users[row['user_id']] = row['data'];
+        }
+        return { users };
+    }
+
+    public async storeUserActivity(userId: string, activity: UserActivity) {
+        const stmt = PgDataStore.BuildUpsertStatement(
+            'user_activity',
+            '(user_id)',
+            [ 'user_id', 'data' ],
+        );
+        await this.pgPool.query(stmt, [userId, JSON.stringify(activity)]);
+    }
+
     public async storePass(userId: string, domain: string, pass: string, encrypt = true): Promise<void> {
         let password = pass;
         if (encrypt) {

src/datastore/postgres/schema/v7.ts

@@ -0,0 +1,10 @@
+import { PoolClient } from "pg";
+
+export async function runSchema(connection: PoolClient) {
+    await connection.query(`
+    CREATE TABLE user_activity (
+        user_id TEXT UNIQUE,
+        data JSON
+    );
+    `);
+}