@aztec/wallet-sdk 0.0.1-commit.fffb133c → 0.0.1-dev

package/src/extension/provider/extension_wallet.ts

@@ -6,7 +6,17 @@ import { schemaHasMethod } from '@aztec/foundation/schemas';
 import type { FunctionsOf } from '@aztec/foundation/types';
 
 import { type EncryptedPayload, decrypt, encrypt } from '../../crypto.js';
-import { type WalletMessage, WalletMessageType, type WalletResponse } from '../../types.js';
+import {
+  DEFAULT_HEARTBEAT_DEAD_AFTER_MS,
+  DEFAULT_HEARTBEAT_INTERVAL_MS,
+  type DisconnectCallback,
+  type HeartbeatOptions,
+  NOOP_LOGGER,
+  type WalletMessage,
+  WalletMessageType,
+  type WalletResponse,
+  type WalletSdkLogger,
+} from '../../types.js';
 
 /**
  * Internal type representing a wallet method call before encryption.
@@ -19,11 +29,6 @@ type WalletMethodCall = {
   args: unknown[];
 };
 
-/**
- * Callback type for wallet disconnect events.
- */
-export type DisconnectCallback = () => void;
-
 /**
  * A wallet implementation that communicates with browser extension wallets
  * using an encrypted MessageChannel.
@@ -60,6 +65,11 @@ export class ExtensionWallet {
   private inFlight = new Map<string, PromiseWithResolvers<unknown>>();
   private disconnected = false;
   private disconnectCallbacks: DisconnectCallback[] = [];
+  private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
+  private lastInboundAt = 0;
+  private log: WalletSdkLogger;
+  private heartbeatIntervalMs: number;
+  private heartbeatDeadAfterMs: number;
 
   /**
    * Private constructor - use {@link ExtensionWallet.create} to instantiate.
@@ -68,6 +78,8 @@ export class ExtensionWallet {
    * @param extensionId - The unique identifier of the target wallet extension
    * @param port - The MessagePort for private communication with the wallet
    * @param sharedKey - The derived AES-256-GCM shared key for encryption
+   * @param logger - Optional logger; defaults to a no-op logger
+   * @param heartbeatOptions - Optional heartbeat tuning (mostly useful for tests)
    */
   private constructor(
     private chainInfo: ChainInfo,
@@ -75,7 +87,13 @@ export class ExtensionWallet {
     private extensionId: string,
     private port: MessagePort,
     private sharedKey: CryptoKey,
-  ) {}
+    logger?: WalletSdkLogger,
+    heartbeatOptions?: HeartbeatOptions,
+  ) {
+    this.log = logger ?? NOOP_LOGGER;
+    this.heartbeatIntervalMs = heartbeatOptions?.intervalMs ?? DEFAULT_HEARTBEAT_INTERVAL_MS;
+    this.heartbeatDeadAfterMs = heartbeatOptions?.deadAfterMs ?? DEFAULT_HEARTBEAT_DEAD_AFTER_MS;
+  }
 
   /**
    * Creates a Wallet that communicates with a browser extension
@@ -86,6 +104,8 @@ export class ExtensionWallet {
    * @param sharedKey - The derived AES-256-GCM shared key for encryption
    * @param chainInfo - The chain information (chainId and version) for request context
    * @param appId - Application identifier used to identify the requesting dApp to the wallet
+   * @param logger - Optional logger; defaults to a no-op logger to keep extension/page bundles small
+   * @param heartbeatOptions - Optional override for heartbeat tuning (mostly useful for tests)
    * @returns A Wallet interface where all method calls are encrypted
    *
    * @example
@@ -109,13 +129,17 @@ export class ExtensionWallet {
     sharedKey: CryptoKey,
     chainInfo: ChainInfo,
     appId: string,
-  ): Wallet {
-    const wallet = new ExtensionWallet(chainInfo, appId, extensionId, port, sharedKey);
+    logger?: WalletSdkLogger,
+    heartbeatOptions?: HeartbeatOptions,
+  ): ExtensionWallet {
+    const wallet = new ExtensionWallet(chainInfo, appId, extensionId, port, sharedKey, logger, heartbeatOptions);
 
     // Set up message handler for encrypted responses and unencrypted control messages
     wallet.port.onmessage = (event: MessageEvent) => {
       const data = event.data;
-      // Check for unencrypted disconnect notification
+      // Any inbound traffic counts as proof of liveness.
+      wallet.lastInboundAt = Date.now();
+
       if (data && typeof data === 'object' && 'type' in data && data.type === WalletMessageType.DISCONNECT) {
         wallet.handleDisconnect();
         return;
@@ -127,8 +151,10 @@ export class ExtensionWallet {
     wallet.port.start();
 
     return new Proxy(wallet, {
-      get: (target, prop) => {
-        if (schemaHasMethod(WalletSchema, prop.toString())) {
+      get: (target, prop, receiver) => {
+        if (prop === 'asWallet') {
+          return () => receiver as unknown as Wallet;
+        } else if (schemaHasMethod(WalletSchema, prop.toString())) {
           return async (...args: unknown[]) => {
             const result = await target.postMessage({
               type: prop.toString() as keyof FunctionsOf<Wallet>,
@@ -140,7 +166,13 @@ export class ExtensionWallet {
           return target[prop as keyof ExtensionWallet];
         }
       },
-    }) as unknown as Wallet;
+    });
+  }
+
+  asWallet(): Wallet {
+    // Overridden by the proxy in create() to return the proxy itself.
+    // This body is never reached when accessed through create().
+    return this as unknown as Wallet;
   }
 
   /**
@@ -181,8 +213,10 @@ export class ExtensionWallet {
         resolve(result);
       }
       this.inFlight.delete(messageId);
-      // eslint-disable-next-line no-empty
-    } catch {}
+      this.maybeStopHeartbeat();
+    } catch (err) {
+      this.log.warn('Failed to decrypt wallet response', { err });
+    }
   }
 
   /**
@@ -220,9 +254,59 @@ export class ExtensionWallet {
 
     const { promise, resolve, reject } = promiseWithResolvers<unknown>();
     this.inFlight.set(messageId, { promise, resolve, reject });
+    this.startHeartbeat();
     return promise;
   }
 
+  /**
+   * Start the heartbeat probe loop while at least one request is in flight.
+   * Idempotent — calling while already running is a no-op.
+   *
+   * Heartbeat is opt-in via wire protocol: PINGs are unencrypted control messages
+   * (like DISCONNECT). Older wallets that do not understand PING simply drop it,
+   * which is safe — we only declare disconnect when **no** inbound traffic of any
+   * kind (PONG, encrypted response, DISCONNECT) arrives within the dead window.
+   * A wallet that is processing a slow request will reset the timer when it
+   * eventually responds, so this never causes false disconnects on legacy peers.
+   */
+  private startHeartbeat(): void {
+    if (this.heartbeatTimer !== null || this.disconnected) {
+      return;
+    }
+    this.lastInboundAt = Date.now();
+    this.heartbeatTimer = setInterval(() => this.heartbeatTick(), this.heartbeatIntervalMs);
+  }
+
+  private maybeStopHeartbeat(): void {
+    if (this.inFlight.size === 0 && this.heartbeatTimer !== null) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+  }
+
+  private heartbeatTick(): void {
+    if (this.disconnected || this.inFlight.size === 0) {
+      this.maybeStopHeartbeat();
+      return;
+    }
+
+    const idleMs = Date.now() - this.lastInboundAt;
+    if (idleMs >= this.heartbeatDeadAfterMs) {
+      this.log.warn('Wallet channel unresponsive — declaring disconnect', {
+        idleMs,
+        inFlight: this.inFlight.size,
+      });
+      this.handleDisconnect();
+      return;
+    }
+
+    try {
+      this.port.postMessage({ type: WalletMessageType.PING });
+    } catch (err) {
+      this.log.warn('Failed to send heartbeat PING', { err });
+    }
+  }
+
   /**
    * Handles wallet disconnection.
    * Rejects all pending requests and notifies registered callbacks.
@@ -234,6 +318,11 @@ export class ExtensionWallet {
     }
     this.disconnected = true;
 
+    if (this.heartbeatTimer !== null) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+
     if (this.port) {
       this.port.onmessage = null;
       this.port.close();

package/src/extension/provider/index.ts

@@ -1,4 +1,4 @@
-export { ExtensionWallet, type DisconnectCallback } from './extension_wallet.js';
+export { ExtensionWallet } from './extension_wallet.js';
 export {
   ExtensionProvider,
   type DiscoveredWallet,

package/src/iframe/handlers/iframe_connection_handler.ts

@@ -0,0 +1,342 @@
+/**
+ * IframeConnectionHandler — wallet-side of the cross-origin iframe protocol.
+ *
+ * This mirrors {@link BackgroundConnectionHandler} from `@aztec/wallet-sdk/extension/handlers`
+ * but uses `window.postMessage` instead of browser.runtime messaging.
+ *
+ * Message flow (wallet receives):
+ *   parent → DISCOVERY            → show approval UI → send DISCOVERY_RESPONSE
+ *   parent → KEY_EXCHANGE_REQUEST → ECDH             → send KEY_EXCHANGE_RESPONSE
+ *   parent → SECURE_MESSAGE       → decrypt → Wallet → encrypt → SECURE_RESPONSE
+ *   parent → DISCONNECT           → terminate session
+ *
+ * The wallet announces itself by posting WALLET_READY as soon as the handler starts,
+ * so the dApp knows it can send a discovery request.
+ */
+import type { ChainInfo } from '@aztec/aztec.js/account';
+import type { Wallet } from '@aztec/aztec.js/wallet';
+import { WalletSchema } from '@aztec/aztec.js/wallet';
+import { jsonStringify } from '@aztec/foundation/json-rpc';
+import { parseWithOptionals, schemaHasMethod } from '@aztec/foundation/schemas';
+
+import {
+  type EncryptedPayload,
+  decrypt,
+  deriveSessionKeys,
+  encrypt,
+  exportPublicKey,
+  generateKeyPair,
+  importPublicKey,
+} from '../../crypto.js';
+import { type WalletMessage, WalletMessageType, type WalletResponse, type WalletSdkLogger } from '../../types.js';
+
+/**
+ * A pending discovery request from a dApp (before user approval).
+ */
+export interface PendingSession {
+  /** Unique request identifier */
+  requestId: string;
+  /** Application identifier */
+  appId: string;
+  /** Origin URL of the requesting page */
+  origin: string;
+  /** Approval status */
+  status: 'pending' | 'approved';
+}
+
+/**
+ * An active session (after key exchange).
+ */
+export interface ActiveSession {
+  /** Session identifier (same as the discovery requestId) */
+  sessionId: string;
+  /** AES-256-GCM shared key for this session */
+  sharedKey: CryptoKey;
+  /** Verification hash for emoji display */
+  verificationHash: string;
+  /** Origin URL of the connected dApp */
+  origin: string;
+  /** Application identifier */
+  appId: string;
+}
+
+/**
+ * Configuration for the iframe connection handler.
+ */
+export interface IframeConnectionConfig {
+  /** Unique wallet identifier */
+  walletId: string;
+  /** Display name for the wallet */
+  walletName: string;
+  /** Wallet version string */
+  walletVersion: string;
+  /** Optional wallet icon URL */
+  walletIcon?: string;
+  /** Origins allowed to connect. If empty or undefined, all origins are allowed (dev mode). */
+  allowedOrigins?: string[];
+  /** Logger used for diagnostics. */
+  logger: WalletSdkLogger;
+}
+
+/**
+ * Event callbacks for the iframe connection handler.
+ */
+export interface IframeConnectionCallbacks {
+  /** Called when a new discovery request arrives — wallet can show approval UI */
+  onPendingDiscovery?: (session: PendingSession) => void;
+  /** Called when a session is established (key exchange complete) */
+  onSessionEstablished?: (session: ActiveSession) => void;
+  /** Called when a session is terminated */
+  onSessionTerminated?: (sessionId: string) => void;
+  /** Called when a key exchange completes — show verificationHash as emojis to the user */
+  onVerificationHash?: (verificationHash: string) => void;
+  /**
+   * Resolves the Wallet instance to use for a given dApp and chain.
+   * Called when an encrypted message arrives and needs to be dispatched.
+   */
+  getWallet: (appId: string, chainInfo: ChainInfo) => Promise<Wallet>;
+}
+
+/**
+ * Handles the wallet side of the cross-origin iframe protocol.
+ *
+ * Manages the full lifecycle: discovery, ECDH key exchange, encrypted message
+ * dispatch to a {@link Wallet} instance, and session termination.
+ *
+ * @example
+ * ```typescript
+ * const handler = new IframeConnectionHandler(
+ *   { walletId: 'my-wallet', walletName: 'My Wallet', walletVersion: '1.0.0', logger: console },
+ *   {
+ *     onPendingDiscovery: (session) => showApprovalUI(session),
+ *     getWallet: (appId, chainInfo) => createWalletForApp(appId, chainInfo),
+ *   },
+ * );
+ * handler.start();
+ * ```
+ */
+export class IframeConnectionHandler {
+  private pendingSessions = new Map<string, PendingSession>();
+  private activeSessions = new Map<string, ActiveSession>();
+  private log: WalletSdkLogger;
+
+  constructor(
+    private config: IframeConnectionConfig,
+    private callbacks: IframeConnectionCallbacks,
+  ) {
+    this.log = config.logger;
+  }
+
+  start(): void {
+    window.addEventListener('message', this.handleMessage);
+    this.postToParent({ type: WalletMessageType.WALLET_READY });
+    this.log.info('IframeConnectionHandler started, posted WALLET_READY');
+  }
+
+  stop(): void {
+    window.removeEventListener('message', this.handleMessage);
+  }
+
+  approveDiscovery(requestId: string): void {
+    const pending = this.pendingSessions.get(requestId);
+    if (!pending || pending.status !== 'pending') {
+      return;
+    }
+
+    pending.status = 'approved';
+    this.postToOrigin(pending.origin, {
+      type: WalletMessageType.DISCOVERY_RESPONSE,
+      requestId,
+      walletInfo: {
+        id: this.config.walletId,
+        name: this.config.walletName,
+        version: this.config.walletVersion,
+        icon: this.config.walletIcon,
+      },
+    });
+    this.log.info(`Discovery approved for requestId=${requestId}`);
+  }
+
+  rejectDiscovery(requestId: string): void {
+    this.pendingSessions.delete(requestId);
+  }
+
+  terminateSession(sessionId: string): void {
+    const session = this.activeSessions.get(sessionId);
+    if (session) {
+      this.postToOrigin(session.origin, {
+        type: WalletMessageType.SESSION_DISCONNECTED,
+        sessionId,
+      });
+      this.activeSessions.delete(sessionId);
+      this.callbacks.onSessionTerminated?.(sessionId);
+    }
+  }
+
+  getPendingSessions(): PendingSession[] {
+    return Array.from(this.pendingSessions.values()).filter(s => s.status === 'pending');
+  }
+
+  private handleMessage = (event: MessageEvent): void => {
+    void this.handleMessageAsync(event);
+  };
+
+  private async handleMessageAsync(event: MessageEvent): Promise<void> {
+    if (this.config.allowedOrigins && this.config.allowedOrigins.length > 0) {
+      if (!this.config.allowedOrigins.includes(event.origin)) {
+        return;
+      }
+    }
+
+    const msg = event.data;
+    if (!msg || typeof msg !== 'object' || !msg.type) {
+      return;
+    }
+
+    switch (msg.type) {
+      case WalletMessageType.DISCOVERY:
+        this.handleDiscoveryRequest(msg, event.origin);
+        break;
+      case WalletMessageType.KEY_EXCHANGE_REQUEST:
+        await this.handleKeyExchangeRequest(msg, event.origin);
+        break;
+      case WalletMessageType.SECURE_MESSAGE:
+        await this.handleSecureMessage(msg);
+        break;
+      case WalletMessageType.DISCONNECT:
+        this.terminateSession(msg.sessionId);
+        break;
+      case WalletMessageType.PING:
+        this.handlePing(msg.sessionId);
+        break;
+    }
+  }
+
+  private handlePing(sessionId: string): void {
+    const session = this.activeSessions.get(sessionId);
+    if (!session) {
+      return;
+    }
+    this.postToOrigin(session.origin, { type: WalletMessageType.PONG, sessionId });
+  }
+
+  private handleDiscoveryRequest(msg: Record<string, unknown>, origin: string): void {
+    // eslint-disable-next-line jsdoc/require-jsdoc
+    const { requestId, appId } = msg as { requestId: string; appId: string };
+    const pending: PendingSession = { requestId, appId, origin, status: 'pending' };
+    this.pendingSessions.set(requestId, pending);
+    this.log.info(`Discovery request from appId=${appId} origin=${origin}`);
+    this.callbacks.onPendingDiscovery?.(pending);
+  }
+
+  private async handleKeyExchangeRequest(msg: Record<string, unknown>, origin: string): Promise<void> {
+    const { requestId, publicKey: appPublicKeyRaw } = msg as {
+      // eslint-disable-next-line jsdoc/require-jsdoc
+      requestId: string;
+      // eslint-disable-next-line jsdoc/require-jsdoc
+      publicKey: { kty: string; crv: string; x: string; y: string };
+    };
+    const pending = this.pendingSessions.get(requestId);
+    if (!pending || pending.status !== 'approved') {
+      this.log.warn(`Key exchange for unknown/unapproved requestId=${requestId}`);
+      return;
+    }
+
+    try {
+      const keyPair = await generateKeyPair();
+      const walletPublicKey = await exportPublicKey(keyPair.publicKey);
+      const appPublicKey = await importPublicKey(appPublicKeyRaw);
+      const sessionKeys = await deriveSessionKeys(keyPair, appPublicKey, false);
+
+      const session: ActiveSession = {
+        sessionId: requestId,
+        sharedKey: sessionKeys.encryptionKey,
+        verificationHash: sessionKeys.verificationHash,
+        origin: pending.origin,
+        appId: pending.appId,
+      };
+
+      this.activeSessions.set(requestId, session);
+      this.pendingSessions.delete(requestId);
+
+      this.postToOrigin(origin, {
+        type: WalletMessageType.KEY_EXCHANGE_RESPONSE,
+        requestId,
+        publicKey: walletPublicKey,
+        verificationHash: sessionKeys.verificationHash,
+      });
+
+      this.callbacks.onVerificationHash?.(sessionKeys.verificationHash);
+      this.callbacks.onSessionEstablished?.(session);
+      this.log.info(`Key exchange complete, sessionId=${requestId}`);
+    } catch (err) {
+      this.log.error(`Key exchange failed: ${err}`);
+    }
+  }
+
+  private async handleSecureMessage(msg: Record<string, unknown>): Promise<void> {
+    // eslint-disable-next-line jsdoc/require-jsdoc
+    const { sessionId, encrypted } = msg as { sessionId: string; encrypted: EncryptedPayload };
+    const session = this.activeSessions.get(sessionId);
+    if (!session) {
+      return;
+    }
+
+    let walletMessage: WalletMessage;
+    try {
+      walletMessage = await decrypt<WalletMessage>(session.sharedKey, encrypted);
+    } catch {
+      this.log.warn(`Decryption failed for sessionId=${sessionId}`);
+      return;
+    }
+
+    const { messageId, type, args, chainInfo, appId } = walletMessage;
+
+    let result: unknown;
+    let error: string | undefined;
+
+    try {
+      const wallet = await this.callbacks.getWallet(appId, chainInfo);
+
+      if (!schemaHasMethod(WalletSchema, type)) {
+        throw new Error(`Unknown wallet method: ${type}`);
+      }
+      // Zod's AnyZodTuple rejects optional tuple items typed as `T | undefined`
+      const sanitizedArgs = await parseWithOptionals(args, WalletSchema[type].parameters() as any);
+      result = await (wallet as Record<string, (...a: unknown[]) => Promise<unknown>>)[type](...sanitizedArgs);
+    } catch (err: unknown) {
+      error = err instanceof Error ? err.message : String(err);
+      this.log.error(`Error handling ${type}: ${error}`);
+    }
+
+    const response: WalletResponse = {
+      messageId,
+      walletId: this.config.walletId,
+      result,
+      error,
+    };
+
+    try {
+      const encryptedResponse = await encrypt(session.sharedKey, jsonStringify(response));
+      this.postToOrigin(session.origin, {
+        type: WalletMessageType.SECURE_RESPONSE,
+        sessionId,
+        encrypted: encryptedResponse,
+      });
+    } catch (err) {
+      this.log.error(`Encryption of response failed: ${err}`);
+    }
+  }
+
+  private postToParent(msg: object): void {
+    if (window.parent !== window) {
+      window.parent.postMessage(msg, '*');
+    }
+  }
+
+  private postToOrigin(origin: string, msg: object): void {
+    if (window.parent !== window) {
+      window.parent.postMessage(msg, origin);
+    }
+  }
+}

package/src/iframe/handlers/index.ts

@@ -0,0 +1,7 @@
+export {
+  IframeConnectionHandler,
+  type IframeConnectionConfig,
+  type IframeConnectionCallbacks,
+  type PendingSession,
+  type ActiveSession,
+} from './iframe_connection_handler.js';