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

package/src/extension/handlers/background_connection_handler.ts

@@ -0,0 +1,423 @@
+import type { ChainInfo } from '@aztec/aztec.js/account';
+
+import {
+  type EncryptedPayload,
+  decrypt,
+  deriveSessionKeys,
+  encrypt,
+  exportPublicKey,
+  generateKeyPair,
+  importPublicKey,
+} from '../../crypto.js';
+import {
+  type DiscoveryRequest,
+  type KeyExchangeRequest,
+  type KeyExchangeResponse,
+  type WalletInfo,
+  type WalletMessage,
+  WalletMessageType,
+  type WalletResponse,
+} from '../../types.js';
+import {
+  type BackgroundMessage,
+  type ContentScriptMessage,
+  InternalMessageType,
+  MessageOrigin,
+  type MessageSender,
+} from './internal_message_types.js';
+
+/**
+ * Status of a pending discovery request.
+ */
+export type DiscoveryStatus = 'pending' | 'approved' | 'rejected';
+
+/**
+ * A pending discovery request from a dApp.
+ */
+export interface PendingDiscovery {
+  /** Unique request identifier. */
+  requestId: string;
+  /** Application identifier. */
+  appId: string;
+  /** Optional application name. */
+  appName?: string;
+  /** Origin URL of the requesting page. */
+  origin: string;
+  /** Network information. */
+  chainInfo: ChainInfo;
+  /** Browser tab ID. */
+  tabId: number;
+  /** Request timestamp. */
+  timestamp: number;
+  /** Current status. */
+  status: DiscoveryStatus;
+}
+
+/**
+ * An active session with a connected dApp.
+ * Created after key exchange completes.
+ */
+export interface ActiveSession {
+  /** Unique session identifier. */
+  sessionId: string;
+  /** Derived AES-GCM encryption key. */
+  sharedKey: CryptoKey;
+  /** Hex-encoded verification hash for visual comparison. */
+  verificationHash: string;
+  /** Browser tab ID. */
+  tabId: number;
+  /** Origin URL of the connected app. */
+  origin: string;
+  /** Application identifier. */
+  appId: string;
+  /** Connection timestamp. */
+  connectedAt: number;
+  /** Network information. */
+  chainInfo: ChainInfo;
+}
+
+/**
+ * Transport interface for background script communication.
+ */
+export interface BackgroundTransport {
+  /**
+   * Send a message to a specific tab.
+   * Typically `(tabId, message) => browser.tabs.sendMessage(tabId, message)`.
+   */
+  sendToTab: (tabId: number, message: BackgroundMessage) => void;
+
+  /**
+   * Register a listener for messages from content scripts.
+   * Typically `browser.runtime.onMessage.addListener`.
+   */
+  addContentListener: (handler: (message: unknown, sender: MessageSender) => void) => void;
+}
+
+/**
+ * Event callbacks for the background connection handler.
+ * All callbacks are optional.
+ */
+export interface BackgroundConnectionCallbacks {
+  /**
+   * Called when a new discovery request is received and stored as pending.
+   */
+  onPendingDiscovery?: (discovery: PendingDiscovery) => 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 decrypted wallet message is received.
+   */
+  onWalletMessage?: (session: ActiveSession, message: WalletMessage) => void;
+}
+
+/**
+ * Configuration for the background connection handler.
+ */
+export interface BackgroundConnectionConfig {
+  /** Unique wallet identifier. */
+  walletId: string;
+  /** Display name for the wallet. */
+  walletName: string;
+  /** Wallet version string. */
+  walletVersion: string;
+  /** Optional wallet icon URL. */
+  walletIcon?: string;
+}
+
+/**
+ * Handles wallet session flow in the extension background script.
+ *
+ * This class manages:
+ * - Pending discovery requests (before user approval)
+ * - Active sessions (after key exchange)
+ * - Per-session ECDH key exchange
+ * - Message encryption/decryption
+ *
+ * @example
+ * ```typescript
+ * const handler = new BackgroundConnectionHandler(
+ *   {
+ *     walletId: 'my-wallet',
+ *     walletName: 'My Wallet',
+ *     walletVersion: '1.0.0',
+ *   },
+ *   {
+ *     sendToTab: (tabId, message) => browser.tabs.sendMessage(tabId, message),
+ *     addContentListener: (handler) => browser.runtime.onMessage.addListener(handler),
+ *   },
+ *   {
+ *     onPendingDiscovery: (discovery) => updateBadge(),
+ *     onSessionEstablished: (session) => console.log('Connected:', session.sessionId),
+ *     onWalletMessage: (session, message) => nativePort.postMessage(message),
+ *   }
+ * );
+ *
+ * handler.initialize();
+ * ```
+ */
+export class BackgroundConnectionHandler {
+  private pendingDiscoveries = new Map<string, PendingDiscovery>();
+  private activeSessions = new Map<string, ActiveSession>();
+
+  constructor(
+    private config: BackgroundConnectionConfig,
+    private transport: BackgroundTransport,
+    private callbacks: BackgroundConnectionCallbacks = {},
+  ) {}
+
+  initialize(): void {
+    this.transport.addContentListener(this.handleMessage);
+  }
+
+  private handleMessage = (message: unknown, sender: MessageSender): void => {
+    const msg = message as ContentScriptMessage;
+    if (msg.origin !== MessageOrigin.CONTENT_SCRIPT) {
+      return;
+    }
+
+    const tabId = sender.tab?.id;
+    const tabOrigin = sender.tab?.url ? new URL(sender.tab.url).origin : 'unknown';
+    if (!tabId) {
+      return;
+    }
+
+    const { type, sessionId, content } = msg;
+
+    switch (type) {
+      case InternalMessageType.DISCOVERY_REQUEST:
+        this.handleDiscoveryRequest(content as DiscoveryRequest, tabId, tabOrigin);
+        break;
+      case InternalMessageType.KEY_EXCHANGE_REQUEST:
+        if (sessionId) {
+          this.handleKeyExchangeRequest(sessionId, content as KeyExchangeRequest).catch(() => {
+            // Key exchange failed - session won't be established
+          });
+        }
+        break;
+      case InternalMessageType.DISCONNECT_REQUEST:
+        if (sessionId) {
+          this.terminateSession(sessionId);
+        }
+        break;
+      case InternalMessageType.SECURE_MESSAGE:
+        if (sessionId) {
+          void this.handleEncryptedMessage(sessionId, content as EncryptedPayload);
+        }
+        break;
+    }
+  };
+
+  getWalletInfo(): WalletInfo {
+    return {
+      id: this.config.walletId,
+      name: this.config.walletName,
+      version: this.config.walletVersion,
+      icon: this.config.walletIcon,
+    };
+  }
+
+  handleDiscoveryRequest(request: DiscoveryRequest, tabId: number, origin: string): void {
+    const discovery: PendingDiscovery = {
+      requestId: request.requestId,
+      appId: request.appId,
+      origin,
+      chainInfo: request.chainInfo,
+      tabId,
+      timestamp: Date.now(),
+      status: 'pending',
+    };
+
+    this.pendingDiscoveries.set(request.requestId, discovery);
+    this.callbacks.onPendingDiscovery?.(discovery);
+  }
+
+  approveDiscovery(requestId: string): boolean {
+    const discovery = this.pendingDiscoveries.get(requestId);
+    if (!discovery || discovery.status !== 'pending') {
+      return false;
+    }
+
+    // The discovery requestId becomes our sessionId
+    // This is what will be used internally to correlate
+    // content<->background messages
+    const sessionId = requestId;
+
+    discovery.status = 'approved';
+    this.transport.sendToTab(discovery.tabId, {
+      origin: MessageOrigin.BACKGROUND,
+      type: InternalMessageType.DISCOVERY_APPROVED,
+      sessionId,
+      content: this.getWalletInfo(),
+    });
+
+    return true;
+  }
+
+  rejectDiscovery(requestId: string): boolean {
+    const discovery = this.pendingDiscoveries.get(requestId);
+    if (!discovery || discovery.status !== 'pending') {
+      return false;
+    }
+
+    discovery.status = 'rejected';
+    this.pendingDiscoveries.delete(requestId);
+
+    return true;
+  }
+
+  async handleKeyExchangeRequest(sessionId: string, request: KeyExchangeRequest): Promise<void> {
+    const discovery = this.pendingDiscoveries.get(sessionId);
+    if (!discovery || discovery.status !== 'approved') {
+      return;
+    }
+
+    try {
+      const keyPair = await generateKeyPair();
+      const publicKey = await exportPublicKey(keyPair.publicKey);
+
+      const appPublicKey = await importPublicKey(request.publicKey);
+      const sessionKeys = await deriveSessionKeys(keyPair, appPublicKey, false);
+
+      const session: ActiveSession = {
+        sessionId,
+        sharedKey: sessionKeys.encryptionKey,
+        verificationHash: sessionKeys.verificationHash,
+        tabId: discovery.tabId,
+        origin: discovery.origin,
+        appId: discovery.appId,
+        connectedAt: Date.now(),
+        chainInfo: discovery.chainInfo,
+      };
+
+      this.activeSessions.set(sessionId, session);
+      this.pendingDiscoveries.delete(sessionId);
+
+      const response: KeyExchangeResponse = {
+        type: WalletMessageType.KEY_EXCHANGE_RESPONSE,
+        requestId: sessionId, // Public protocol uses requestId
+        publicKey,
+      };
+
+      this.transport.sendToTab(discovery.tabId, {
+        origin: MessageOrigin.BACKGROUND,
+        type: InternalMessageType.KEY_EXCHANGE_RESPONSE,
+        sessionId,
+        content: response,
+      });
+
+      this.callbacks.onSessionEstablished?.(session);
+    } catch {
+      // Key exchange failed silently - session won't be established
+    }
+  }
+
+  async handleEncryptedMessage(sessionId: string, encrypted: EncryptedPayload): Promise<void> {
+    const session = this.activeSessions.get(sessionId);
+    if (!session) {
+      return;
+    }
+
+    try {
+      const message = await decrypt<WalletMessage>(session.sharedKey, encrypted);
+      this.callbacks.onWalletMessage?.(session, message);
+    } catch {
+      // Decryption failed - ignore malformed message
+    }
+  }
+
+  async sendResponse(sessionId: string, response: WalletResponse): Promise<void> {
+    const session = this.activeSessions.get(sessionId);
+    if (!session) {
+      return;
+    }
+
+    try {
+      const encrypted = await encrypt(session.sharedKey, JSON.stringify(response));
+      this.transport.sendToTab(session.tabId, {
+        origin: MessageOrigin.BACKGROUND,
+        type: InternalMessageType.SECURE_RESPONSE,
+        sessionId,
+        content: encrypted,
+      });
+    } catch {
+      // Encryption failed - response won't be sent
+    }
+  }
+
+  terminateSession(sessionId: string): void {
+    const session = this.activeSessions.get(sessionId);
+    if (session) {
+      // Notify the content script (and ultimately the dApp) that the session is disconnected
+      this.transport.sendToTab(session.tabId, {
+        origin: MessageOrigin.BACKGROUND,
+        type: InternalMessageType.SESSION_DISCONNECTED,
+        sessionId,
+      });
+
+      this.activeSessions.delete(sessionId);
+      this.callbacks.onSessionTerminated?.(sessionId);
+
+      // Restore discovery to approved state so user can retry key exchange
+      const discovery: PendingDiscovery = {
+        requestId: sessionId,
+        appId: session.appId,
+        origin: session.origin,
+        chainInfo: session.chainInfo,
+        tabId: session.tabId,
+        timestamp: Date.now(),
+        status: 'approved',
+      };
+      this.pendingDiscoveries.set(sessionId, discovery);
+    }
+  }
+
+  terminateForTab(tabId: number): void {
+    for (const [sessionId, session] of this.activeSessions) {
+      if (session.tabId === tabId) {
+        this.terminateSession(sessionId);
+      }
+    }
+    for (const [requestId, discovery] of this.pendingDiscoveries) {
+      if (discovery.tabId === tabId) {
+        this.pendingDiscoveries.delete(requestId);
+      }
+    }
+  }
+
+  clearAll(): void {
+    for (const sessionId of this.activeSessions.keys()) {
+      this.callbacks.onSessionTerminated?.(sessionId);
+    }
+    this.activeSessions.clear();
+    this.pendingDiscoveries.clear();
+  }
+
+  getPendingDiscoveries(): PendingDiscovery[] {
+    return Array.from(this.pendingDiscoveries.values()).filter(d => d.status === 'pending');
+  }
+
+  getPendingDiscoveryCount(): number {
+    return this.getPendingDiscoveries().length;
+  }
+
+  getActiveSessions(): ActiveSession[] {
+    return Array.from(this.activeSessions.values());
+  }
+
+  getSession(sessionId: string): ActiveSession | undefined {
+    return this.activeSessions.get(sessionId);
+  }
+
+  getPendingDiscovery(requestId: string): PendingDiscovery | undefined {
+    return this.pendingDiscoveries.get(requestId);
+  }
+}

