@aztec/wallet-sdk 0.0.1-commit.d431d1c → 0.0.1-commit.d58ff9d0

package/src/{providers/extension → extension/provider}/extension_wallet.ts

@@ -2,11 +2,21 @@ import type { ChainInfo } from '@aztec/aztec.js/account';
 import { type Wallet, WalletSchema } from '@aztec/aztec.js/wallet';
 import { jsonStringify } from '@aztec/foundation/json-rpc';
 import { type PromiseWithResolvers, promiseWithResolvers } from '@aztec/foundation/promise';
-import { schemaHasMethod } from '@aztec/foundation/schemas';
+import { getSchemaReturnType, schemaHasMethod } from '@aztec/foundation/schemas';
 import type { FunctionsOf } from '@aztec/foundation/types';
 
 import { type EncryptedPayload, decrypt, encrypt } from '../../crypto.js';
-import { type WalletInfo, 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,37 +29,32 @@ type WalletMethodCall = {
   args: unknown[];
 };
 
-/**
- * Callback type for wallet disconnect events.
- */
-export type DisconnectCallback = () => void;
-
 /**
  * A wallet implementation that communicates with browser extension wallets
- * using a secure encrypted MessageChannel.
+ * using an encrypted MessageChannel.
  *
- * This class uses a pre-established secure channel from the discovery phase:
+ * This class uses a secure channel established after discovery:
  *
- * 1. **MessageChannel**: A private communication channel created during discovery,
- *    not visible to other scripts on the page (unlike window.postMessage).
+ * 1. **MessageChannel**: Created during discovery and transferred via window.postMessage.
+ *    Note: The port transfer is visible to page scripts, but security comes from encryption.
  *
- * 2. **ECDH Key Exchange**: The shared secret was derived during discovery using
- *    Elliptic Curve Diffie-Hellman key exchange.
+ * 2. **ECDH Key Exchange**: The shared secret was derived after discovery using
+ *    Elliptic Curve Diffie-Hellman key exchange over the MessagePort.
  *
  * 3. **AES-GCM Encryption**: All messages are encrypted using AES-256-GCM,
- *    providing both confidentiality and authenticity.
+ *    providing both confidentiality and authenticity. This is what secures the channel.
  *
  * @example
  * ```typescript
- * // Discovery returns wallets with secure channel components
- * const wallets = await ExtensionProvider.discoverExtensions(chainInfo);
- * const { info, port, sharedKey } = wallets[0];
+ * // Discover and establish secure channel to a wallet
+ * const discoveredWallets = await ExtensionProvider.discoverWallets(chainInfo, { appId: 'my-dapp' });
+ * const connection = await discoveredWallets[0].establishSecureChannel();
  *
  * // User can verify emoji if desired
- * console.log('Verify:', hashToEmoji(info.verificationHash!));
+ * console.log('Verify:', hashToEmoji(connection.info.verificationHash!));
  *
- * // Create wallet using the discovered components
- * const wallet = await ExtensionWallet.create(info, chainInfo, port, sharedKey, 'my-dapp');
+ * // Create wallet using the connection
+ * const wallet = ExtensionWallet.create(connection.info.id, connection.port, connection.sharedKey, chainInfo, 'my-dapp');
  *
  * // All subsequent calls are encrypted
  * const accounts = await wallet.getAccounts();
@@ -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,99 +87,92 @@ export class ExtensionWallet {
     private extensionId: string,
     private port: MessagePort,
     private sharedKey: CryptoKey,
-  ) {}
-
-  /** Cached Wallet proxy instance */
-  private walletProxy: Wallet | null = null;
+    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 an ExtensionWallet instance that communicates with a browser extension
+   * Creates a Wallet that communicates with a browser extension
    * over a secure encrypted MessageChannel.
    *
-   * @param walletInfo - The wallet info from ExtensionProvider.discoverExtensions()
-   * @param chainInfo - The chain information (chainId and version) for request context
-   * @param port - The MessagePort for private communication with the wallet
+   * @param extensionId - The unique identifier of the wallet extension
+   * @param port - The MessagePort for encrypted communication with the wallet
    * @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
-   * @returns The ExtensionWallet instance. Use {@link getWallet} to get the Wallet interface.
+   * @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
    * ```typescript
-   * const wallets = await ExtensionProvider.discoverExtensions(chainInfo);
-   * const { info, port, sharedKey } = wallets[0];
-   * const extensionWallet = ExtensionWallet.create(
-   *   info,
-   *   { chainId: Fr(31337), version: Fr(0) },
-   *   port,
-   *   sharedKey,
+   * const discoveredWallets = await ExtensionProvider.discoverWallets(chainInfo, { appId: 'my-defi-app' });
+   * const connection = await discoveredWallets[0].establishSecureChannel();
+   * const wallet = ExtensionWallet.create(
+   *   connection.info.id,
+   *   connection.port,
+   *   connection.sharedKey,
+   *   chainInfo,
    *   'my-defi-app'
    * );
    *
-   * // Register disconnect handler
-   * extensionWallet.onDisconnect(() => console.log('Disconnected!'));
-   *
-   * // Get the Wallet interface for dApp usage
-   * const wallet = extensionWallet.getWallet();
    * const accounts = await wallet.getAccounts();
    * ```
    */
   static create(
-    walletInfo: WalletInfo,
-    chainInfo: ChainInfo,
+    extensionId: string,
     port: MessagePort,
     sharedKey: CryptoKey,
+    chainInfo: ChainInfo,
     appId: string,
+    logger?: WalletSdkLogger,
+    heartbeatOptions?: HeartbeatOptions,
   ): ExtensionWallet {
-    const wallet = new ExtensionWallet(chainInfo, appId, walletInfo.id, port, sharedKey);
+    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;
+      // Any inbound traffic counts as proof of liveness.
+      wallet.lastInboundAt = Date.now();
 
-    // Set up message handler - all messages are now encrypted
-    wallet.port.onmessage = (event: MessageEvent<EncryptedPayload>) => {
-      void wallet.handleEncryptedResponse(event.data);
+      if (data && typeof data === 'object' && 'type' in data && data.type === WalletMessageType.DISCONNECT) {
+        wallet.handleDisconnect();
+        return;
+      }
+      // Otherwise treat as encrypted payload
+      void wallet.handleEncryptedResponse(data as EncryptedPayload);
     };
 
     wallet.port.start();
 
-    return wallet;
-  }
-
-  /**
-   * Returns a Wallet interface that proxies all method calls through the secure channel.
-   *
-   * The returned Wallet can be used directly by dApps - all method calls are automatically
-   * encrypted and sent to the wallet extension.
-   *
-   * @returns A Wallet implementation that encrypts all communication
-   *
-   * @example
-   * ```typescript
-   * const extensionWallet = ExtensionWallet.create(info, chainInfo, port, sharedKey, 'my-app');
-   * const wallet = extensionWallet.getWallet();
-   * const accounts = await wallet.getAccounts();
-   * ```
-   */
-  getWallet(): Wallet {
-    if (this.walletProxy) {
-      return this.walletProxy;
-    }
-
-    // Create a Proxy that intercepts wallet method calls and forwards them to the extension
-    this.walletProxy = new Proxy(this, {
-      get: (target, prop) => {
-        if (schemaHasMethod(WalletSchema, prop.toString())) {
+    return new Proxy(wallet, {
+      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>,
               args,
             });
-            return WalletSchema[prop.toString() as keyof typeof WalletSchema].returnType().parseAsync(result);
+            return getSchemaReturnType(WalletSchema[prop.toString() as keyof typeof WalletSchema]).parseAsync(result);
           };
         } else {
           return target[prop as keyof ExtensionWallet];
         }
       },
-    }) as unknown as Wallet;
+    });
+  }
 
