fix: reject inbound messages on expired sessions (WAPI-1130) (#72)

Author: chakra-guy

apps/integration-tests/src/end-to-end.integration.test.ts

@@ -1,6 +1,6 @@
 /** biome-ignore-all lint/suspicious/noExplicitAny: test code */
 /** biome-ignore-all lint/suspicious/noShadowRestrictedNames: test code */
-import { type ConnectionMode, type IKeyManager, type IKVStore, type KeyPair, type SessionRequest, SessionStore, WebSocketTransport } from "@metamask/mobile-wallet-protocol-core";
+import { type ConnectionMode, ErrorCode, type IKeyManager, type IKVStore, type KeyPair, type SessionRequest, SessionStore, WebSocketTransport } from "@metamask/mobile-wallet-protocol-core";
 import { DappClient, type OtpRequiredPayload } from "@metamask/mobile-wallet-protocol-dapp-client";
 import { WalletClient } from "@metamask/mobile-wallet-protocol-wallet-client";
 import { decrypt, encrypt, PrivateKey, PublicKey } from "eciesjs";
@@ -190,6 +190,34 @@ t.describe("E2E Integration Test", () => {
 		await t.expect(messageFromWalletPromise).resolves.toEqual(responsePayload);
 	});
 
+	t.test("should reject inbound messages on an expired session", async () => {
+		await connectClients(dappClient, walletClient, "trusted");
+
+		// Verify a message works pre-expiry
+		const preExpiryPayload = { method: "pre_expiry_check" };
+		const preExpiryPromise = new Promise((resolve) => walletClient.on("message", resolve));
+		await dappClient.sendRequest(preExpiryPayload);
+		await t.expect(preExpiryPromise).resolves.toEqual(preExpiryPayload);
+
+		// Force-expire the wallet's session
+		(walletClient as any).session.expiresAt = Date.now() - 1000;
+
+		const errorPromise = new Promise<any>((resolve) => {
+			walletClient.once("error", resolve);
+		});
+
+		// Send another message from dapp
+		await dappClient.sendRequest({ method: "post_expiry_check" });
+
+		// Wallet should emit SESSION_EXPIRED
+		const error = await errorPromise;
+		t.expect(error.code).toBe(ErrorCode.SESSION_EXPIRED);
+
+		// Wait briefly, confirm the message was NOT delivered
+		const walletMessagePromise = new Promise((resolve) => walletClient.once("message", resolve));
+		await assertPromiseNotResolve(walletMessagePromise, 500, "Wallet should not receive messages on expired session");
+	});
+
 	t.test("should successfully resume a previously established session", async () => {
 		await connectClients(dappClient, walletClient, "untrusted");
 		const sessionId = (await dappSessionStore.list())[0].id;

packages/core/CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Reject inbound messages on expired sessions instead of processing them
 - Fix `SessionStore` race conditions and fire-and-forget garbage collection ([#71](https://github.com/MetaMask/mobile-wallet-protocol/pull/71))
 - Guard against `NaN` in session expiry timestamps ([#70](https://github.com/MetaMask/mobile-wallet-protocol/pull/70))
 

packages/core/src/base-client.integration.test.ts

@@ -7,6 +7,7 @@ import * as t from "vitest";
 import WebSocket from "ws";
 import { BaseClient } from "./base-client";
 import { ClientState } from "./domain/client-state";
+import { ErrorCode } from "./domain/errors";
 import type { IKeyManager } from "./domain/key-manager";
 import type { KeyPair } from "./domain/key-pair";
 import type { IKVStore } from "./domain/kv-store";
@@ -368,6 +369,49 @@ t.describe("BaseClient", () => {
 		publishSpy.mockRestore();
 	});
 
+	t.test("should reject inbound messages on an expired session", async () => {
+		const keyManagerA = new KeyManager();
+		const keyManagerB = new KeyManager();
+		const keyPairA = keyManagerA.generateKeyPair();
+		const keyPairB = keyManagerB.generateKeyPair();
+
+		const sessionA: Session = {
+			id: "session-inbound-expiry",
+			channel,
+			keyPair: keyPairA,
+			theirPublicKey: keyPairB.publicKey,
+			expiresAt: Date.now() + 60000,
+		};
+		const sessionB: Session = {
+			id: "session-inbound-expiry",
+			channel,
+			keyPair: keyPairB,
+			theirPublicKey: keyPairA.publicKey,
+			expiresAt: Date.now() - 1000, // Already expired
+		};
+
+		clientA.setSession(sessionA);
+		clientB.setSession(sessionB);
+
+		await clientA["transport"].subscribe(channel);
+		await clientB["transport"].subscribe(channel);
+
+		const errorPromise = new Promise<any>((resolve) => {
+			clientB.once("error", resolve);
+		});
+
+		const messageToSend: ProtocolMessage = { type: "message", payload: { method: "should_be_rejected" } };
+		await clientA.sendMessage(channel, messageToSend);
+
+		const error = await errorPromise;
+		t.expect(error.code).toBe(ErrorCode.SESSION_EXPIRED);
+
+		// Give a small window to ensure no message processing occurs
+		await new Promise((resolve) => setTimeout(resolve, 200));
+		t.expect(clientB.receivedMessages).toHaveLength(0);
+		t.expect(clientB.getSession()).toBeNull();
+	});
+
 	t.test("should reject resume() when client is already connected", async () => {
 		// 1. Create and store a valid session
 		const keyManagerA = new KeyManager();

packages/core/src/base-client.ts

@@ -45,6 +45,10 @@ export abstract class BaseClient extends EventEmitter {
 
 		this.transport.on("message", async (payload) => {
 			if (!this.session?.keyPair.privateKey) return;
+			if (await this.checkSessionExpiry()) {
+				this.emit("error", new SessionError(ErrorCode.SESSION_EXPIRED, "Session expired"));
+				return;
+			}
 			const message = await this.decryptMessage(payload.data);
 			if (message) this.handleMessage(message);
 		});
@@ -140,23 +144,22 @@ export abstract class BaseClient extends EventEmitter {
 	 */
 	protected async sendMessage(channel: string, message: ProtocolMessage): Promise<void> {
 		if (!this.session) throw new SessionError(ErrorCode.SESSION_INVALID_STATE, "Cannot send message: session is not initialized.");
-		await this.checkSessionExpiry();
+		if (await this.checkSessionExpiry()) throw new SessionError(ErrorCode.SESSION_EXPIRED, "Session expired");
 		const plaintext = JSON.stringify(message);
 		const encrypted = await this.keymanager.encrypt(plaintext, this.session.theirPublicKey);
 		const ok = await this.transport.publish(channel, encrypted);
 		if (!ok) throw new TransportError(ErrorCode.TRANSPORT_DISCONNECTED, "Message could not be sent because the transport is disconnected.");
 	}
 
 	/**
-	 * Checks if the current session is expired. If it is, triggers a disconnect.
-	 * @throws {SessionError} if the session is expired.
+	 * Checks if the current session has expired. If so, triggers a disconnect.
+	 *
+	 * @returns true if the session was expired (and cleanup was triggered), false otherwise.
 	 */
-	private async checkSessionExpiry(): Promise<void> {
-		if (!this.session) return;
-		if (this.session.expiresAt < Date.now()) {
-			await this.disconnect();
-			throw new SessionError(ErrorCode.SESSION_EXPIRED, "Session expired");
-		}
+	private async checkSessionExpiry(): Promise<boolean> {
+		if (!this.session || this.session.expiresAt >= Date.now()) return false;
+		await this.disconnect();
+		return true;
 	}
 
 	/**