@aztec/wallet-sdk 0.0.1-commit.f295ac2 → 0.0.1-commit.fc805bf

package/src/extension/handlers/internal_message_types.ts

@@ -0,0 +1,69 @@
+/**
+ * Internal message types for content script ↔ background communication.
+ * These are NOT part of the public wallet protocol - they are implementation
+ * details for coordinating between extension components.
+ */
+export const InternalMessageType = {
+  // Content script → Background
+  DISCOVERY_REQUEST: 'discovery-request',
+  KEY_EXCHANGE_REQUEST: 'key-exchange-request',
+  SECURE_MESSAGE: 'secure-message',
+  DISCONNECT_REQUEST: 'disconnect-request',
+  // Background → Content script
+  DISCOVERY_APPROVED: 'discovery-approved',
+  KEY_EXCHANGE_RESPONSE: 'key-exchange-response',
+  SECURE_RESPONSE: 'secure-response',
+  SESSION_DISCONNECTED: 'session-disconnected',
+} as const;
+
+/**
+ * Message origins for internal extension communication.
+ */
+export const MessageOrigin = {
+  BACKGROUND: 'background',
+  CONTENT_SCRIPT: 'content-script',
+} as const;
+
+/** Union type of message origins. */
+export type MessageOriginType = (typeof MessageOrigin)[keyof typeof MessageOrigin];
+
+/**
+ * Message sent from content script to background.
+ */
+export interface ContentScriptMessage {
+  /** Message source identifier. */
+  origin: typeof MessageOrigin.CONTENT_SCRIPT;
+  /** Message type. */
+  type: string;
+  /** Optional session identifier. */
+  sessionId?: string;
+  /** Optional message payload. */
+  content?: unknown;
+}
+
+/**
+ * Message sent from background to content script.
+ */
+export interface BackgroundMessage {
+  /** Message source identifier. */
+  origin: typeof MessageOrigin.BACKGROUND;
+  /** Message type. */
+  type: string;
+  /** Session identifier. */
+  sessionId: string;
+  /** Optional message payload. */
+  content?: unknown;
+}
+
+/**
+ * Sender information for messages from browser runtime.
+ */
+export interface MessageSender {
+  /** Tab information if available. */
+  tab?: {
+    /** Tab identifier. */
+    id?: number;
+    /** Tab URL. */
+    url?: string;
+  };
+}

package/src/extension/provider/extension_provider.ts

