@aztec/wallet-sdk 0.0.1-commit.21caa21 → 0.0.1-commit.2ed92850

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/extension/provider/extension_wallet.ts

@@ -0,0 +1,321 @@
+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 type { FunctionsOf } from '@aztec/foundation/types';
+
+import { type EncryptedPayload, decrypt, encrypt } from '../../crypto.js';
+import { type WalletMessage, WalletMessageType, type WalletResponse } from '../../types.js';
+
+/**
+ * Internal type representing a wallet method call before encryption.
+ * @internal
+ */
+type WalletMethodCall = {
+  /** The wallet method name to invoke */
+  type: keyof FunctionsOf<Wallet>;
+  /** Arguments to pass to the wallet method */
+  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.
+ *
+ * This class uses a secure channel established after discovery:
+ *
+ * 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 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. This is what secures the channel.
+ *
+ * @example
+ * ```typescript
+ * // 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(connection.info.verificationHash!));
+ *
+ * // 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();
+ * ```
+ */
+export class ExtensionWallet {
+  /** Map of pending requests awaiting responses, keyed by message ID */
+  private inFlight = new Map<string, PromiseWithResolvers<unknown>>();
+  private disconnected = false;
+  private disconnectCallbacks: DisconnectCallback[] = [];
+
+  /**
+   * Private constructor - use {@link ExtensionWallet.create} to instantiate.
+   * @param chainInfo - The chain information (chainId and version)
+   * @param appId - Application identifier for the requesting dApp
+   * @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
+   */
+  private constructor(
+    private chainInfo: ChainInfo,
+    private appId: string,
+    private extensionId: string,
+    private port: MessagePort,
+    private sharedKey: CryptoKey,
+  ) {}
+
+  /**
+   * Creates a Wallet that communicates with a browser extension
+   * over a secure encrypted MessageChannel.
+   *
+   * @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 A Wallet interface where all method calls are encrypted
+   *
+   * @example
+   * ```typescript
+   * 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'
+   * );
+   *
+   * const accounts = await wallet.getAccounts();
+   * ```
+   */
+  static create(
+    extensionId: string,
+    port: MessagePort,
+    sharedKey: CryptoKey,
+    chainInfo: ChainInfo,
+    appId: string,
+  ): 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 new Proxy(wallet, {
+      get: (target, prop) => {
+        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);
+          };
+        } else {
+          return target[prop as keyof ExtensionWallet];
+        }
+      },
+    }) as unknown as Wallet;
+  }
+
+  /**
+   * Handles an encrypted response received from the wallet extension.
+   *
+   * Decrypts the response using the shared AES key and resolves or rejects
+   * the corresponding pending promise based on the response content.
+   *
+   * @param encrypted - The encrypted response from the wallet
+   */
+  private async handleEncryptedResponse(encrypted: EncryptedPayload): Promise<void> {
+    if (!this.sharedKey) {
+      return;
+    }
+
+    try {
+      const response = await decrypt<WalletResponse>(this.sharedKey, encrypted);
+
+      const { messageId, result, error, walletId: responseWalletId } = response;
+
+      if (!messageId || !responseWalletId) {
+        return;
+      }
+
+      if (this.extensionId !== responseWalletId) {
+        return;
+      }
+
+      if (!this.inFlight.has(messageId)) {
+        return;
+      }
+
+      const { resolve, reject } = this.inFlight.get(messageId)!;
+
+      if (error) {
+        reject(new Error(jsonStringify(error)));
+      } else {
+        resolve(result);
+      }
+      this.inFlight.delete(messageId);
+      // eslint-disable-next-line no-empty
+    } catch {}
+  }
+
+  /**
+   * Sends an encrypted wallet method call over the secure MessageChannel.
+   *
+   * The message is encrypted using AES-256-GCM with the shared key derived
+   * during discovery. A unique message ID is generated to correlate
+   * the response.
+   *
+   * @param call - The wallet method call containing method name and arguments
+   * @returns A Promise that resolves with the decrypted result from the wallet
+   *
+   * @throws Error if the secure channel has not been established or wallet is disconnected
+   */
+  private async postMessage(call: WalletMethodCall): Promise<unknown> {
+    if (this.disconnected) {
+      throw new Error('Wallet has been disconnected');
+    }
+    if (!this.port || !this.sharedKey) {
+      throw new Error('Secure channel not established');
+    }
+
+    const messageId = globalThis.crypto.randomUUID();
+    const message: WalletMessage = {
+      type: call.type,
+      args: call.args,
+      messageId,
+      chainInfo: this.chainInfo,
+      appId: this.appId,
+      walletId: this.extensionId,
+    };
+
+    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 });
+    return promise;
+  }
+
+  /**
+   * Handles wallet disconnection.
+   * Rejects all pending requests and notifies registered callbacks.
+   * @internal
+   */
+  private handleDisconnect(): void {
+    if (this.disconnected) {
+      return;
+    }
+    this.disconnected = true;
+
+    if (this.port) {
+      this.port.onmessage = null;
+      this.port.close();
+    }
+
+    const error = new Error('Wallet disconnected');
+    for (const { reject } of this.inFlight.values()) {
+      reject(error);
+    }
+    this.inFlight.clear();
+
+    for (const callback of this.disconnectCallbacks) {
+      try {
+        callback();
+      } catch {
+        // Ignore errors on disconnect callbacks
+      }
+    }
+  }
+
+  /**
+   * Registers a callback to be invoked when the wallet disconnects.
+   *
+   * @param callback - Function to call when wallet disconnects
+   * @returns A function to unregister the callback
+   *
+   * @example
+   * ```typescript
+   * const wallet = await ExtensionWallet.create(...);
+   * const unsubscribe = wallet.onDisconnect(() => {
+   *   console.log('Wallet disconnected! Please reconnect.');
+   *   // Clean up UI, prompt user to reconnect, etc.
+   * });
+   * // Later: unsubscribe(); to stop receiving notifications
+   * ```
+   */
+  onDisconnect(callback: DisconnectCallback): () => void {
+    this.disconnectCallbacks.push(callback);
+    return () => {
+      const index = this.disconnectCallbacks.indexOf(callback);
+      if (index !== -1) {
+        this.disconnectCallbacks.splice(index, 1);
+      }
+    };
+  }
+
+  /**
+   * Returns whether the wallet has been disconnected.
+   *
+   * @returns true if the wallet is no longer connected
+   */
+  isDisconnected(): boolean {
+    return this.disconnected;
+  }
+
+  /**
+   * Disconnects from the wallet and cleans up resources.
+   *
+   * This method notifies the wallet extension that the session is ending,
+   * allowing it to clean up its state. After calling this method, the wallet
+   * instance can no longer be used and any pending requests will be rejected.
+   *
+   * @example
+   * ```typescript
+   * const extensionWallet = ExtensionWallet.create(extensionId, port, sharedKey, chainInfo, 'my-app');
+   * // ... use wallet ...
+   * 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;
+    }
+
+    if (this.port) {
+      // Send unencrypted disconnect - control messages don't need encryption
+      this.port.postMessage({
+        type: WalletMessageType.DISCONNECT,
+      });
+    }
+
+    this.handleDisconnect();
+  }
+}

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,20 +5,8 @@ export type {
   WebWalletConfig,
   WalletProviderType,
   WalletProvider,
+  PendingConnection,
+  ProviderDisconnectionCallback,
   DiscoverWalletsOptions,
+  DiscoverySession,
 } from './types.js';
-
-// Re-export types from providers for convenience
-export type {
-  WalletInfo,
-  WalletMessage,
-  WalletResponse,
-  DiscoveryRequest,
-  DiscoveryResponse,
-} from '../providers/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';