package/src/extension/handlers/content_script_connection_handler.ts

@@ -0,0 +1,246 @@
+import type { EncryptedPayload, ExportedPublicKey } from '../../crypto.js';
+import { type DiscoveryRequest, type DiscoveryResponse, type WalletInfo, WalletMessageType } from '../../types.js';
+import {
+  type BackgroundMessage,
+  type ContentScriptMessage,
+  InternalMessageType,
+  MessageOrigin,
+} from './internal_message_types.js';
+
+/**
+ * Key exchange request sent over MessageChannel.
+ */
+interface KeyExchangeRequest {
+  /** Message type identifier. */
+  type: typeof WalletMessageType.KEY_EXCHANGE_REQUEST;
+  /** Request identifier for correlation. */
+  requestId: string;
+  /** ECDH public key in JWK format. */
+  publicKey: ExportedPublicKey;
+}
+
+/**
+ * Key exchange response sent over MessageChannel.
+ */
+interface KeyExchangeResponse {
+  /** Message type identifier. */
+  type: typeof WalletMessageType.KEY_EXCHANGE_RESPONSE;
+  /** Request identifier for correlation. */
+  requestId: string;
+  /** ECDH public key in JWK format. */
+  publicKey: ExportedPublicKey;
+}
+
+/**
+ * Port connection stored by the content script.
+ */
+interface PortConnection {
+  /** MessagePort for communication with page. */
+  port: MessagePort;
+  /** Session identifier. */
+  sessionId: string;
+}
+
+/**
+ * Transport interface for content script communication.
+ */
+export interface ContentScriptTransport {
+  /**
+   * Send a message to the background script.
+   * Typically `browser.runtime.sendMessage`.
+   */
+  sendToBackground: (message: ContentScriptMessage) => void;
+
+  /**
+   * Register a listener for messages from the background script.
+   * Typically `browser.runtime.onMessage.addListener`.
+   */
+  addBackgroundListener: (handler: (message: BackgroundMessage) => void) => void;
+}
+
+/**
+ * Handles wallet connection flow in the extension content script.
+ *
+ * This class manages:
+ * - Listening for discovery requests from the page
+ * - Creating MessageChannels after discovery approval
+ * - Relaying key exchange messages between page and background
+ * - Relaying encrypted messages between page and background
+ *
+ * The content script acts as a pure relay - it never has access to
+ * private keys or shared secrets.
+ *
+ * @example
+ * ```typescript
+ * const handler = new ContentScriptConnectionHandler({
+ *   sendToBackground: (message) => browser.runtime.sendMessage(message),
+ *   addBackgroundListener: (handler) => browser.runtime.onMessage.addListener(handler),
+ * });
+ *
+ * handler.start();
+ * ```
+ */
+export class ContentScriptConnectionHandler {
+  private ports = new Map<string, PortConnection>();
+  private listening = false;
+  private pageMessageHandler: ((event: MessageEvent) => void) | null = null;
+
+  constructor(private transport: ContentScriptTransport) {}
+
+  start(): void {
+    if (this.listening) {
+      return;
+    }
+
+    this.transport.addBackgroundListener(this.handleBackgroundMessage);
+
+    this.pageMessageHandler = (event: MessageEvent) => {
+      if (event.source !== window) {
+        return;
+      }
+
+      let data: DiscoveryRequest;
+      try {
+        data = JSON.parse(event.data);
+      } catch {
+        return;
+      }
+
+      if (!data?.type) {
+        return;
+      }
+
+      if (data.type === WalletMessageType.DISCOVERY) {
+        this.handleDiscoveryRequest(data);
+      }
+    };
+
+    window.addEventListener('message', this.pageMessageHandler);
+    this.listening = true;
+  }
+
+  private handleBackgroundMessage = (message: BackgroundMessage): void => {
+    if (message.origin !== MessageOrigin.BACKGROUND) {
+      return;
+    }
+
+    const { type, sessionId, content } = message;
+
+    switch (type) {
+      case InternalMessageType.DISCOVERY_APPROVED:
+        this.handleDiscoveryApproved(sessionId, content as WalletInfo);
+        break;
+      case InternalMessageType.KEY_EXCHANGE_RESPONSE:
+        this.handleKeyExchangeResponse(sessionId, content as KeyExchangeResponse);
+        break;
+      case InternalMessageType.SECURE_RESPONSE:
+        this.handleSecureResponse(sessionId, content as EncryptedPayload);
+        break;
+      case InternalMessageType.SESSION_DISCONNECTED:
+        this.handleSessionDisconnected(sessionId);
+        break;
+    }
+  };
+
+  private handleDiscoveryRequest(request: DiscoveryRequest): void {
+    this.transport.sendToBackground({
+      origin: MessageOrigin.CONTENT_SCRIPT,
+      type: InternalMessageType.DISCOVERY_REQUEST,
+      content: request,
+    });
+  }
+
+  private handleDiscoveryApproved(sessionId: string, walletInfo: WalletInfo): void {
+    const channel = new MessageChannel();
+
+    this.ports.set(sessionId, {
+      port: channel.port1,
+      sessionId,
+    });
+
+    channel.port1.onmessage = (event: MessageEvent) => {
+      const data = event.data;
+
+      switch (data?.type) {
+        case WalletMessageType.KEY_EXCHANGE_REQUEST:
+          this.transport.sendToBackground({
+            origin: MessageOrigin.CONTENT_SCRIPT,
+            type: InternalMessageType.KEY_EXCHANGE_REQUEST,
+            sessionId,
+            content: data as KeyExchangeRequest,
+          });
+          break;
+        case WalletMessageType.DISCONNECT:
+          this.transport.sendToBackground({
+            origin: MessageOrigin.CONTENT_SCRIPT,
+            type: InternalMessageType.DISCONNECT_REQUEST,
+            sessionId,
+            content: data,
+          });
+          break;
+        default:
+          this.transport.sendToBackground({
+            origin: MessageOrigin.CONTENT_SCRIPT,
+            type: InternalMessageType.SECURE_MESSAGE,
+            sessionId,
+            content: data as EncryptedPayload,
+          });
+          break;
+      }
+    };
+
+    channel.port1.start();
+
+    const response: DiscoveryResponse = {
+      type: WalletMessageType.DISCOVERY_RESPONSE,
+      requestId: sessionId,
+      walletInfo,
+    };
+    window.postMessage(JSON.stringify(response), '*', [channel.port2]);
+  }
+
+  private handleKeyExchangeResponse(sessionId: string, response: KeyExchangeResponse): void {
+    const connection = this.ports.get(sessionId);
+    if (!connection) {
+      return;
+    }
+    connection.port.postMessage(response);
+  }
+
+  private handleSecureResponse(sessionId: string, content: EncryptedPayload): void {
+    const connection = this.ports.get(sessionId);
+    if (!connection) {
+      return;
+    }
+    connection.port.postMessage(content);
+  }
+
+  private handleSessionDisconnected(sessionId: string): void {
+    const connection = this.ports.get(sessionId);
+    if (!connection) {
+      return;
+    }
+    connection.port.postMessage({ type: WalletMessageType.DISCONNECT });
+    connection.port.close();
+    this.ports.delete(sessionId);
+  }
+
+  closeConnection(sessionId: string): void {
+    const connection = this.ports.get(sessionId);
+    if (connection) {
+      connection.port.close();
+      this.ports.delete(sessionId);
+    }
+  }
+
+  closeAllConnections(): void {
+    for (const [sessionId, connection] of this.ports) {
+      connection.port.close();
+      this.ports.delete(sessionId);
+    }
+  }
+
+  getConnectionCount(): number {
+    return this.ports.size;
+  }
+}

package/src/extension/handlers/index.ts

@@ -0,0 +1,25 @@
+/**
+ * Browser-safe exports for extension background and content scripts.
+ *
+ * This module contains ONLY handlers that work in service worker/content script
+ * environments without Node.js polyfills.
+ *
+ * @packageDocumentation
+ */
+export {
+  BackgroundConnectionHandler,
+  type PendingDiscovery,
+  type ActiveSession,
+  type DiscoveryStatus,
+  type BackgroundConnectionCallbacks,
+  type BackgroundConnectionConfig,
+  type BackgroundTransport,
+} from './background_connection_handler.js';
+export { ContentScriptConnectionHandler, type ContentScriptTransport } from './content_script_connection_handler.js';
+export {
+  type ContentScriptMessage,
+  type BackgroundMessage,
+  type MessageSender,
+  type MessageOriginType,
+  MessageOrigin,
+} from './internal_message_types.js';