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

package/src/manager/types.ts

@@ -36,6 +36,11 @@ export interface WalletManagerConfig {
  */
 export type WalletProviderType = 'extension' | 'web' | 'embedded';
 
+/**
+ * Callback type for wallet disconnect events at the provider level.
+ */
+export type ProviderDisconnectionCallback = () => void;
+
 /**
  * A wallet provider that can connect to create a wallet instance.
  * Chain information is already baked in from the discovery process.
@@ -56,6 +61,23 @@ export interface WalletProvider {
    * @param appId - Application identifier for the requesting dapp
    */
   connect(appId: string): Promise<Wallet>;
+  /**
+   * Disconnects the current wallet and cleans up resources.
+   * After calling this, the wallet returned from connect() should no longer be used.
+   * @returns A promise that resolves when disconnection is complete
+   */
+  disconnect?(): Promise<void>;
+  /**
+   * Registers a callback to be invoked when the wallet disconnects unexpectedly.
+   * @param callback - Function to call when wallet disconnects
+   * @returns A function to unregister the callback
+   */
+  onDisconnect?(callback: ProviderDisconnectionCallback): () => void;
+  /**
+   * Returns whether the provider's wallet connection has been disconnected.
+   * @returns true if the wallet is no longer connected
+   */
+  isDisconnected?(): boolean;
 }
 
 /**

package/src/manager/wallet_manager.ts

@@ -1,5 +1,11 @@
 import { ExtensionProvider, ExtensionWallet } from '../providers/extension/index.js';
-import type { DiscoverWalletsOptions, ExtensionWalletConfig, WalletManagerConfig, WalletProvider } from './types.js';
+import type {
+  DiscoverWalletsOptions,
+  ExtensionWalletConfig,
+  ProviderDisconnectionCallback,
+  WalletManagerConfig,
+  WalletProvider,
+} from './types.js';
 
 /**
  * Manager for wallet discovery, configuration, and connection
@@ -35,27 +41,51 @@ export class WalletManager {
     const providers: WalletProvider[] = [];
     const { chainInfo } = options;
 
-    // Discover extension wallets
     if (this.config.extensions?.enabled) {
-      const extensions = await ExtensionProvider.discoverExtensions(chainInfo, options.timeout);
+      const discoveredWallets = await ExtensionProvider.discoverExtensions(chainInfo, options.timeout);
       const extensionConfig = this.config.extensions;
 
-      for (const ext of extensions) {
-        // Apply allow/block lists
-        if (!this.isExtensionAllowed(ext.id, extensionConfig)) {
+      for (const { info, port, sharedKey } of discoveredWallets) {
+        if (!this.isExtensionAllowed(info.id, extensionConfig)) {
           continue;
         }
 
-        providers.push({
-          id: ext.id,
+        let extensionWallet: ExtensionWallet | null = null;
+
+        const provider: WalletProvider = {
+          id: info.id,
           type: 'extension',
-          name: ext.name,
-          icon: ext.icon,
+          name: info.name,
+          icon: info.icon,
           metadata: {
-            version: ext.version,
+            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;
+            }
           },
-          connect: (appId: string) => Promise.resolve(ExtensionWallet.create(chainInfo, appId, ext.id)),
-        });
+          onDisconnect: (callback: ProviderDisconnectionCallback) => {
+            if (extensionWallet) {
+              return extensionWallet.onDisconnect(callback);
+            }
+            return () => {};
+          },
+          isDisconnected: () => {
+            if (extensionWallet) {
+              return extensionWallet.isDisconnected();
+            }
+            return true;
+          },
+        };
+
+        providers.push(provider);
       }
     }
 
@@ -70,17 +100,14 @@ export class WalletManager {
    * @param config - Extension wallet configuration containing allow/block lists
    */
   private isExtensionAllowed(extensionId: string, config: ExtensionWalletConfig): boolean {
-    // Check block list first
     if (config.blockList && config.blockList.includes(extensionId)) {
       return false;
     }
 
-    // If allow list exists, extension must be in it
     if (config.allowList && config.allowList.length > 0) {
       return config.allowList.includes(extensionId);
     }
 
-    // If no allow list, extension is allowed (unless blocked)
     return true;
   }
 }

package/src/providers/extension/extension_provider.ts

@@ -2,34 +2,95 @@ import type { ChainInfo } from '@aztec/aztec.js/account';
 import { jsonStringify } from '@aztec/foundation/json-rpc';
 import { promiseWithResolvers } from '@aztec/foundation/promise';
 
-import type { DiscoveryRequest, DiscoveryResponse, WalletInfo } from '../types.js';
+import { deriveSharedKey, exportPublicKey, generateKeyPair, hashSharedSecret, importPublicKey } from '../../crypto.js';
+import { type DiscoveryRequest, type DiscoveryResponse, type WalletInfo, WalletMessageType } from '../../types.js';
 
 /**
- * Provider for discovering and managing Aztec wallet extensions
+ * 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 discoveredExtensions: Map<string, WalletInfo> = new Map();
   private static discoveryInProgress = false;
 
   /**
-   * Discovers all installed Aztec wallet extensions
+   * 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 extension information
+   * @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<WalletInfo[]> {
-    // If discovery is in progress, wait for it to complete
+  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 Array.from(this.discoveredExtensions.values());
+      return [];
     }
 
     this.discoveryInProgress = true;
-    this.discoveredExtensions.clear();
 
-    const { promise, resolve } = promiseWithResolvers<WalletInfo[]>();
+    // 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: WalletInfo[] = [];
+    const responses: DiscoveredWallet[] = [];
 
     // Set up listener for discovery responses
     const handleMessage = (event: MessageEvent) => {
@@ -37,26 +98,60 @@ export class ExtensionProvider {
         return;
       }
 
-      let data: DiscoveryResponse;
+      let data: DiscoveryResponseWithPort;
       try {
         data = JSON.parse(event.data);
       } catch {
         return;
       }
 
-      if (data.type === 'aztec-wallet-discovery-response' && data.requestId === requestId) {
-        responses.push(data.walletInfo);
-        this.discoveredExtensions.set(data.walletInfo.id, data.walletInfo);
+      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
+    // Send discovery message with our public key
     const discoveryMessage: DiscoveryRequest = {
-      type: 'aztec-wallet-discovery',
+      type: WalletMessageType.DISCOVERY,
       requestId,
       chainInfo,
+      publicKey: exportedPublicKey,
     };
     window.postMessage(jsonStringify(discoveryMessage), '*');