@@ -0,0 +1,233 @@
+import type { ChainInfo } from '@aztec/aztec.js/account';
+import { jsonStringify } from '@aztec/foundation/json-rpc';
+import { promiseWithResolvers } from '@aztec/foundation/promise';
+
+import { deriveSessionKeys, exportPublicKey, generateKeyPair, importPublicKey } from '../../crypto.js';
+import {
+  type ConnectedWalletInfo,
+  type DiscoveryRequest,
+  type DiscoveryResponse,
+  type KeyExchangeRequest,
+  type KeyExchangeResponse,
+  type WalletInfo,
+  WalletMessageType,
+} from '../../types.js';
+
+/** Default discovery timeout - long to give users time to approve */
+const DEFAULT_DISCOVERY_TIMEOUT_MS = 60000; // 60 seconds
+
+/** Key exchange timeout - short, wallet should respond quickly after discovery approval */
+const KEY_EXCHANGE_TIMEOUT_MS = 2000; // 2 seconds
+
+/**
+ * A discovered wallet before key exchange.
+ * Has basic info and MessagePort, but no shared key yet.
+ *
+ * Call {@link establishSecureChannel} to perform key exchange and get a connected wallet.
+ */
+export class DiscoveredWallet {
+  constructor(
+    /** Basic wallet information (id, name, icon, version) */
+    public readonly info: WalletInfo,
+    /** The MessagePort for private communication with the wallet */
+    public readonly port: MessagePort,
+    /** Request ID for correlation */
+    public readonly requestId: string,
+  ) {}
+
+  /**
+   * Establishes a secure connection with this wallet.
+   *
+   * This method:
+   * 1. Generates an ECDH key pair
+   * 2. Sends public key to wallet over the MessagePort
+   * 3. Receives wallet's public key
+   * 4. Derives shared secret and computes verification hash locally
+   *
+   * **IMPORTANT**: Has a 2 second timeout for MITM defense.
+   * Both parties must exchange keys relatively quickly.
+   *
+   * The verification hash is computed independently by both parties
+   * and should be displayed to the user for visual comparison.
+   *
+   * @returns Connected wallet with shared key and verification hash
+   * @throws Error if key exchange fails or times out
+   */
+  async establishSecureChannel(): Promise<ConnectedWallet> {
+    const keyPair = await generateKeyPair();
+    const exportedPublicKey = await exportPublicKey(keyPair.publicKey);
+
+    const { promise, resolve, reject } = promiseWithResolvers<ConnectedWallet>();
+
+    const timeoutId = setTimeout(() => {
+      reject(new Error('Key exchange timeout'));
+    }, KEY_EXCHANGE_TIMEOUT_MS);
+
+    this.port.onmessage = async (event: MessageEvent) => {
+      const data = event.data as KeyExchangeResponse;
+
+      if (data.type === WalletMessageType.KEY_EXCHANGE_RESPONSE && data.requestId === this.requestId) {
+        clearTimeout(timeoutId);
+
+        try {
+          const walletPublicKey = await importPublicKey(data.publicKey);
+          const session = await deriveSessionKeys(keyPair, walletPublicKey, true);
+
+          const connectedInfo: ConnectedWalletInfo = {
+            ...this.info,
+            publicKey: data.publicKey,
+            verificationHash: session.verificationHash,
+          };
+
+          resolve({
+            info: connectedInfo,
+            port: this.port,
+            sharedKey: session.encryptionKey,
+          });
+        } catch (err) {
+          reject(new Error(`Key exchange failed: ${err}`));
+        }
+      }
+    };
+
+    this.port.start();
+
+    const keyExchangeRequest: KeyExchangeRequest = {
+      type: WalletMessageType.KEY_EXCHANGE_REQUEST,
+      requestId: this.requestId,
+      publicKey: exportedPublicKey,
+    };
+    this.port.postMessage(keyExchangeRequest);
+
+    return promise;
+  }
+}
+
+/**
+ * A fully connected wallet with secure channel established.
+ * Available after key exchange completes.
+ */
+export interface ConnectedWallet {
+  /** Full wallet info including public key and verification hash */
+  info: ConnectedWalletInfo;
+  /** The MessagePort for encrypted communication */
+  port: MessagePort;
+  /** The derived AES-256-GCM shared key for encryption */
+  sharedKey: CryptoKey;
+}
+
+/**
+ * Options for wallet discovery.
+ */
+export interface DiscoveryOptions {
+  /** Application ID making the request */
+  appId: string;
+  /** How long to wait for user approval (ms). Default: 60000 (60s) */
+  timeout?: number;
+  /**
+   * Callback invoked when a wallet is discovered.
+   * Wallets are streamed as users approve them.
+   */
+  onWalletDiscovered?: (wallet: DiscoveredWallet) => void;
+  /**
+   * AbortSignal for cancelling discovery early.
+   * When aborted, cleanup happens immediately instead of waiting for timeout.
+   */
+  signal?: AbortSignal;
+}
+
+/**
+ * Provider for discovering Aztec wallet extensions.
+ *
+ * NOTE: Most users should use WalletManager instead of this class directly.
+ * WalletManager provides a higher-level API with streaming support.
+ *
+ * The connection flow is split into two phases for security:
+ *
+ * 1. **Discovery Phase** ({@link discoverWallets}):
+ *    - Broadcasts a discovery request (NO public keys)
+ *    - Wallet shows pending request to user
+ *    - User must approve before wallet reveals itself
+ *    - Wallets are streamed via callback as they're approved
+ *
+ * 2. **Secure Channel Phase** ({@link DiscoveredWallet.establishSecureChannel}):
+ *    - Performs ECDH key exchange over private MessageChannel
+ *    - Both parties compute verification hash locally
+ *    - Has a 2s timeout for MITM defense
+ *    - Returns connected wallet with shared key and verification hash
+ */
+export class ExtensionProvider {
+  /**
+   * Discovers wallet extensions that user has approved.
+   *
+   * Wallets are streamed via the `onWalletDiscovered` callback as users approve them.
+   * The promise resolves when the timeout expires or signal is aborted.
+   *
+   * @param chainInfo - Chain information to check if extensions support this network
+   * @param options - Discovery options including appId, appName, timeout, and callback
+   * @returns Promise that resolves when discovery completes
+   */
+  static discoverWallets(chainInfo: ChainInfo, options: DiscoveryOptions): Promise<void> {
+    if (options.signal?.aborted) {
+      return Promise.resolve();
+    }
+    const timeout = options.timeout ?? DEFAULT_DISCOVERY_TIMEOUT_MS;
+
+    return new Promise(resolve => {
+      const requestId = globalThis.crypto.randomUUID();
+      let timeoutId: ReturnType<typeof setTimeout> | null = null;
+      let completed = false;
+
+      const finish = () => {
+        if (completed) {
+          return;
+        }
+        completed = true;
+
+        if (timeoutId !== null) {
+          clearTimeout(timeoutId);
+        }
+        window.removeEventListener('message', onMessage);
+        options.signal?.removeEventListener('abort', onAbort);
+        resolve();
+      };
+
+      const onAbort = () => finish();
+
+      const onMessage = (event: MessageEvent) => {
+        if (completed || event.source !== window) {
+          return;
+        }
+
+        let data: DiscoveryResponse;
+        try {
+          data = JSON.parse(event.data);
+        } catch {
+          return;
+        }
+
+        if (data.type !== WalletMessageType.DISCOVERY_RESPONSE || data.requestId !== requestId) {
+          return;
+        }
+
+        const port = event.ports?.[0];
+        if (port) {
+          options.onWalletDiscovered?.(new DiscoveredWallet(data.walletInfo, port, requestId));
+        }
+      };
+
+      options.signal?.addEventListener('abort', onAbort, { once: true });
+      window.addEventListener('message', onMessage);
+
+      const discoveryMessage: DiscoveryRequest = {
+        type: WalletMessageType.DISCOVERY,
+        requestId,
+        appId: options.appId,
+        chainInfo,
+      };
+      window.postMessage(jsonStringify(discoveryMessage), '*');
+
+      timeoutId = setTimeout(finish, timeout);
+    });
+  }
+}

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

