Improve logging of event encryption in RustCrypto (#3898)

* Improve logging of event encryption in `RustCrypto`

* fix tests

spec/unit/rust-crypto/KeyClaimManager.spec.ts

@@ -23,6 +23,7 @@ import { OutgoingRequestProcessor } from "../../../src/rust-crypto/OutgoingReque
 import { KeyClaimManager } from "../../../src/rust-crypto/KeyClaimManager";
 import { TypedEventEmitter } from "../../../src/models/typed-event-emitter";
 import { HttpApiEvent, HttpApiEventHandlerMap, MatrixHttpApi } from "../../../src";
+import { logger, LogSpan } from "../../../src/logger";
 
 afterEach(() => {
     fetchMock.mockReset();
@@ -93,7 +94,7 @@ describe("KeyClaimManager", () => {
         olmMachine.markRequestAsSent.mockResolvedValueOnce(undefined);
 
         // fire off the request
-        await keyClaimManager.ensureSessionsForUsers([u1, u2]);
+        await keyClaimManager.ensureSessionsForUsers(new LogSpan(logger, "test"), [u1, u2]);
 
         // check that all the calls were made
         expect(olmMachine.getMissingSessions).toHaveBeenCalledWith([u1, u2]);
@@ -119,12 +120,13 @@ describe("KeyClaimManager", () => {
         let markRequestAsSentPromise = awaitCallToMarkRequestAsSent();
 
         // fire off two requests, and keep track of whether their promises resolve
+        const span = new LogSpan(logger, "test");
         let req1Resolved = false;
-        keyClaimManager.ensureSessionsForUsers([u1]).then(() => {
+        keyClaimManager.ensureSessionsForUsers(span, [u1]).then(() => {
             req1Resolved = true;
         });
         let req2Resolved = false;
-        const req2 = keyClaimManager.ensureSessionsForUsers([u2]).then(() => {
+        const req2 = keyClaimManager.ensureSessionsForUsers(span, [u2]).then(() => {
             req2Resolved = true;
         });
 

src/logger.ts

@@ -18,7 +18,18 @@ limitations under the License.
 import loglevel from "loglevel";
 
 /** Logger interface used within the js-sdk codebase */
-export interface Logger {
+export interface Logger extends BaseLogger {
+    /**
+     * Create a child logger.
+     *
+     * @param namespace - name to add to the current logger to generate the child. Some implementations of `Logger`
+     *    use this as a prefix; others use a different mechanism.
+     */
+    getChild(namespace: string): Logger;
+}
+
+/** The basic interface for a logger which doesn't support children */
+interface BaseLogger {
     /**
      * Output trace message to the logger, with stack trace.
      *
@@ -53,14 +64,6 @@ export interface Logger {
      * @param msg - Data to log.
      */
     error(...msg: any[]): void;
-
-    /**
-     * Create a child logger.
-     *
-     * @param namespace - name to add to the current logger to generate the child. Some implementations of `Logger`
-     *    use this as a prefix; others use a different mechanism.
-     */
-    getChild(namespace: string): Logger;
 }
 
 // This is to demonstrate, that you can use any namespace you want.
@@ -139,3 +142,41 @@ function getPrefixedLogger(prefix: string): PrefixedLogger {
 export const logger = loglevel.getLogger(DEFAULT_NAMESPACE) as PrefixedLogger;
 logger.setLevel(loglevel.levels.DEBUG, false);
 extendLogger(logger);
+
+/**
+ * A "span" for grouping related log lines together.
+ *
+ * The current implementation just adds the name at the start of each log line.
+ *
+ * This offers a lighter-weight alternative to 'child' loggers returned by {@link Logger#getChild}. In particular,
+ * it's not possible to apply individual filters to the LogSpan such as setting the verbosity level. On the other hand,
+ * no reference to the LogSpan is retained in the logging framework, so it is safe to make lots of them over the course
+ * of an application's life and just drop references to them when the job is done.
+ */
+export class LogSpan implements BaseLogger {
+    private readonly name;
+
+    public constructor(private readonly parent: BaseLogger, name: string) {
+        this.name = name + ":";
+    }
+
+    public trace(...msg: any[]): void {
+        this.parent.trace(this.name, ...msg);
+    }
+
+    public debug(...msg: any[]): void {
+        this.parent.debug(this.name, ...msg);
+    }
+
+    public info(...msg: any[]): void {
+        this.parent.info(this.name, ...msg);
+    }
+
+    public warn(...msg: any[]): void {
+        this.parent.warn(this.name, ...msg);
+    }
+
+    public error(...msg: any[]): void {
+        this.parent.error(this.name, ...msg);
+    }
+}

src/rust-crypto/KeyClaimManager.ts

@@ -17,6 +17,7 @@ limitations under the License.
 import { OlmMachine, UserId } from "@matrix-org/matrix-sdk-crypto-wasm";
 
 import { OutgoingRequestProcessor } from "./OutgoingRequestProcessor";
+import { LogSpan } from "../logger";
 
 /**
  * KeyClaimManager: linearises calls to OlmMachine.getMissingSessions to avoid races
@@ -52,7 +53,7 @@ export class KeyClaimManager {
      *
      * @param userList - list of userIDs to claim
      */
-    public ensureSessionsForUsers(userList: Array<UserId>): Promise<void> {
+    public ensureSessionsForUsers(logger: LogSpan, userList: Array<UserId>): Promise<void> {
         // The Rust-SDK requires that we only have one getMissingSessions process in flight at once. This little dance
         // ensures that, by only having one call to ensureSessionsForUsersInner active at once (and making them
         // queue up in order).
@@ -61,19 +62,22 @@ export class KeyClaimManager {
                 // any errors in the previous claim will have been reported already, so there is nothing to do here.
                 // we just throw away the error and start anew.
             })
-            .then(() => this.ensureSessionsForUsersInner(userList));
+            .then(() => this.ensureSessionsForUsersInner(logger, userList));
         this.currentClaimPromise = prom;
         return prom;
     }
 
-    private async ensureSessionsForUsersInner(userList: Array<UserId>): Promise<void> {
+    private async ensureSessionsForUsersInner(logger: LogSpan, userList: Array<UserId>): Promise<void> {
         // bail out quickly if we've been stopped.
         if (this.stopped) {
             throw new Error(`Cannot ensure Olm sessions: shutting down`);
         }
+        logger.info("Checking for missing Olm sessions");
         const claimRequest = await this.olmMachine.getMissingSessions(userList);
         if (claimRequest) {
+            logger.info("Making /keys/claim request");
             await this.outgoingRequestProcessor.makeOutgoingRequest(claimRequest);
         }
+        logger.info("Olm sessions prepared");
     }
 }

src/rust-crypto/RoomEncryptor.ts

@@ -28,7 +28,7 @@ import * as RustSdkCryptoJs from "@matrix-org/matrix-sdk-crypto-wasm";
 import { EventType } from "../@types/event";
 import { IContent, MatrixEvent } from "../models/event";
 import { Room } from "../models/room";
-import { Logger, logger } from "../logger";
+import { Logger, logger, LogSpan } from "../logger";
 import { KeyClaimManager } from "./KeyClaimManager";
 import { RoomMember } from "../models/room-member";
 import { HistoryVisibility } from "../@types/partials";
@@ -110,12 +110,27 @@ export class RoomEncryptor {
      *
      * @param globalBlacklistUnverifiedDevices - When `true`, it will not send encrypted messages to unverified devices
      */
-    public async ensureEncryptionSession(globalBlacklistUnverifiedDevices: boolean): Promise<void> {
+    public async prepareForEncryption(globalBlacklistUnverifiedDevices: boolean): Promise<void> {
+        const logger = new LogSpan(this.prefixedLogger, "prepareForEncryption");
+        await this.ensureEncryptionSession(logger, globalBlacklistUnverifiedDevices);
+    }
+
+    /**
+     * Prepare to encrypt events in this room.
+     *
+     * This ensures that we have a megolm session ready to use and that we have shared its key with all the devices
+     * in the room.
+     *
+     * @param logger - a place to write diagnostics to
+     * @param globalBlacklistUnverifiedDevices - When `true`, it will not send encrypted messages to unverified devices
+     */
+    private async ensureEncryptionSession(logger: LogSpan, globalBlacklistUnverifiedDevices: boolean): Promise<void> {
         if (this.encryptionSettings.algorithm !== "m.megolm.v1.aes-sha2") {
             throw new Error(
                 `Cannot encrypt in ${this.room.roomId} for unsupported algorithm '${this.encryptionSettings.algorithm}'`,
             );
         }
+        logger.debug("Starting encryption");
 
         const members = await this.room.getEncryptionTargetMembers();
 
@@ -128,7 +143,7 @@ export class RoomEncryptor {
         if (!this.lazyLoadedMembersResolved) {
             await this.olmMachine.updateTrackedUsers(members.map((u) => new RustSdkCryptoJs.UserId(u.userId)));
             this.lazyLoadedMembersResolved = true;
-            this.prefixedLogger.debug(`Updated tracked users for room ${this.room.roomId}`);
+            logger.debug(`Updated tracked users`);
         }
 
         // Query keys in case we don't have them for newly tracked members.
@@ -138,17 +153,16 @@ export class RoomEncryptor {
         // process the pending requests from the olmMachine. (usually these are processed
         // at the end of the sync, but we can't wait for that).
         // XXX future improvement process only KeysQueryRequests for the tracked users.
+        logger.debug(`Processing outgoing requests`);
         await this.outgoingRequestManager.doProcessOutgoingRequests();
 
-        this.prefixedLogger.debug(
+        logger.debug(
             `Encrypting for users (shouldEncryptForInvitedMembers: ${this.room.shouldEncryptForInvitedMembers()}):`,
             members.map((u) => `${u.userId} (${u.membership})`),
         );
 
         const userList = members.map((u) => new UserId(u.userId));
-        await this.keyClaimManager.ensureSessionsForUsers(userList);
-
-        this.prefixedLogger.debug("Sessions for users are ready; now sharing room key");
+        await this.keyClaimManager.ensureSessionsForUsers(logger, userList);
 
         const rustEncryptionSettings = new EncryptionSettings();
         rustEncryptionSettings.historyVisibility = toRustHistoryVisibility(this.room.getHistoryVisibility());
@@ -204,8 +218,10 @@ export class RoomEncryptor {
      * @param globalBlacklistUnverifiedDevices - When `true`, it will not send encrypted messages to unverified devices
      */
     public async encryptEvent(event: MatrixEvent, globalBlacklistUnverifiedDevices: boolean): Promise<void> {
-        await this.ensureEncryptionSession(globalBlacklistUnverifiedDevices);
+        const logger = new LogSpan(this.prefixedLogger, event.getTxnId() ?? "");
+        await this.ensureEncryptionSession(logger, globalBlacklistUnverifiedDevices);
 
+        logger.debug("Encrypting actual message content");
         const encryptedContent = await this.olmMachine.encryptRoomEvent(
             new RoomId(this.room.roomId),
             event.getType(),
@@ -218,6 +234,8 @@ export class RoomEncryptor {
             this.olmMachine.identityKeys.curve25519.toBase64(),
             this.olmMachine.identityKeys.ed25519.toBase64(),
         );
+
+        logger.debug("Encrypted event successfully");
     }
 }
 

src/rust-crypto/rust-crypto.ts

@@ -375,7 +375,7 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, RustCryptoEv
         const encryptor = this.roomEncryptors[room.roomId];
 
         if (encryptor) {
-            encryptor.ensureEncryptionSession(this.globalBlacklistUnverifiedDevices);
+            encryptor.prepareForEncryption(this.globalBlacklistUnverifiedDevices);
         }
     }