-    return this.walletProxy;
+  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;
   }
 
   /**
@@ -188,18 +193,6 @@ export class ExtensionWallet {
 
       const { messageId, result, error, walletId: responseWalletId } = response;
 
-      // Check for disconnect notification from the wallet backend
-      // This is sent as an encrypted error response with a special type
-      if (
-        error &&
-        typeof error === 'object' &&
-        'type' in error &&
-        (error.type as WalletMessageType) === WalletMessageType.SESSION_DISCONNECTED
-      ) {
-        this.handleDisconnect();
-        return;
-      }
-
       if (!messageId || !responseWalletId) {
         return;
       }
@@ -220,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 });
+    }
   }
 
   /**
@@ -254,15 +249,64 @@ export class ExtensionWallet {
       walletId: this.extensionId,
     };
 
-    // Encrypt the message and send over the private MessageChannel
-    const encrypted = await encrypt(this.sharedKey, message);
+    const encrypted = await encrypt(this.sharedKey, jsonStringify(message));
     this.port.postMessage(encrypted);
 
     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.
@@ -274,27 +318,27 @@ export class ExtensionWallet {
     }
     this.disconnected = true;
 
-    // Close the port to prevent any further messages
+    if (this.heartbeatTimer !== null) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+
     if (this.port) {
       this.port.onmessage = null;
       this.port.close();
     }
 
-    // Reject all pending requests
-    // Note: These rejections should be caught by the callers, but we log them
-    // here to help with debugging if they become unhandled
     const error = new Error('Wallet disconnected');
     for (const { reject } of this.inFlight.values()) {
       reject(error);
     }
     this.inFlight.clear();
 
-    // Notify registered callbacks
     for (const callback of this.disconnectCallbacks) {
       try {
         callback();
       } catch {
-        // Ignore errors in callbacks
+        // Ignore errors on disconnect callbacks
       }
     }
   }
@@ -343,37 +387,24 @@ export class ExtensionWallet {
    *
    * @example
    * ```typescript
-   * const wallet = await provider.connect('my-app');
+   * const extensionWallet = ExtensionWallet.create(extensionId, port, sharedKey, chainInfo, 'my-app');
    * // ... use wallet ...
-   * await wallet.disconnect(); // Clean disconnect when done
+   * await extensionWallet.disconnect(); // Clean disconnect when done
    * ```
    */
+  // eslint-disable-next-line require-await -- async for interface compatibility
   async disconnect(): Promise<void> {
     if (this.disconnected) {
       return;
     }
 
-    // Send disconnect message to extension before closing
-    if (this.port && this.sharedKey) {
-      try {
-        const message = {
-          type: WalletMessageType.DISCONNECT,
-          messageId: globalThis.crypto.randomUUID(),
-          chainInfo: this.chainInfo,
-          appId: this.appId,
-          walletId: this.extensionId,
-          args: [],
-        };
-        const encrypted = await encrypt(this.sharedKey, message);
-        this.port.postMessage(encrypted);
-      } catch {
-        // Ignore errors sending disconnect message
-      }
+    if (this.port) {
+      // Send unencrypted disconnect - control messages don't need encryption
+      this.port.postMessage({
+        type: WalletMessageType.DISCONNECT,
+      });
     }
 
     this.handleDisconnect();
-    if (this.port) {
-      this.port.close();
-    }
   }
 }

package/src/extension/provider/index.ts

@@ -0,0 +1,7 @@
+export { ExtensionWallet } from './extension_wallet.js';
+export {
+  ExtensionProvider,
+  type DiscoveredWallet,
+  type ConnectedWallet,
+  type DiscoveryOptions,
+} from './extension_provider.js';