feat: fetch keychain per wallet once

Ticket: CSI-1548

Author: abk404

modules/bitgo/test/v2/unit/wallets.ts

@@ -3708,6 +3708,99 @@ describe('V2 Wallets:', function () {
         assert.equal(error.message, 'Password shared is incorrect for this wallet');
       }
     });
+
+    it('should fetch keychain only ONCE when sharing with multiple users', async () => {
+      const walletPassphrase = 'bitgo1234';
+      const pub = 'Zo1ggzTUKMY5bYnDvT5mtVeZxzf2FaLTbKkmvGUhUQk';
+      const prv1 = 'xprv1';
+
+      // Create multiple users to share with
+      const users = [
+        {
+          userId: 'user1@example.com',
+          permissions: ['view', 'spend'],
+          pubKey: '02705a6d33a2459feb537e7abe36aaad8c11532cdbffa3a2e4e58868467d51f532',
+          path: 'm/999999/1/1',
+        },
+        {
+          userId: 'user2@example.com',
+          permissions: ['view', 'spend'],
+          pubKey: '03705a6d33a2459feb537e7abe36aaad8c11532cdbffa3a2e4e58868467d51f533',
+          path: 'm/999999/1/2',
+        },
+        {
+          userId: 'user3@example.com',
+          permissions: ['view', 'spend'],
+          pubKey: '02815a6d33a2459feb537e7abe36aaad8c11532cdbffa3a2e4e58868467d51f534',
+          path: 'm/999999/1/3',
+        },
+      ];
+
+      const params: BulkWalletShareOptions = {
+        walletPassphrase,
+        keyShareOptions: users,
+      };
+
+      const getDecryptedKeychainStub = sinon.stub(wallet, 'getDecryptedKeychainForSharing').resolves({
+        prv: prv1,
+        pub,
+      });
+
+      const encryptPrvForUserStub = sinon
+        .stub(wallet, 'encryptPrvForUser')
+        .callsFake((prv, pubKey, userPubKey, path) => {
+          return {
+            pub: pubKey,
+            encryptedPrv: 'dummyEncryptedPrv',
+            fromPubKey: 'dummyFromPubKey',
+            toPubKey: userPubKey,
+            path: path,
+          };
+        });
+
+      sinon.stub(wallet, 'createBulkKeyShares').resolves({
+        shares: users.map((user) => ({
+          id: user.userId,
+          coin: walletData.coin,
+          wallet: walletData.id,
+          fromUser: 'fromUser',
+          toUser: user.userId,
+          permissions: user.permissions,
+          keychain: {
+            pub: 'dummyPub',
+            encryptedPrv: 'dummyEncryptedPrv',
+            fromPubKey: 'dummyFromPubKey',
+            toPubKey: user.pubKey,
+            path: user.path,
+          },
+        })),
+      });
+
+      await wallet.createBulkWalletShare(params);
+
+      // Verify keychain is fetched only ONCE, not once per user
+      assert.strictEqual(
+        getDecryptedKeychainStub.callCount,
+        1,
+        'getDecryptedKeychainForSharing should be called exactly once, not once per user'
+      );
+
+      // Verify encryption happens for EACH user (once per user)
+      assert.strictEqual(
+        encryptPrvForUserStub.callCount,
+        users.length,
+        `encryptPrvForUser should be called once per user (${users.length} times)`
+      );
+
+      // Verify each call had the correct user's pubKey and path
+      users.forEach((user, index) => {
+        const call = encryptPrvForUserStub.getCall(index);
+        assert.strictEqual(call.args[0], prv1, 'Should use the same decrypted private key for all users');
+        assert.strictEqual(call.args[1], pub, 'Should use the same public key for all users');
+        assert.strictEqual(call.args[2], user.pubKey, `Should use user${index + 1}'s public key`);
+        assert.strictEqual(call.args[3], user.path, `Should use user${index + 1}'s path`);
+      });
+    });
   });
 
   describe('List Wallets:', function () {

modules/sdk-core/src/bitgo/wallet/iWallet.ts

@@ -645,6 +645,11 @@ export interface BulkWalletShareOptions {
 
 export type WalletShareState = 'active' | 'accepted' | 'canceled' | 'rejected' | 'pendingapproval';
 
+export interface DecryptedKeychainData {
+  prv: string;
+  pub: string;
+}
+
 export interface BulkWalletShareKeychain {
   pub: string;
   encryptedPrv: string;

modules/sdk-core/src/bitgo/wallet/wallet.ts

@@ -70,6 +70,7 @@ import {
   CreatePolicyRuleOptions,
   CreateShareOptions,
   CrossChainUTXO,
+  DecryptedKeychainData,
   DeployForwardersOptions,
   DownloadKeycardOptions,
   FanoutUnspentsOptions,
@@ -1685,8 +1686,27 @@ export class Wallet implements IWallet {
     if (params.keyShareOptions.length === 0) {
       throw new Error('No share options provided');
     }
+
     const bulkCreateShareOptions: BulkCreateShareOption[] = [];
 
+    // Check if any share option needs a keychain (has 'spend' permission)
+    const anyNeedsKeychain = params.keyShareOptions.some((opt) => opt.permissions && opt.permissions.includes('spend'));
+
+    // Fetch and decrypt the keychain ONCE for all users
+    let decryptedKeychain: DecryptedKeychainData | undefined;
+
+    if (anyNeedsKeychain) {
+      try {
+        decryptedKeychain = await this.getDecryptedKeychainForSharing(params.walletPassphrase);
+      } catch (e) {
+        if (e instanceof MissingEncryptedKeychainError) {
+          decryptedKeychain = undefined;
+        } else {
+          throw e;
+        }
+      }
+    }
+
     for (const shareOption of params.keyShareOptions) {
       try {
         common.validateParams(shareOption, ['userId', 'pubKey', 'path'], []);
@@ -1699,36 +1719,22 @@ export class Wallet implements IWallet {
 
       const needsKeychain = shareOption.permissions && shareOption.permissions.includes('spend');
 
-      if (needsKeychain) {
-        const sharedKeychain = await this.prepareSharedKeychain(
-          params.walletPassphrase,
+      if (needsKeychain && decryptedKeychain) {
+        const sharedKeychain = this.encryptPrvForUser(
+          decryptedKeychain.prv,
+          decryptedKeychain.pub,
           shareOption.pubKey,
           shareOption.path
         );
-        const keychain = Object.keys(sharedKeychain ?? {}).length === 0 ? undefined : sharedKeychain;
-        if (keychain) {
-          assert(keychain.pub, 'pub must be defined for sharing');
-          assert(keychain.encryptedPrv, 'encryptedPrv must be defined for sharing');
-          assert(keychain.fromPubKey, 'fromPubKey must be defined for sharing');
-          assert(keychain.toPubKey, 'toPubKey must be defined for sharing');
-          assert(keychain.path, 'path must be defined for sharing');
-
-          const bulkKeychain: BulkWalletShareKeychain = {
-            pub: keychain.pub,
-            encryptedPrv: keychain.encryptedPrv,
-            fromPubKey: keychain.fromPubKey,
-            toPubKey: keychain.toPubKey,
-            path: keychain.path,
-          };
 
-          bulkCreateShareOptions.push({
-            user: shareOption.userId,
-            permissions: shareOption.permissions,
-            keychain: bulkKeychain,
-          });
-        }
+        bulkCreateShareOptions.push({
+          user: shareOption.userId,
+          permissions: shareOption.permissions,
+          keychain: sharedKeychain,
+        });
       }
     }
+
     return await this.createBulkKeyShares(bulkCreateShareOptions);
   }
 
@@ -1776,62 +1782,107 @@ export class Wallet implements IWallet {
     }
   }
 
-  async prepareSharedKeychain(
-    walletPassphrase: string | undefined,
-    pubkey: string,
-    path: string
-  ): Promise<SharedKeyChain> {
-    let sharedKeychain: SharedKeyChain = {};
+  /**
+   * Fetches and decrypts the wallet keychain for sharing.
+   * This method fetches the keychain once and decrypts it, returning the decrypted
+   * private key and public key info needed for sharing with multiple users.
+   *
+   * @param walletPassphrase - The passphrase to decrypt the keychain
+   * @returns Object containing decrypted prv and pub, or undefined for cold wallets
+   */
+  async getDecryptedKeychainForSharing(
+    walletPassphrase: string | undefined
+  ): Promise<DecryptedKeychainData | undefined> {
+    const keychain = await this.getEncryptedWalletKeychainForWalletSharing();
 
-    try {
-      const keychain = await this.getEncryptedWalletKeychainForWalletSharing();
+    if (!keychain.encryptedPrv) {
+      return undefined;
+    }
 
-      // Decrypt the user key with a passphrase
-      if (keychain.encryptedPrv) {
-        if (!walletPassphrase) {
-          throw new Error('Missing walletPassphrase argument');
-        }
+    if (!walletPassphrase) {
+      throw new Error('Missing walletPassphrase argument');
+    }
 
-        const userPrv = decryptKeychainPrivateKey(this.bitgo, keychain, walletPassphrase);
-        if (!userPrv) {
-          throw new IncorrectPasswordError('Password shared is incorrect for this wallet');
-        }
+    const prv = decryptKeychainPrivateKey(this.bitgo, keychain, walletPassphrase);
+    if (!prv) {
+      throw new IncorrectPasswordError('Password shared is incorrect for this wallet');
+    }
 
-        keychain.prv = userPrv;
-        const eckey = makeRandomKey();
-        const secret = getSharedSecret(eckey, Buffer.from(pubkey, 'hex')).toString('hex');
-        const newEncryptedPrv = this.bitgo.encrypt({
-          password: secret,
-          input: keychain.prv,
-        });
+    // Only one of pub/commonPub/commonKeychain should be present in the keychain
+    let pub = keychain.pub ?? keychain.commonPub;
+    if (keychain.commonKeychain) {
+      pub =
+        this.baseCoin.getMPCAlgorithm() === 'eddsa'
+          ? EddsaUtils.getPublicKeyFromCommonKeychain(keychain.commonKeychain)
+          : EcdsaUtils.getPublicKeyFromCommonKeychain(keychain.commonKeychain);
+    }
 
-        // Only one of pub/commonPub/commonKeychain should be present in the keychain
-        let pub = keychain.pub ?? keychain.commonPub;
-        if (keychain.commonKeychain) {
-          pub =
-            this.baseCoin.getMPCAlgorithm() === 'eddsa'
-              ? EddsaUtils.getPublicKeyFromCommonKeychain(keychain.commonKeychain)
-              : EcdsaUtils.getPublicKeyFromCommonKeychain(keychain.commonKeychain);
-        }
+    if (!pub) {
+      throw new Error('Unable to determine public key from keychain');
+    }
 
-        sharedKeychain = {
-          pub,
-          encryptedPrv: newEncryptedPrv,
-          fromPubKey: eckey.publicKey.toString('hex'),
-          toPubKey: pubkey,
-          path: path,
-        };
+    return { prv, pub };
+  }
+
+  /**
+   * Encrypts a decrypted private key for sharing with a specific user.
+   * This is the pure encryption step - no API calls, no decryption.
+   *
+   * @param decryptedPrv - The already-decrypted private key
+   * @param pub - The wallet's public key
+   * @param userPubkey - The recipient user's public key
+   * @param path - The key path
+   * @returns The encrypted keychain for the recipient with all required fields
+   */
+  encryptPrvForUser(decryptedPrv: string, pub: string, userPubkey: string, path: string): BulkWalletShareKeychain {
+    const eckey = makeRandomKey();
+    const secret = getSharedSecret(eckey, Buffer.from(userPubkey, 'hex')).toString('hex');
+    const newEncryptedPrv = this.bitgo.encrypt({ password: secret, input: decryptedPrv });
+
+    const keychain: BulkWalletShareKeychain = {
+      pub,
+      encryptedPrv: newEncryptedPrv,
+      fromPubKey: eckey.publicKey.toString('hex'),
+      toPubKey: userPubkey,
+      path: path,
+    };
+
+    assert(keychain.pub, 'pub must be defined for sharing');
+    assert(keychain.encryptedPrv, 'encryptedPrv must be defined for sharing');
+    assert(keychain.fromPubKey, 'fromPubKey must be defined for sharing');
+    assert(keychain.toPubKey, 'toPubKey must be defined for sharing');
+    assert(keychain.path, 'path must be defined for sharing');
+
+    return keychain;
+  }
+
+  /**
+   * Prepares a keychain for sharing with another user.
+   * Fetches the wallet keychain, decrypts it, and encrypts it for the recipient.
+   *
+   * @param walletPassphrase - The passphrase to decrypt the keychain
+   * @param pubkey - The recipient's public key
+   * @param path - The key path
+   * @returns The encrypted keychain for the recipient
+   */
+  async prepareSharedKeychain(
+    walletPassphrase: string | undefined,
+    pubkey: string,
+    path: string
+  ): Promise<SharedKeyChain> {
+    try {
+      const decryptedKeychain = await this.getDecryptedKeychainForSharing(walletPassphrase);
+      if (!decryptedKeychain) {
+        return {};
       }
+      return this.encryptPrvForUser(decryptedKeychain.prv, decryptedKeychain.pub, pubkey, path);
     } catch (e) {
       if (e instanceof MissingEncryptedKeychainError) {
-        sharedKeychain = {};
         // ignore this error because this looks like a cold wallet
-      } else {
-        throw e;
+        return {};
       }
+      throw e;
     }
-
-    return sharedKeychain;
   }
 
   /**