diff --git a/apps/integration-tests/src/end-to-end.integration.test.ts b/apps/integration-tests/src/end-to-end.integration.test.ts index eb98ba6..a7c59de 100644 --- a/apps/integration-tests/src/end-to-end.integration.test.ts +++ b/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((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; diff --git a/packages/core/CHANGELOG.md b/packages/core/CHANGELOG.md index 17b077b..1678344 100644 --- a/packages/core/CHANGELOG.md +++ b/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)) diff --git a/packages/core/src/base-client.integration.test.ts b/packages/core/src/base-client.integration.test.ts index 99af4a9..189af7d 100644 --- a/packages/core/src/base-client.integration.test.ts +++ b/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((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(); diff --git a/packages/core/src/base-client.ts b/packages/core/src/base-client.ts index ede243f..3628843 100644 --- a/packages/core/src/base-client.ts +++ b/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,7 +144,7 @@ export abstract class BaseClient extends EventEmitter { */ protected async sendMessage(channel: string, message: ProtocolMessage): Promise { 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); @@ -148,15 +152,14 @@ export abstract class BaseClient extends EventEmitter { } /** - * 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 { - 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 { + if (!this.session || this.session.expiresAt >= Date.now()) return false; + await this.disconnect(); + return true; } /**