@@ -6,7 +6,7 @@ import { 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 { type WalletMessage, WalletMessageType, type WalletResponse } from '../../types.js';
 
 /**
  * Internal type representing a wallet method call before encryption.
@@ -26,30 +26,30 @@ 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();
@@ -77,81 +77,56 @@ export class ExtensionWallet {
     private sharedKey: CryptoKey,
   ) {}
 
-  /** Cached Wallet proxy instance */
-  private walletProxy: Wallet | null = null;
-
   /**
-   * 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.
+   * @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,
-  ): ExtensionWallet {
-    const wallet = new ExtensionWallet(chainInfo, appId, walletInfo.id, port, sharedKey);
-
-    // Set up message handler - all messages are now encrypted
-    wallet.port.onmessage = (event: MessageEvent<EncryptedPayload>) => {
-      void wallet.handleEncryptedResponse(event.data);
+  ): Wallet {
+    const wallet = new ExtensionWallet(chainInfo, appId, extensionId, port, sharedKey);
+
+    // 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
+      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, {
+    return new Proxy(wallet, {
       get: (target, prop) => {
         if (schemaHasMethod(WalletSchema, prop.toString())) {
           return async (...args: unknown[]) => {
@@ -166,8 +141,6 @@ export class ExtensionWallet {
         }
       },
     }) as unknown as Wallet;
-
-    return this.walletProxy;
   }
 
   /**
@@ -188,18 +161,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;
       }
@@ -254,8 +215,7 @@ 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>();
@@ -274,27 +234,22 @@ export class ExtensionWallet {
     }
     this.disconnected = true;
 
-    // Close the port to prevent any further messages
     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 +298,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, type DisconnectCallback } from './extension_wallet.js';
+export {
+  ExtensionProvider,
+  type DiscoveredWallet,
+  type ConnectedWallet,
+  type DiscoveryOptions,
+} from './extension_provider.js';

package/src/manager/index.ts

@@ -5,16 +5,8 @@ export type {
   WebWalletConfig,
   WalletProviderType,
   WalletProvider,
+  PendingConnection,
   ProviderDisconnectionCallback,
   DiscoverWalletsOptions,
+  DiscoverySession,
 } from './types.js';
-
-// Re-export types and enums from providers for convenience
-export { WalletMessageType } from '../types.js';
-export type { WalletInfo, WalletMessage, WalletResponse, DiscoveryRequest, DiscoveryResponse } from '../types.js';
-
-// Re-export commonly needed utilities for wallet integration
-export { ChainInfoSchema } from '@aztec/aztec.js/account';
-export type { ChainInfo } from '@aztec/aztec.js/account';
-export { WalletSchema } from '@aztec/aztec.js/wallet';
-export { jsonStringify } from '@aztec/foundation/json-rpc';