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

package/src/manager/wallet_manager.ts

@@ -1,48 +1,14 @@
-import type { ChainInfo } from '@aztec/aztec.js/account';
-import { promiseWithResolvers } from '@aztec/foundation/promise';
-
-import { type DiscoveredWallet, ExtensionProvider, ExtensionWallet } from '../extension/provider/index.js';
-import { WalletMessageType } from '../types.js';
+import { ExtensionProvider, ExtensionWallet } from '../providers/extension/index.js';
 import type {
   DiscoverWalletsOptions,
-  DiscoverySession,
   ExtensionWalletConfig,
-  PendingConnection,
   ProviderDisconnectionCallback,
   WalletManagerConfig,
   WalletProvider,
 } from './types.js';
 
 /**
- * Manager for wallet discovery, configuration, and connection.
- *
- * This is the main entry point for dApps to discover and connect to wallets.
- *
- * @example Basic usage with async iterator
- * ```typescript
- * const discovery = WalletManager.configure({ extensions: { enabled: true } })
- *   .getAvailableWallets({ chainInfo, appId: 'my-app' });
- *
- * // Iterate over discovered wallets
- * for await (const provider of discovery.wallets) {
- *   console.log(`Found wallet: ${provider.name}`);
- * }
- *
- * // Or cancel early when done
- * discovery.cancel();
- * ```
- *
- * @example With callback for discovered wallets
- * ```typescript
- * const discovery = manager.getAvailableWallets({
- *   chainInfo,
- *   appId: 'my-app',
- *   onWalletDiscovered: (provider) => console.log(`Found: ${provider.name}`),
- * });
- *
- * // Wait for discovery to complete or cancel it
- * await discovery.done;
- * ```
+ * Manager for wallet discovery, configuration, and connection
  */
 export class WalletManager {
   private config: WalletManagerConfig = {
@@ -67,177 +33,65 @@ export class WalletManager {
 
   /**
    * Discovers all available wallets for a given chain and version.
-   *
-   * Returns a `DiscoverySession` with:
-   * - `wallets`: AsyncIterable to iterate over discovered wallets
-   * - `done`: Promise that resolves when discovery completes or is cancelled
-   * - `cancel()`: Function to stop discovery immediately
-   *
-   * If `onWalletDiscovered` callback is provided, wallets are also streamed via callback.
-   *
-   * @param options - Discovery options including chain info, appId, and timeout
-   * @returns A cancellable discovery session
+   * Only returns wallets that support the requested chain and version.
+   * @param options - Discovery options including chain info and timeout
+   * @returns Array of wallet providers with baked-in chain info
    */
-  getAvailableWallets(options: DiscoverWalletsOptions): DiscoverySession {
-    const { chainInfo, appId } = options;
-    const abortController = new AbortController();
-
-    const pendingProviders: WalletProvider[] = [];
-    let pendingResolve: ((result: IteratorResult<WalletProvider>) => void) | null = null;
-    let completed = false;
-
-    const { promise: donePromise, resolve: resolveDone } = promiseWithResolvers<void>();
-
-    const markComplete = () => {
-      completed = true;
-      resolveDone();
-      if (pendingResolve) {
-        const resolve = pendingResolve;
-        pendingResolve = null;
-        resolve({ value: undefined, done: true });
-      }
-    };
+  async getAvailableWallets(options: DiscoverWalletsOptions): Promise<WalletProvider[]> {
+    const providers: WalletProvider[] = [];
+    const { chainInfo } = options;
 
     if (this.config.extensions?.enabled) {
+      const discoveredWallets = await ExtensionProvider.discoverExtensions(chainInfo, options.timeout);
       const extensionConfig = this.config.extensions;
 
-      void ExtensionProvider.discoverWallets(chainInfo, {
-        appId,
-        timeout: options.timeout,
-        signal: abortController.signal,
-        onWalletDiscovered: discoveredWallet => {
-          const provider = this.createProviderFromDiscoveredWallet(discoveredWallet, chainInfo, extensionConfig);
-          if (!provider) {
-            return;
-          }
-
-          // Call user's callback if provided
-          options.onWalletDiscovered?.(provider);
+      for (const { info, port, sharedKey } of discoveredWallets) {
+        if (!this.isExtensionAllowed(info.id, extensionConfig)) {
+          continue;
+        }
 
-          // Also queue for async iterator
-          if (pendingResolve) {
-            const resolve = pendingResolve;
-            pendingResolve = null;
-            resolve({ value: provider, done: false });
-          } else {
-            pendingProviders.push(provider);
-          }
-        },
-      }).then(markComplete);
-    } else {
-      markComplete();
-    }
+        let extensionWallet: ExtensionWallet | null = null;
 
-    const wallets: AsyncIterable<WalletProvider> = {
-      // eslint-disable-next-line jsdoc/require-jsdoc
-      [Symbol.asyncIterator](): AsyncIterator<WalletProvider> {
-        return {
-          // eslint-disable-next-line jsdoc/require-jsdoc
-          next(): Promise<IteratorResult<WalletProvider>> {
-            if (pendingProviders.length > 0) {
-              return Promise.resolve({ value: pendingProviders.shift()!, done: false });
+        const provider: WalletProvider = {
+          id: info.id,
+          type: 'extension',
+          name: info.name,
+          icon: info.icon,
+          metadata: {
+            version: info.version,
+            verificationHash: info.verificationHash,
+          },
+          connect: (appId: string) => {
+            extensionWallet = ExtensionWallet.create(info, chainInfo, port, sharedKey, appId);
+            return Promise.resolve(extensionWallet.getWallet());
+          },
+          disconnect: async () => {
+            if (extensionWallet) {
+              await extensionWallet.disconnect();
+              extensionWallet = null;
             }
-
-            if (completed) {
-              return Promise.resolve({ value: undefined, done: true });
+          },
+          onDisconnect: (callback: ProviderDisconnectionCallback) => {
+            if (extensionWallet) {
+              return extensionWallet.onDisconnect(callback);
             }
-
-            return new Promise(resolve => {
-              pendingResolve = resolve;
-            });
+            return () => {};
+          },
+          isDisconnected: () => {
+            if (extensionWallet) {
+              return extensionWallet.isDisconnected();
+            }
+            return true;
           },
         };
-      },
-    };
 
-    return {
-      wallets,
-      done: donePromise,
-      cancel: () => abortController.abort(),
-    };
-  }
-
-  /**
-   * Creates a WalletProvider from a discovered wallet.
-   * Returns null if the wallet is not allowed by config.
-   * @param discoveredWallet - The discovered wallet from extension discovery.
-   * @param chainInfo - Network information.
-   * @param extensionConfig - Extension wallet configuration.
-   */
-  private createProviderFromDiscoveredWallet(
-    discoveredWallet: DiscoveredWallet,
-    chainInfo: ChainInfo,
-    extensionConfig: ExtensionWalletConfig,
-  ): WalletProvider | null {
-    const { info } = discoveredWallet;
-
-    if (!this.isExtensionAllowed(info.id, extensionConfig)) {
-      return null;
+        providers.push(provider);
+      }
     }
 
-    let extensionWallet: ExtensionWallet | null = null;
-
-    const provider: WalletProvider = {
-      id: info.id,
-      type: 'extension',
-      name: info.name,
-      icon: info.icon,
-      metadata: {
-        version: info.version,
-      },
-      establishSecureChannel: async (connectAppId: string): Promise<PendingConnection> => {
-        const connection = await discoveredWallet.establishSecureChannel();
-
-        provider.metadata = {
-          ...provider.metadata,
-          verificationHash: connection.info.verificationHash,
-        };
-
-        return {
-          verificationHash: connection.info.verificationHash!,
-          confirm: () => {
-            return Promise.resolve(
-              ExtensionWallet.create(
-                connection.info.id,
-                connection.port,
-                connection.sharedKey,
-                chainInfo,
-                connectAppId,
-              ),
-            );
-          },
-          cancel: () => {
-            // Send disconnect to terminate the session on the extension side
-            // but keep the port open so we can retry key exchange
-            connection.port.postMessage({
-              type: WalletMessageType.DISCONNECT,
-              requestId: discoveredWallet.requestId,
-            });
-            // Don't close the port - allow retry with fresh key exchange
-          },
-        };
-      },
-      disconnect: async () => {
-        if (extensionWallet) {
-          await extensionWallet.disconnect();
-          extensionWallet = null;
-        }
-      },
-      onDisconnect: (callback: ProviderDisconnectionCallback) => {
-        if (extensionWallet) {
-          return extensionWallet.onDisconnect(callback);
-        }
-        return () => {};
-      },
-      isDisconnected: () => {
-        if (extensionWallet) {
-          return extensionWallet.isDisconnected();
-        }
-        return true;
-      },
-    };
+    // TODO: Add web wallet discovery when implemented
 
-    return provider;
+    return providers;
   }
 
   /**

package/src/providers/extension/extension_provider.ts

@@ -0,0 +1,167 @@
+import type { ChainInfo } from '@aztec/aztec.js/account';
+import { jsonStringify } from '@aztec/foundation/json-rpc';
+import { promiseWithResolvers } from '@aztec/foundation/promise';
+
+import { deriveSharedKey, exportPublicKey, generateKeyPair, hashSharedSecret, importPublicKey } from '../../crypto.js';
+import { type DiscoveryRequest, type DiscoveryResponse, type WalletInfo, WalletMessageType } from '../../types.js';
+
+/**
+ * A discovered wallet with its secure channel components.
+ * Returned by {@link ExtensionProvider.discoverExtensions}.
+ */
+export interface DiscoveredWallet {
+  /** Basic wallet information (id, name, icon, version, publicKey, verificationHash) */
+  info: WalletInfo;
+  /** The MessagePort for private communication with the wallet */
+  port: MessagePort;
+  /** The derived AES-256-GCM shared key for encryption */
+  sharedKey: CryptoKey;
+}
+
+/**
+ * Internal type for discovery response with MessagePort
+ * @internal
+ */
+interface DiscoveryResponseWithPort extends DiscoveryResponse {
+  /** The MessagePort transferred from the wallet */
+  port?: MessagePort;
+}
+
+/**
+ * Provider for discovering Aztec wallet extensions.
+ *
+ * This class handles the discovery phase of wallet communication:
+ * 1. Broadcasts a discovery request with the dApp's public key
+ * 2. Receives responses from installed wallets with their public keys
+ * 3. Derives shared secrets and computes verification hashes
+ * 4. Returns discovered wallets with their secure channel components
+ *
+ * @example
+ * ```typescript
+ * const wallets = await ExtensionProvider.discoverExtensions(chainInfo);
+ * // Display wallets to user with optional emoji verification
+ * for (const wallet of wallets) {
+ *   const emoji = hashToEmoji(wallet.info.verificationHash!);
+ *   console.log(`${wallet.info.name}: ${emoji}`);
+ * }
+ * // User selects a wallet after verifying
+ * const wallet = await ExtensionWallet.create(wallets[0], chainInfo, 'my-app');
+ * ```
+ */
+export class ExtensionProvider {
+  private static discoveryInProgress = false;
+
+  /**
+   * Discovers all installed Aztec wallet extensions and establishes secure channel components.
+   *
+   * This method:
+   * 1. Generates an ECDH key pair for this discovery session
+   * 2. Broadcasts a discovery request with the dApp's public key
+   * 3. Receives responses from wallets with their public keys and MessagePorts
+   * 4. Derives shared secrets and computes verification hashes
+   *
+   * @param chainInfo - Chain information to check if extensions support this network
+   * @param timeout - How long to wait for extensions to respond (ms)
+   * @returns Array of discovered wallets with their secure channel components
+   *
+   * @example
+   * ```typescript
+   * const wallets = await ExtensionProvider.discoverExtensions({
+   *   chainId: Fr(31337),
+   *   version: Fr(0)
+   * });
+   * // Access wallet info and secure channel
+   * const { info, port, sharedKey } = wallets[0];
+   * ```
+   */
+  static async discoverExtensions(chainInfo: ChainInfo, timeout: number = 1000): Promise<DiscoveredWallet[]> {
+    // If discovery is already in progress, wait and return empty
+    // (caller should retry or handle appropriately)
+    if (this.discoveryInProgress) {
+      await new Promise(resolve => setTimeout(resolve, timeout));
+      return [];
+    }
+
+    this.discoveryInProgress = true;
+
+    // Generate key pair for this discovery session
+    const keyPair = await generateKeyPair();
+    const exportedPublicKey = await exportPublicKey(keyPair.publicKey);
+
+    const { promise, resolve } = promiseWithResolvers<DiscoveredWallet[]>();
+    const requestId = globalThis.crypto.randomUUID();
+    const responses: DiscoveredWallet[] = [];
+
+    // Set up listener for discovery responses
+    const handleMessage = (event: MessageEvent) => {
+      if (event.source !== window) {
+        return;
+      }
+
+      let data: DiscoveryResponseWithPort;
+      try {
+        data = JSON.parse(event.data);
+      } catch {
+        return;
+      }
+
+      if (data.type === WalletMessageType.DISCOVERY_RESPONSE && data.requestId === requestId) {
+        // Get the MessagePort from the event
+        const port = event.ports?.[0];
+        if (!port) {
+          return;
+        }
+
+        // Derive shared key from wallet's public key
+        const walletPublicKey = data.walletInfo.publicKey;
+        if (!walletPublicKey) {
+          return;
+        }
+
+        void (async () => {
+          try {
+            const importedWalletKey = await importPublicKey(walletPublicKey);
+            const sharedKey = await deriveSharedKey(keyPair.privateKey, importedWalletKey);
+
+            // Compute verification hash
+            const verificationHash = await hashSharedSecret(sharedKey);
+
+            // Create wallet info with verification hash
+            const walletInfo: WalletInfo = {
+              ...data.walletInfo,
+              verificationHash,
+            };
+
+            responses.push({
+              info: walletInfo,
+              port,
+              sharedKey,
+            });
+          } catch {
+            // Failed to derive key, skip this wallet
+          }
+        })();
+      }
+    };
+
+    window.addEventListener('message', handleMessage);
+
+    // Send discovery message with our public key
+    const discoveryMessage: DiscoveryRequest = {
+      type: WalletMessageType.DISCOVERY,
+      requestId,
+      chainInfo,
+      publicKey: exportedPublicKey,
+    };
+    window.postMessage(jsonStringify(discoveryMessage), '*');
+
+    // Wait for responses
+    setTimeout(() => {
+      window.removeEventListener('message', handleMessage);
+      this.discoveryInProgress = false;
+      resolve(responses);
+    }, timeout);
+
+    return promise;
+  }
+}